# coding=utf-8 import os import subprocess import sys import uuid_utils.compat as uuid from textwrap import dedent from diskcache import Cache from maxkb.const import BASE_DIR from maxkb.const import PROJECT_DIR python_directory = '/opt/maxkb-app/sandbox/py3/bin/python' class ToolExecutor: 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) 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}' path = r'' + self.sandbox_path + '' _exec_code = f""" try: import os import sys path_to_exclude = '/opt/py3/lib/python3.11/site-packages' sys.path = [p for p in sys.path if p != path_to_exclude] 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')): 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) from diskcache import Cache cache = Cache({path!a}) cache.set({_id!a},{success}) except Exception as e: from diskcache import Cache cache = Cache({path!a}) cache.set({_id!a},{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) cache = Cache(self.sandbox_path) result = cache.get(_id) cache.delete(_id) 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}/{_id}.py' with open(exec_python_file, 'w') as file: file.write(_code) os.system(f"chown {self.user}:{self.user} {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([sys.executable, '-c', _code], text=True, capture_output=True)