estel/dify
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

chore: refactor the http executor node (#5212)

Browse Source
This commit is contained in:
非法操作
2024-06-24 16:14:59 +08:00
committed by GitHub
parent 1e28a8c033
commit f7900f298f
8 changed files with 249 additions and 230 deletions
Download Patch File Download Diff File Expand all files Collapse all files

75
api/core/helper/ssrf_proxy.py
View File

@@ -1,65 +1,48 @@
"""
Proxy requests to avoid SSRF
"""
import os
from httpx import get as _get
from httpx import head as _head
from httpx import options as _options
from httpx import patch as _patch
from httpx import post as _post
from httpx import put as _put
from requests import delete as _delete
import httpx
SSRF_PROXY_ALL_URL = os.getenv('SSRF_PROXY_ALL_URL', '')
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 = {
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 make_request(method, url, **kwargs):
if SSRF_PROXY_ALL_URL:
return httpx.request(method=method, url=url, proxy=SSRF_PROXY_ALL_URL, **kwargs)
elif proxies:
return httpx.request(method=method, url=url, proxies=proxies, **kwargs)
else:
return httpx.request(method=method, url=url, **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 get(url, **kwargs):
return make_request('GET', url, **kwargs)
def delete(url, *args, **kwargs):
if 'follow_redirects' in kwargs:
if kwargs['follow_redirects']:
kwargs['allow_redirects'] = kwargs['follow_redirects']
kwargs.pop('follow_redirects')
if 'timeout' in kwargs:
timeout = kwargs['timeout']
if timeout is None:
kwargs.pop('timeout')
elif isinstance(timeout, tuple):
# check length of tuple
if len(timeout) == 2:
kwargs['timeout'] = timeout
elif len(timeout) == 1:
kwargs['timeout'] = timeout[0]
elif len(timeout) > 2:
kwargs['timeout'] = (timeout[0], timeout[1])
else:
kwargs['timeout'] = (timeout, timeout)
return _delete(url=url, *args, proxies=requests_proxies, **kwargs)
def head(url, *args, **kwargs):
return _head(url=url, *args, proxies=httpx_proxies, **kwargs)
def post(url, **kwargs):
return make_request('POST', url, **kwargs)
def options(url, *args, **kwargs):
return _options(url=url, *args, proxies=httpx_proxies, **kwargs)
def put(url, **kwargs):
return make_request('PUT', url, **kwargs)
def patch(url, **kwargs):
return make_request('PATCH', url, **kwargs)
def delete(url, **kwargs):
return make_request('DELETE', url, **kwargs)
def head(url, **kwargs):
return make_request('HEAD', url, **kwargs)

116
api/core/tools/tool/api_tool.py
View File

@@ -1,11 +1,9 @@
import json
from json import dumps
from os import getenv
from typing import Any, Union
from typing import Any
from urllib.parse import urlencode
import httpx
import requests
import core.helper.ssrf_proxy as ssrf_proxy
from core.tools.entities.tool_bundle import ApiToolBundle
@@ -18,12 +16,14 @@ API_TOOL_DEFAULT_TIMEOUT = (
int(getenv('API_TOOL_DEFAULT_READ_TIMEOUT', '60'))
)
class ApiTool(Tool):
api_bundle: ApiToolBundle
"""
Api tool
"""
def fork_tool_runtime(self, runtime: dict[str, Any]) -> 'Tool':
"""
fork a new tool with meta data
@@ -38,8 +38,9 @@ class ApiTool(Tool):
api_bundle=self.api_bundle.model_copy() if self.api_bundle else None,
runtime=Tool.Runtime(**runtime)
)
def validate_credentials(self, credentials: dict[str, Any], parameters: dict[str, Any], format_only: bool = False) -> str:
def validate_credentials(self, credentials: dict[str, Any], parameters: dict[str, Any],
format_only: bool = False) -> str:
"""
validate the credentials for Api tool
"""
@@ -47,7 +48,7 @@ class ApiTool(Tool):
headers = self.assembling_request(parameters)
if format_only:
return
return ''
response = self.do_http_request(self.api_bundle.server_url, self.api_bundle.method, headers, parameters)
# validate response
@@ -68,12 +69,12 @@ class ApiTool(Tool):
if 'api_key_header' in credentials:
api_key_header = credentials['api_key_header']
if 'api_key_value' not in credentials:
raise ToolProviderCredentialValidationError('Missing api_key_value')
elif not isinstance(credentials['api_key_value'], str):
raise ToolProviderCredentialValidationError('api_key_value must be a string')
if 'api_key_header_prefix' in credentials:
api_key_header_prefix = credentials['api_key_header_prefix']
if api_key_header_prefix == 'basic' and credentials['api_key_value']:
@@ -82,20 +83,20 @@ class ApiTool(Tool):
credentials['api_key_value'] = f'Bearer {credentials["api_key_value"]}'
elif api_key_header_prefix == 'custom':
pass
headers[api_key_header] = credentials['api_key_value']
needed_parameters = [parameter for parameter in self.api_bundle.parameters if parameter.required]
for parameter in needed_parameters:
if parameter.required and parameter.name not in parameters:
raise ToolParameterValidationError(f"Missing required parameter {parameter.name}")
if parameter.default is not None and parameter.name not in parameters:
parameters[parameter.name] = parameter.default
return headers
def validate_and_parse_response(self, response: Union[httpx.Response, requests.Response]) -> str:
def validate_and_parse_response(self, response: httpx.Response) -> str:
"""
validate the response
"""
@@ -112,23 +113,20 @@ class ApiTool(Tool):
return json.dumps(response)
except Exception as e:
return response.text
elif isinstance(response, requests.Response):
if not response.ok:
raise ToolInvokeError(f"Request failed with status code {response.status_code} and {response.text}")
if not response.content:
return 'Empty response from the tool, please check your parameters and try again.'
try:
response = response.json()
try:
return json.dumps(response, ensure_ascii=False)
except Exception as e:
return json.dumps(response)
except Exception as e:
return response.text
else:
raise ValueError(f'Invalid response type {type(response)}')
def do_http_request(self, url: str, method: str, headers: dict[str, Any], parameters: dict[str, Any]) -> httpx.Response:
@staticmethod
def get_parameter_value(parameter, parameters):
if parameter['name'] in parameters:
return parameters[parameter['name']]
elif parameter.get('required', False):
raise ToolParameterValidationError(f"Missing required parameter {parameter['name']}")
else:
return (parameter.get('schema', {}) or {}).get('default', '')
def do_http_request(self, url: str, method: str, headers: dict[str, Any],
parameters: dict[str, Any]) -> httpx.Response:
"""
do http request depending on api bundle
"""
@@ -141,44 +139,17 @@ class ApiTool(Tool):
# check parameters
for parameter in self.api_bundle.openapi.get('parameters', []):
value = self.get_parameter_value(parameter, parameters)
if parameter['in'] == 'path':
value = ''
if parameter['name'] in parameters:
value = parameters[parameter['name']]
elif parameter['required']:
raise ToolParameterValidationError(f"Missing required parameter {parameter['name']}")
else:
value = (parameter.get('schema', {}) or {}).get('default', '')
path_params[parameter['name']] = value
elif parameter['in'] == 'query':
value = ''
if parameter['name'] in parameters:
value = parameters[parameter['name']]
elif parameter.get('required', False):
raise ToolParameterValidationError(f"Missing required parameter {parameter['name']}")
else:
value = (parameter.get('schema', {}) or {}).get('default', '')
params[parameter['name']] = value
elif parameter['in'] == 'cookie':
value = ''
if parameter['name'] in parameters:
value = parameters[parameter['name']]
elif parameter.get('required', False):
raise ToolParameterValidationError(f"Missing required parameter {parameter['name']}")
else:
value = (parameter.get('schema', {}) or {}).get('default', '')
cookies[parameter['name']] = value
elif parameter['in'] == 'header':
value = ''
if parameter['name'] in parameters:
value = parameters[parameter['name']]
elif parameter.get('required', False):
raise ToolParameterValidationError(f"Missing required parameter {parameter['name']}")
else:
value = (parameter.get('schema', {}) or {}).get('default', '')
headers[parameter['name']] = value
# check if there is a request body and handle it
@@ -203,7 +174,7 @@ class ApiTool(Tool):
else:
body[name] = None
break
# replace path parameters
for name, value in path_params.items():
url = url.replace(f'{{{name}}}', f'{value}')
@@ -211,33 +182,21 @@ class ApiTool(Tool):
# parse http body data if needed, for GET/HEAD/OPTIONS/TRACE, the body is ignored
if 'Content-Type' in headers:
if headers['Content-Type'] == 'application/json':
body = dumps(body)
body = json.dumps(body)
elif headers['Content-Type'] == 'application/x-www-form-urlencoded':
body = urlencode(body)
else:
body = body
# do http request
if method == 'get':
response = ssrf_proxy.get(url, params=params, headers=headers, cookies=cookies, timeout=API_TOOL_DEFAULT_TIMEOUT, follow_redirects=True)
elif method == 'post':
response = ssrf_proxy.post(url, params=params, headers=headers, cookies=cookies, data=body, timeout=API_TOOL_DEFAULT_TIMEOUT, follow_redirects=True)
elif method == 'put':
response = ssrf_proxy.put(url, params=params, headers=headers, cookies=cookies, data=body, timeout=API_TOOL_DEFAULT_TIMEOUT, follow_redirects=True)
elif method == 'delete':
response = ssrf_proxy.delete(url, params=params, headers=headers, cookies=cookies, data=body, timeout=API_TOOL_DEFAULT_TIMEOUT, allow_redirects=True)
elif method == 'patch':
response = ssrf_proxy.patch(url, params=params, headers=headers, cookies=cookies, data=body, timeout=API_TOOL_DEFAULT_TIMEOUT, follow_redirects=True)
elif method == 'head':
response = ssrf_proxy.head(url, params=params, headers=headers, cookies=cookies, timeout=API_TOOL_DEFAULT_TIMEOUT, follow_redirects=True)
elif method == 'options':
response = ssrf_proxy.options(url, params=params, headers=headers, cookies=cookies, timeout=API_TOOL_DEFAULT_TIMEOUT, follow_redirects=True)
if method in ('get', 'head', 'post', 'put', 'delete', 'patch'):
response = getattr(ssrf_proxy, method)(url, params=params, headers=headers, cookies=cookies, data=body,
timeout=API_TOOL_DEFAULT_TIMEOUT, follow_redirects=True)
return response
else:
raise ValueError(f'Invalid http method {method}')
return response
def _convert_body_property_any_of(self, property: dict[str, Any], value: Any, any_of: list[dict[str, Any]], max_recursive=10) -> Any:
raise ValueError(f'Invalid http method {self.method}')
def _convert_body_property_any_of(self, property: dict[str, Any], value: Any, any_of: list[dict[str, Any]],
max_recursive=10) -> Any:
if max_recursive <= 0:
raise Exception("Max recursion depth reached")
for option in any_of or []:
@@ -322,4 +281,3 @@ class ApiTool(Tool):
# assemble invoke message
return self.create_text_message(response)

67
api/core/workflow/nodes/http_request/http_executor.py
View File

@@ -6,7 +6,6 @@ from typing import Any, Optional, Union
from urllib.parse import urlencode
import httpx
import requests
import core.helper.ssrf_proxy as ssrf_proxy
from core.workflow.entities.variable_entities import VariableSelector
@@ -22,14 +21,11 @@ READABLE_MAX_TEXT_SIZE = f'{MAX_TEXT_SIZE / 1024 / 1024:.2f}MB'
class HttpExecutorResponse:
headers: dict[str, str]
response: Union[httpx.Response, requests.Response]
response: httpx.Response
def __init__(self, response: Union[httpx.Response, requests.Response] = None):
self.headers = {}
if isinstance(response, httpx.Response | requests.Response):
for k, v in response.headers.items():
self.headers[k] = v
def __init__(self, response: httpx.Response = None):
self.response = response
self.headers = dict(response.headers) if isinstance(self.response, httpx.Response) else {}
@property
def is_file(self) -> bool:
@@ -42,10 +38,8 @@ class HttpExecutorResponse:
return any(v in content_type for v in file_content_types)
def get_content_type(self) -> str:
if 'content-type' in self.headers:
return self.headers.get('content-type')
else:
return self.headers.get('Content-Type') or ""
return self.headers.get('content-type', '')
def extract_file(self) -> tuple[str, bytes]:
"""
@@ -58,46 +52,31 @@ class HttpExecutorResponse:
@property
def content(self) -> str:
"""
get content
"""
if isinstance(self.response, httpx.Response | requests.Response):
if isinstance(self.response, httpx.Response):
return self.response.text
else:
raise ValueError(f'Invalid response type {type(self.response)}')
@property
def body(self) -> bytes:
"""
get body
"""
if isinstance(self.response, httpx.Response | requests.Response):
if isinstance(self.response, httpx.Response):
return self.response.content
else:
raise ValueError(f'Invalid response type {type(self.response)}')
@property
def status_code(self) -> int:
"""
get status code
"""
if isinstance(self.response, httpx.Response | requests.Response):
if isinstance(self.response, httpx.Response):
return self.response.status_code
else:
raise ValueError(f'Invalid response type {type(self.response)}')
@property
def size(self) -> int:
"""
get size
"""
return len(self.body)
@property
def readable_size(self) -> str:
"""
get readable size
"""
if self.size < 1024:
return f'{self.size} bytes'
elif self.size < 1024 * 1024:
@@ -148,13 +127,9 @@ class HttpExecutor:
return False
@staticmethod
def _to_dict(convert_item: str, convert_text: str, maxsplit: int = -1):
def _to_dict(convert_text: str):
"""
Convert the string like `aa:bb\n cc:dd` to dict `{aa:bb, cc:dd}`
:param convert_item: A label for what item to be converted, params, headers or body.
:param convert_text: The string containing key-value pairs separated by '\n'.
:param maxsplit: The maximum number of splits allowed for the ':' character in each key-value pair. Default is -1 (no limit).
:return: A dictionary containing the key-value pairs from the input string.
"""
kv_paris = convert_text.split('\n')
result = {}
@@ -162,15 +137,11 @@ class HttpExecutor:
if not kv.strip():
continue
kv = kv.split(':', maxsplit=maxsplit)
if len(kv) >= 3:
k, v = kv[0], ":".join(kv[1:])
elif len(kv) == 2:
k, v = kv
elif len(kv) == 1:
kv = kv.split(':', maxsplit=1)
if len(kv) == 1:
k, v = kv[0], ''
else:
raise ValueError(f'Invalid {convert_item} {kv}')
k, v = kv
result[k.strip()] = v
return result
@@ -181,11 +152,11 @@ class HttpExecutor:
# extract all template in params
params, params_variable_selectors = self._format_template(node_data.params, variable_pool)
self.params = self._to_dict("params", params)
self.params = self._to_dict(params)
# extract all template in headers
headers, headers_variable_selectors = self._format_template(node_data.headers, variable_pool)
self.headers = self._to_dict("headers", headers)
self.headers = self._to_dict(headers)
# extract all template in body
body_data_variable_selectors = []
@@ -203,7 +174,7 @@ class HttpExecutor:
self.headers['Content-Type'] = 'application/x-www-form-urlencoded'
if node_data.body.type in ['form-data', 'x-www-form-urlencoded']:
body = self._to_dict("body", body_data, 1)
body = self._to_dict(body_data)
if node_data.body.type == 'form-data':
self.files = {
@@ -242,11 +213,11 @@ class HttpExecutor:
return headers
def _validate_and_parse_response(self, response: Union[httpx.Response, requests.Response]) -> HttpExecutorResponse:
def _validate_and_parse_response(self, response: httpx.Response) -> HttpExecutorResponse:
"""
validate the response
"""
if isinstance(response, httpx.Response | requests.Response):
if isinstance(response, httpx.Response):
executor_response = HttpExecutorResponse(response)
else:
raise ValueError(f'Invalid response type {type(response)}')
@@ -274,9 +245,7 @@ class HttpExecutor:
'follow_redirects': True
}
if self.method in ('get', 'head', 'options'):
response = getattr(ssrf_proxy, self.method)(**kwargs)
elif self.method in ('post', 'put', 'delete', 'patch'):
if self.method in ('get', 'head', 'post', 'put', 'delete', 'patch'):
response = getattr(ssrf_proxy, self.method)(data=self.body, files=self.files, **kwargs)
else:
raise ValueError(f'Invalid http method {self.method}')