refactor: avoid using temp files in sandbox.

Author: liqiang-fit2cloud

apps/application/chat_pipeline/step/chat_step/impl/base_chat_step.py

@@ -200,14 +200,6 @@ def execute(self, message_list: List[BaseMessage],
                                       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 @@ def _handle_mcp_request(self, mcp_enable, tool_enable, mcp_source, mcp_servers,
         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 @@ def _handle_mcp_request(self, mcp_enable, tool_enable, mcp_source, mcp_servers,
                         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:

apps/application/flow/step_node/ai_chat_step_node/impl/base_chat_node.py

@@ -233,7 +233,6 @@ def _handle_mcp_request(self, mcp_enable, tool_enable, mcp_source, mcp_servers,
         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 @@ def _handle_mcp_request(self, mcp_enable, tool_enable, mcp_source, mcp_servers,
                         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 @@ def reset_message_list(message_list: List[BaseMessage], answer_text):
         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,

apps/common/utils/tool_code.py

@@ -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 @@ def __init__(self, sandbox=False):
             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 @@ def exec_code(self, code_str, keywords):
         _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 @@ def exec_code(self, code_str, keywords):
     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 @@ def generate_mcp_server_code(self, code_str, params):
 """
 
     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 @@ def get_tool_mcp_config(self, code, params):
         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):

installer/start-maxkb.sh

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