Feat/show detailed custom api response when testing (#2400)

This commit is contained in:
Yeuoly 2024-02-05 18:48:30 +08:00 committed by GitHub
parent d4cfd3e7ac
commit bf736bc55d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 61 additions and 15 deletions

View File

@ -121,3 +121,6 @@ HOSTED_ANTHROPIC_PAID_ENABLED=false
ETL_TYPE=dify ETL_TYPE=dify
UNSTRUCTURED_API_URL= UNSTRUCTURED_API_URL=
SSRF_PROXY_HTTP_URL=
SSRF_PROXY_HTTPS_URL=

View File

@ -30,10 +30,10 @@ class APIBasedExtensionRequestor:
try: try:
# proxy support for security # proxy support for security
proxies = None 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 = { proxies = {
'http': os.environ.get("API_BASED_EXTENSION_HTTP_PROXY"), 'http': os.environ.get("SSRF_PROXY_HTTP_URL"),
'https': os.environ.get("API_BASED_EXTENSION_HTTPS_PROXY"), 'https': os.environ.get("SSRF_PROXY_HTTPS_URL"),
} }
response = requests.request( response = requests.request(

View File

@ -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)

View File

@ -4,6 +4,7 @@ from typing import Any, Dict, List, Union
import httpx import httpx
import requests import requests
import core.helper.ssrf_proxy as ssrf_proxy
from core.tools.entities.tool_bundle import ApiBasedToolBundle from core.tools.entities.tool_bundle import ApiBasedToolBundle
from core.tools.entities.tool_entities import ToolInvokeMessage from core.tools.entities.tool_entities import ToolInvokeMessage
from core.tools.errors import ToolProviderCredentialValidationError from core.tools.errors import ToolProviderCredentialValidationError
@ -31,7 +32,7 @@ class ApiTool(Tool):
runtime=Tool.Runtime(**meta) 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 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) response = self.do_http_request(self.api_bundle.server_url, self.api_bundle.method, headers, parameters)
# validate response # 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]: def assembling_request(self, parameters: Dict[str, Any]) -> Dict[str, Any]:
headers = {} headers = {}
@ -201,23 +202,23 @@ class ApiTool(Tool):
# do http request # do http request
if method == 'get': 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': 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': 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': elif method == 'delete':
""" """
request body data is unsupported for DELETE method in standard http protocol 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 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': 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': 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': 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: else:
raise ValueError(f'Invalid http method {method}') raise ValueError(f'Invalid http method {method}')

View File

@ -521,8 +521,8 @@ class ToolManageService:
'credentials': credentials, 'credentials': credentials,
'tenant_id': tenant_id, 'tenant_id': tenant_id,
}) })
tool.validate_credentials(credentials, parameters) result = tool.validate_credentials(credentials, parameters)
except Exception as e: except Exception as e:
return { 'error': str(e) } return { 'error': str(e) }
return { 'result': 'success' } return { 'result': result or 'empty response' }