From bf736bc55db2e4606e4caa97cd0963742c726d23 Mon Sep 17 00:00:00 2001 From: Yeuoly <45712896+Yeuoly@users.noreply.github.com> Date: Mon, 5 Feb 2024 18:48:30 +0800 Subject: [PATCH] Feat/show detailed custom api response when testing (#2400) --- api/.env.example | 5 ++- .../api_based_extension_requestor.py | 6 +-- api/core/helper/ssrf_proxy.py | 42 +++++++++++++++++++ api/core/tools/tool/api_tool.py | 19 +++++---- api/services/tools_manage_service.py | 4 +- 5 files changed, 61 insertions(+), 15 deletions(-) create mode 100644 api/core/helper/ssrf_proxy.py diff --git a/api/.env.example b/api/.env.example index bf5bf7c4e5..1e99a369ac 100644 --- a/api/.env.example +++ b/api/.env.example @@ -120,4 +120,7 @@ HOSTED_ANTHROPIC_QUOTA_LIMIT=600000 HOSTED_ANTHROPIC_PAID_ENABLED=false ETL_TYPE=dify -UNSTRUCTURED_API_URL= \ No newline at end of file +UNSTRUCTURED_API_URL= + +SSRF_PROXY_HTTP_URL= +SSRF_PROXY_HTTPS_URL= diff --git a/api/core/extension/api_based_extension_requestor.py b/api/core/extension/api_based_extension_requestor.py index c244fe88f1..97e2f4bdf0 100644 --- a/api/core/extension/api_based_extension_requestor.py +++ b/api/core/extension/api_based_extension_requestor.py @@ -30,10 +30,10 @@ class APIBasedExtensionRequestor: try: # proxy support for security proxies = None - if os.environ.get("API_BASED_EXTENSION_HTTP_PROXY") and os.environ.get("API_BASED_EXTENSION_HTTPS_PROXY"): + if os.environ.get("SSRF_PROXY_HTTP_URL") and os.environ.get("SSRF_PROXY_HTTPS_URL"): proxies = { - 'http': os.environ.get("API_BASED_EXTENSION_HTTP_PROXY"), - 'https': os.environ.get("API_BASED_EXTENSION_HTTPS_PROXY"), + 'http': os.environ.get("SSRF_PROXY_HTTP_URL"), + 'https': os.environ.get("SSRF_PROXY_HTTPS_URL"), } response = requests.request( diff --git a/api/core/helper/ssrf_proxy.py b/api/core/helper/ssrf_proxy.py new file mode 100644 index 0000000000..9501f2ce5b --- /dev/null +++ b/api/core/helper/ssrf_proxy.py @@ -0,0 +1,42 @@ +""" +Proxy requests to avoid SSRF +""" + +from httpx import get as _get, post as _post, put as _put, patch as _patch, head as _head, options as _options +from requests import delete as _delete + +import os + +SSRF_PROXY_HTTP_URL = os.getenv('SSRF_PROXY_HTTP_URL', '') +SSRF_PROXY_HTTPS_URL = os.getenv('SSRF_PROXY_HTTPS_URL', '') + +requests_proxies = { + 'http': SSRF_PROXY_HTTP_URL, + 'https': SSRF_PROXY_HTTPS_URL +} if SSRF_PROXY_HTTP_URL and SSRF_PROXY_HTTPS_URL else None + +httpx_proxies = { + 'http://': SSRF_PROXY_HTTP_URL, + 'https://': SSRF_PROXY_HTTPS_URL +} if SSRF_PROXY_HTTP_URL and SSRF_PROXY_HTTPS_URL else None + +def get(url, *args, **kwargs): + return _get(url=url, *args, proxies=httpx_proxies, **kwargs) + +def post(url, *args, **kwargs): + return _post(url=url, *args, proxies=httpx_proxies, **kwargs) + +def put(url, *args, **kwargs): + return _put(url=url, *args, proxies=httpx_proxies, **kwargs) + +def patch(url, *args, **kwargs): + return _patch(url=url, *args, proxies=httpx_proxies, **kwargs) + +def delete(url, *args, **kwargs): + return _delete(url=url, *args, proxies=requests_proxies, **kwargs) + +def head(url, *args, **kwargs): + return _head(url=url, *args, proxies=httpx_proxies, **kwargs) + +def options(url, *args, **kwargs): + return _options(url=url, *args, proxies=httpx_proxies, **kwargs) diff --git a/api/core/tools/tool/api_tool.py b/api/core/tools/tool/api_tool.py index 2d8845e98b..c72a6ec183 100644 --- a/api/core/tools/tool/api_tool.py +++ b/api/core/tools/tool/api_tool.py @@ -4,6 +4,7 @@ from typing import Any, Dict, List, Union import httpx import requests +import core.helper.ssrf_proxy as ssrf_proxy from core.tools.entities.tool_bundle import ApiBasedToolBundle from core.tools.entities.tool_entities import ToolInvokeMessage from core.tools.errors import ToolProviderCredentialValidationError @@ -31,7 +32,7 @@ class ApiTool(Tool): runtime=Tool.Runtime(**meta) ) - def validate_credentials(self, credentials: Dict[str, Any], parameters: Dict[str, Any], format_only: bool = False) -> None: + def validate_credentials(self, credentials: Dict[str, Any], parameters: Dict[str, Any], format_only: bool = False) -> str: """ validate the credentials for Api tool """ @@ -43,7 +44,7 @@ class ApiTool(Tool): response = self.do_http_request(self.api_bundle.server_url, self.api_bundle.method, headers, parameters) # validate response - self.validate_and_parse_response(response) + return self.validate_and_parse_response(response) def assembling_request(self, parameters: Dict[str, Any]) -> Dict[str, Any]: headers = {} @@ -201,23 +202,23 @@ class ApiTool(Tool): # do http request if method == 'get': - response = httpx.get(url, params=params, headers=headers, cookies=cookies, timeout=10, follow_redirects=True) + response = ssrf_proxy.get(url, params=params, headers=headers, cookies=cookies, timeout=10, follow_redirects=True) elif method == 'post': - response = httpx.post(url, params=params, headers=headers, cookies=cookies, data=body, timeout=10, follow_redirects=True) + response = ssrf_proxy.post(url, params=params, headers=headers, cookies=cookies, data=body, timeout=10, follow_redirects=True) elif method == 'put': - response = httpx.put(url, params=params, headers=headers, cookies=cookies, data=body, timeout=10, follow_redirects=True) + response = ssrf_proxy.put(url, params=params, headers=headers, cookies=cookies, data=body, timeout=10, follow_redirects=True) elif method == 'delete': """ request body data is unsupported for DELETE method in standard http protocol however, OpenAPI 3.0 supports request body data for DELETE method, so we support it here by using requests """ - response = requests.delete(url, params=params, headers=headers, cookies=cookies, data=body, timeout=10, allow_redirects=True) + response = ssrf_proxy.delete(url, params=params, headers=headers, cookies=cookies, data=body, timeout=10, allow_redirects=True) elif method == 'patch': - response = httpx.patch(url, params=params, headers=headers, cookies=cookies, data=body, timeout=10, follow_redirects=True) + response = ssrf_proxy.patch(url, params=params, headers=headers, cookies=cookies, data=body, timeout=10, follow_redirects=True) elif method == 'head': - response = httpx.head(url, params=params, headers=headers, cookies=cookies, timeout=10, follow_redirects=True) + response = ssrf_proxy.head(url, params=params, headers=headers, cookies=cookies, timeout=10, follow_redirects=True) elif method == 'options': - response = httpx.options(url, params=params, headers=headers, cookies=cookies, timeout=10, follow_redirects=True) + response = ssrf_proxy.options(url, params=params, headers=headers, cookies=cookies, timeout=10, follow_redirects=True) else: raise ValueError(f'Invalid http method {method}') diff --git a/api/services/tools_manage_service.py b/api/services/tools_manage_service.py index 2a97b3e8b5..ac21d1d9d6 100644 --- a/api/services/tools_manage_service.py +++ b/api/services/tools_manage_service.py @@ -521,8 +521,8 @@ class ToolManageService: 'credentials': credentials, 'tenant_id': tenant_id, }) - tool.validate_credentials(credentials, parameters) + result = tool.validate_credentials(credentials, parameters) except Exception as e: return { 'error': str(e) } - return { 'result': 'success' } \ No newline at end of file + return { 'result': result or 'empty response' } \ No newline at end of file