mirror of
https://github.com/1Panel-dev/MaxKB.git
synced 2025-12-26 01:33:05 +00:00
refactor: avoid using temp files in sandbox.
This commit is contained in:
parent
6d38a71906
commit
8bfce62ad8
|
|
@ -200,14 +200,6 @@ class BaseChatStep(IChatStep):
|
|||
mcp_enable, mcp_tool_ids, mcp_servers, mcp_source, tool_enable, tool_ids, mcp_output_enable)
|
||||
|
||||
def get_details(self, manage, **kwargs):
|
||||
# 删除临时生成的MCP代码文件
|
||||
if self.context.get('execute_ids'):
|
||||
executor = ToolExecutor(CONFIG.get('SANDBOX'))
|
||||
# 清理工具代码文件,延时删除,避免文件被占用
|
||||
for tool_id in self.context.get('execute_ids'):
|
||||
code_path = f'{executor.sandbox_path}/execute/{tool_id}.py'
|
||||
if os.path.exists(code_path):
|
||||
os.remove(code_path)
|
||||
return {
|
||||
'step_type': 'chat_step',
|
||||
'run_time': self.context['run_time'],
|
||||
|
|
@ -254,7 +246,6 @@ class BaseChatStep(IChatStep):
|
|||
if tool_enable:
|
||||
if tool_ids and len(tool_ids) > 0: # 如果有工具ID,则将其转换为MCP
|
||||
self.context['tool_ids'] = tool_ids
|
||||
self.context['execute_ids'] = []
|
||||
for tool_id in tool_ids:
|
||||
tool = QuerySet(Tool).filter(id=tool_id).first()
|
||||
if tool is None or tool.is_active is False:
|
||||
|
|
@ -264,9 +255,8 @@ class BaseChatStep(IChatStep):
|
|||
params = json.loads(rsa_long_decrypt(tool.init_params))
|
||||
else:
|
||||
params = {}
|
||||
_id, tool_config = executor.get_tool_mcp_config(tool.code, params)
|
||||
tool_config = executor.get_tool_mcp_config(tool.code, params)
|
||||
|
||||
self.context['execute_ids'].append(_id)
|
||||
mcp_servers_config[str(tool.id)] = tool_config
|
||||
|
||||
if len(mcp_servers_config) > 0:
|
||||
|
|
|
|||
|
|
@ -233,7 +233,6 @@ class BaseChatNode(IChatNode):
|
|||
if tool_enable:
|
||||
if tool_ids and len(tool_ids) > 0: # 如果有工具ID,则将其转换为MCP
|
||||
self.context['tool_ids'] = tool_ids
|
||||
self.context['execute_ids'] = []
|
||||
for tool_id in tool_ids:
|
||||
tool = QuerySet(Tool).filter(id=tool_id).first()
|
||||
if not tool.is_active:
|
||||
|
|
@ -243,9 +242,8 @@ class BaseChatNode(IChatNode):
|
|||
params = json.loads(rsa_long_decrypt(tool.init_params))
|
||||
else:
|
||||
params = {}
|
||||
_id, tool_config = executor.get_tool_mcp_config(tool.code, params)
|
||||
tool_config = executor.get_tool_mcp_config(tool.code, params)
|
||||
|
||||
self.context['execute_ids'].append(_id)
|
||||
mcp_servers_config[str(tool.id)] = tool_config
|
||||
|
||||
if len(mcp_servers_config) > 0:
|
||||
|
|
@ -307,14 +305,6 @@ class BaseChatNode(IChatNode):
|
|||
return result
|
||||
|
||||
def get_details(self, index: int, **kwargs):
|
||||
# 删除临时生成的MCP代码文件
|
||||
if self.context.get('execute_ids'):
|
||||
executor = ToolExecutor(CONFIG.get('SANDBOX'))
|
||||
# 清理工具代码文件,延时删除,避免文件被占用
|
||||
for tool_id in self.context.get('execute_ids'):
|
||||
code_path = f'{executor.sandbox_path}/execute/{tool_id}.py'
|
||||
if os.path.exists(code_path):
|
||||
os.remove(code_path)
|
||||
return {
|
||||
'name': self.node.properties.get('stepName'),
|
||||
"index": index,
|
||||
|
|
|
|||
|
|
@ -1,24 +1,25 @@
|
|||
# coding=utf-8
|
||||
import ast
|
||||
import base64
|
||||
import gzip
|
||||
import json
|
||||
import os
|
||||
import socket
|
||||
import subprocess
|
||||
import sys
|
||||
from textwrap import dedent
|
||||
import socket
|
||||
import uuid_utils.compat as uuid
|
||||
from common.utils.logger import maxkb_logger
|
||||
from django.utils.translation import gettext_lazy as _
|
||||
from maxkb.const import BASE_DIR, CONFIG
|
||||
from maxkb.const import PROJECT_DIR
|
||||
from common.utils.logger import maxkb_logger
|
||||
import threading
|
||||
from textwrap import dedent
|
||||
|
||||
python_directory = sys.executable
|
||||
|
||||
class ToolExecutor:
|
||||
|
||||
class ToolExecutor:
|
||||
_dir_initialized = False
|
||||
_lock = threading.Lock()
|
||||
|
||||
def __init__(self, sandbox=False):
|
||||
self.sandbox = sandbox
|
||||
if sandbox:
|
||||
|
|
@ -29,29 +30,22 @@ class ToolExecutor:
|
|||
self.user = None
|
||||
self.banned_keywords = CONFIG.get("SANDBOX_PYTHON_BANNED_KEYWORDS", 'nothing_is_banned').split(',');
|
||||
self.sandbox_so_path = f'{self.sandbox_path}/sandbox.so'
|
||||
with ToolExecutor._lock:
|
||||
self._init_dir()
|
||||
self._init_dir()
|
||||
|
||||
def _init_dir(self):
|
||||
if ToolExecutor._dir_initialized:
|
||||
# 只初始化一次
|
||||
return
|
||||
execute_file_path = os.path.join(self.sandbox_path, 'execute')
|
||||
os.makedirs(execute_file_path, 0o500, exist_ok=True)
|
||||
result_file_path = os.path.join(self.sandbox_path, 'result')
|
||||
os.makedirs(result_file_path, 0o300, exist_ok=True)
|
||||
if self.sandbox:
|
||||
os.system(f"chown {self.user}:root {self.sandbox_path}")
|
||||
os.system(f"chown -R {self.user}:root {execute_file_path}")
|
||||
os.system(f"chown -R {self.user}:root {result_file_path}")
|
||||
os.chmod(self.sandbox_path, 0o550)
|
||||
if CONFIG.get("SANDBOX_TMP_DIR_ENABLED", '0') == "1":
|
||||
tmp_dir_path = os.path.join(self.sandbox_path, 'tmp')
|
||||
os.makedirs(tmp_dir_path, 0o700, exist_ok=True)
|
||||
os.system(f"chown -R {self.user}:root {tmp_dir_path}")
|
||||
if os.path.exists(self.sandbox_so_path):
|
||||
os.chmod(self.sandbox_so_path, 0o440)
|
||||
try:
|
||||
if os.path.exists(self.sandbox_so_path):
|
||||
os.chmod(self.sandbox_so_path, 0o440)
|
||||
# 初始化host黑名单
|
||||
banned_hosts_file_path = f'{self.sandbox_path}/.SANDBOX_BANNED_HOSTS'
|
||||
if os.path.exists(banned_hosts_file_path):
|
||||
|
|
@ -74,13 +68,13 @@ class ToolExecutor:
|
|||
_id = str(uuid.uuid7())
|
||||
success = '{"code":200,"msg":"成功","data":exec_result}'
|
||||
err = '{"code":500,"msg":str(e),"data":None}'
|
||||
result_path = f'{self.sandbox_path}/result/{_id}.result'
|
||||
python_paths = CONFIG.get_sandbox_python_package_paths().split(',')
|
||||
_exec_code = f"""
|
||||
try:
|
||||
import os
|
||||
import sys
|
||||
import json
|
||||
import base64
|
||||
path_to_exclude = ['/opt/py3/lib/python3.11/site-packages', '/opt/maxkb-app/apps']
|
||||
sys.path = [p for p in sys.path if p not in path_to_exclude]
|
||||
sys.path += {python_paths}
|
||||
|
|
@ -92,21 +86,21 @@ try:
|
|||
for local in locals_v:
|
||||
globals_v[local] = locals_v[local]
|
||||
exec_result=f(**keywords)
|
||||
with open({result_path!a}, 'w') as file:
|
||||
file.write(json.dumps({success}, default=str))
|
||||
print(f"{_id}:" + base64.b64encode(json.dumps({success}, default=str).encode()).decode())
|
||||
except Exception as e:
|
||||
with open({result_path!a}, 'w') as file:
|
||||
file.write(json.dumps({err}))
|
||||
print(f"{_id}:" + base64.b64encode(json.dumps({err}, default=str).encode()).decode())
|
||||
"""
|
||||
if self.sandbox:
|
||||
subprocess_result = self._exec_sandbox(_exec_code, _id)
|
||||
subprocess_result = self._exec_sandbox(_exec_code)
|
||||
else:
|
||||
subprocess_result = self._exec(_exec_code)
|
||||
if subprocess_result.returncode == 1:
|
||||
raise Exception(subprocess_result.stderr)
|
||||
with open(result_path, 'r') as file:
|
||||
result = json.loads(file.read())
|
||||
os.remove(result_path)
|
||||
lines = subprocess_result.stdout.splitlines()
|
||||
result_line = [line for line in lines if line.startswith(_id)]
|
||||
if not result_line:
|
||||
raise Exception("No result found.")
|
||||
result = json.loads(base64.b64decode(result_line[0].split(":", 1)[1]).decode())
|
||||
if result.get('code') == 200:
|
||||
return result.get('data')
|
||||
raise Exception(result.get('msg'))
|
||||
|
|
@ -194,18 +188,16 @@ exec({dedent(code)!a})
|
|||
"""
|
||||
|
||||
def get_tool_mcp_config(self, code, params):
|
||||
code = self.generate_mcp_server_code(code, params)
|
||||
|
||||
_id = uuid.uuid7()
|
||||
code_path = f'{self.sandbox_path}/execute/{_id}.py'
|
||||
with open(code_path, 'w') as f:
|
||||
f.write(code)
|
||||
_code = self.generate_mcp_server_code(code, params)
|
||||
maxkb_logger.debug(f"Python code of mcp tool: {_code}")
|
||||
compressed_and_base64_encoded_code_str = base64.b64encode(gzip.compress(_code.encode())).decode()
|
||||
if self.sandbox:
|
||||
tool_config = {
|
||||
'command': 'su',
|
||||
'args': [
|
||||
'-s', sys.executable,
|
||||
'-c', f"exec(open('{code_path}', 'r').read())",
|
||||
'-c',
|
||||
f'import base64,gzip; exec(gzip.decompress(base64.b64decode(\'{compressed_and_base64_encoded_code_str}\')).decode())',
|
||||
self.user,
|
||||
],
|
||||
'cwd': self.sandbox_path,
|
||||
|
|
@ -217,24 +209,24 @@ exec({dedent(code)!a})
|
|||
else:
|
||||
tool_config = {
|
||||
'command': sys.executable,
|
||||
'args': [code_path],
|
||||
'args': f'import base64,gzip; exec(gzip.decompress(base64.b64decode(\'{compressed_and_base64_encoded_code_str}\')).decode())',
|
||||
'transport': 'stdio',
|
||||
}
|
||||
return _id, tool_config
|
||||
return tool_config
|
||||
|
||||
def _exec_sandbox(self, _code, _id):
|
||||
exec_python_file = f'{self.sandbox_path}/execute/{_id}.py'
|
||||
with open(exec_python_file, 'w') as file:
|
||||
file.write(_code)
|
||||
def _exec_sandbox(self, _code):
|
||||
kwargs = {'cwd': BASE_DIR}
|
||||
kwargs['env'] = {
|
||||
'LD_PRELOAD': self.sandbox_so_path,
|
||||
}
|
||||
maxkb_logger.debug(f"Sandbox execute code: {_code}")
|
||||
compressed_and_base64_encoded_code_str = base64.b64encode(gzip.compress(_code.encode())).decode()
|
||||
subprocess_result = subprocess.run(
|
||||
['su', '-s', python_directory, '-c', "exec(open('" + exec_python_file + "').read())", self.user],
|
||||
['su', '-s', python_directory, '-c',
|
||||
f'import base64,gzip; exec(gzip.decompress(base64.b64decode(\'{compressed_and_base64_encoded_code_str}\')).decode())',
|
||||
self.user],
|
||||
text=True,
|
||||
capture_output=True, **kwargs)
|
||||
os.remove(exec_python_file)
|
||||
return subprocess_result
|
||||
|
||||
def validate_banned_keywords(self, code_str):
|
||||
|
|
|
|||
|
|
@ -2,8 +2,8 @@
|
|||
|
||||
if [ ! -d /opt/maxkb/logs ]; then
|
||||
mkdir -p /opt/maxkb/logs
|
||||
chmod 700 /opt/maxkb/logs
|
||||
fi
|
||||
chmod -R 700 /opt/maxkb/logs
|
||||
if [ ! -d /opt/maxkb/local ]; then
|
||||
mkdir -p /opt/maxkb/local
|
||||
chmod 700 /opt/maxkb/local
|
||||
|
|
|
|||
Loading…
Reference in New Issue