diff --git a/apps/application/flow/step_node/tool_lib_node/impl/base_tool_lib_node.py b/apps/application/flow/step_node/tool_lib_node/impl/base_tool_lib_node.py index 04ea0d0b..a99f4f5b 100644 --- a/apps/application/flow/step_node/tool_lib_node/impl/base_tool_lib_node.py +++ b/apps/application/flow/step_node/tool_lib_node/impl/base_tool_lib_node.py @@ -17,14 +17,14 @@ from application.flow.i_step_node import NodeResult from application.flow.step_node.tool_lib_node.i_tool_lib_node import IToolLibNode from common.database_model_manage.database_model_manage import DatabaseModelManage from common.exception.app_exception import AppApiException -from common.utils.function_code import FunctionExecutor from common.utils.rsa_util import rsa_long_decrypt +from common.utils.tool_code import ToolExecutor from maxkb.const import CONFIG from system_manage.models import AuthTargetType from system_manage.serializers.user_resource_permission import UserResourcePermissionSerializer from tools.models import Tool -function_executor = FunctionExecutor(CONFIG.get('SANDBOX')) +function_executor = ToolExecutor(CONFIG.get('SANDBOX')) def write_context(step_variable: Dict, global_variable: Dict, node, workflow): diff --git a/apps/application/flow/step_node/tool_node/impl/base_tool_node.py b/apps/application/flow/step_node/tool_node/impl/base_tool_node.py index 67547fb3..9f03a7a8 100644 --- a/apps/application/flow/step_node/tool_node/impl/base_tool_node.py +++ b/apps/application/flow/step_node/tool_node/impl/base_tool_node.py @@ -12,10 +12,10 @@ from typing import Dict from application.flow.i_step_node import NodeResult from application.flow.step_node.tool_node.i_tool_node import IToolNode -from common.utils.function_code import FunctionExecutor +from common.utils.tool_code import ToolExecutor from maxkb.const import CONFIG -function_executor = FunctionExecutor(CONFIG.get('SANDBOX')) +function_executor = ToolExecutor(CONFIG.get('SANDBOX')) def write_context(step_variable: Dict, global_variable: Dict, node, workflow): diff --git a/apps/common/utils/function_code.py b/apps/common/utils/function_code.py deleted file mode 100644 index a4bd9751..00000000 --- a/apps/common/utils/function_code.py +++ /dev/null @@ -1,99 +0,0 @@ -# coding=utf-8 -""" - @project: MaxKB - @Author:虎 - @file: function_code.py - @date:2024/8/7 16:11 - @desc: -""" -import os -import pickle -import subprocess -import sys -import uuid_utils.compat as uuid -from textwrap import dedent - -from maxkb.const import BASE_DIR -from maxkb.const import PROJECT_DIR - -python_directory = sys.executable - - -class FunctionExecutor: - def __init__(self, sandbox=False): - self.sandbox = sandbox - if sandbox: - self.sandbox_path = '/opt/maxkb-app/sandbox' - self.user = 'sandbox' - else: - self.sandbox_path = os.path.join(PROJECT_DIR, 'data', 'sandbox') - self.user = None - self._createdir() - if self.sandbox: - os.system(f"chown -R {self.user}:root {self.sandbox_path}") - - def _createdir(self): - old_mask = os.umask(0o077) - try: - os.makedirs(self.sandbox_path, 0o700, exist_ok=True) - os.makedirs(os.path.join(self.sandbox_path, 'execute'), 0o700, exist_ok=True) - os.makedirs(os.path.join(self.sandbox_path, 'result'), 0o700, exist_ok=True) - finally: - os.umask(old_mask) - - 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' - _exec_code = f""" -try: - import os - env = dict(os.environ) - for key in list(env.keys()): - if key in os.environ and (key.startswith('MAXKB') or key.startswith('POSTGRES') or key.startswith('PG') or key.startswith('REDIS')): - del os.environ[key] - locals_v={'{}'} - keywords={keywords} - globals_v=globals() - exec({dedent(code_str)!a}, globals_v, locals_v) - f_name, f = locals_v.popitem() - for local in locals_v: - globals_v[local] = locals_v[local] - exec_result=f(**keywords) - import pickle - with open({result_path!a}, 'wb') as file: - file.write(pickle.dumps({success})) -except Exception as e: - with open({result_path!a}, 'wb') as file: - file.write(pickle.dumps({err})) -""" - if self.sandbox: - subprocess_result = self._exec_sandbox(_exec_code, _id) - else: - subprocess_result = self._exec(_exec_code) - if subprocess_result.returncode == 1: - raise Exception(subprocess_result.stderr) - with open(result_path, 'rb') as file: - result = pickle.loads(file.read()) - os.remove(result_path) - if result.get('code') == 200: - return result.get('data') - raise Exception(result.get('msg')) - - 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) - os.system(f"chown {self.user}:root {exec_python_file}") - kwargs = {'cwd': BASE_DIR} - subprocess_result = subprocess.run( - ['su', '-s', python_directory, '-c', "exec(open('" + exec_python_file + "').read())", self.user], - text=True, - capture_output=True, **kwargs) - os.remove(exec_python_file) - return subprocess_result - - @staticmethod - def _exec(_code): - return subprocess.run([python_directory, '-c', _code], text=True, capture_output=True)