panyy
/
UnisKB
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

chore: replace pickle with json for data serialization in tool_code.py

v3.2
CaptainB 2025-09-04 16:42:59 +08:00
parent d94871b9c3
commit 817d5bcd5a
1 changed files with 9 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
apps/common/utils/tool_code.py
View File

@ -1,7 +1,7 @@
# coding=utf-8 # coding=utf-8
import ast import ast
import os import os
import pickle import json
import subprocess import subprocess
import sys import sys
from textwrap import dedent from textwrap import dedent
@ -40,7 +40,7 @@ class ToolExecutor:
def exec_code(self, code_str, keywords): def exec_code(self, code_str, keywords):
self.validate_banned_keywords(code_str) self.validate_banned_keywords(code_str)
_id = str(uuid.uuid7()) _id = str(uuid.uuid7())
success = '{"code":200,"msg":"成功","data":exec_result}' success = '{"code":200,"msg":"成功","data":json.dumps(exec_result, default=str)}'
err = '{"code":500,"msg":str(e),"data":None}' err = '{"code":500,"msg":str(e),"data":None}'
result_path = f'{self.sandbox_path}/result/{_id}.result' result_path = f'{self.sandbox_path}/result/{_id}.result'
python_paths = CONFIG.get_sandbox_python_package_paths().split(',') python_paths = CONFIG.get_sandbox_python_package_paths().split(',')
@ -48,7 +48,7 @@ class ToolExecutor:
try: try:
import os import os
import sys import sys
import pickle import json
path_to_exclude = ['/opt/py3/lib/python3.11/site-packages', '/opt/maxkb-app/apps'] 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 = [p for p in sys.path if p not in path_to_exclude]
sys.path += {python_paths} sys.path += {python_paths}
@ -64,11 +64,11 @@ try:
for local in locals_v: for local in locals_v:
globals_v[local] = locals_v[local] globals_v[local] = locals_v[local]
exec_result=f(**keywords) exec_result=f(**keywords)
with open({result_path!a}, 'wb') as file: with open({result_path!a}, 'w') as file:
file.write(pickle.dumps({success})) file.write(json.dumps({success}))
except Exception as e: except Exception as e:
with open({result_path!a}, 'wb') as file: with open({result_path!a}, 'w') as file:
file.write(pickle.dumps({err})) file.write(json.dumps({err}))
""" """
if self.sandbox: if self.sandbox:
subprocess_result = self._exec_sandbox(_exec_code, _id) subprocess_result = self._exec_sandbox(_exec_code, _id)
@ -76,8 +76,8 @@ except Exception as e:
subprocess_result = self._exec(_exec_code) subprocess_result = self._exec(_exec_code)
if subprocess_result.returncode == 1: if subprocess_result.returncode == 1:
raise Exception(subprocess_result.stderr) raise Exception(subprocess_result.stderr)
with open(result_path, 'rb') as file: with open(result_path, 'r') as file:
result = pickle.loads(file.read()) result = json.loads(file.read())
os.remove(result_path) os.remove(result_path)
if result.get('code') == 200: if result.get('code') == 200:
return result.get('data') return result.get('data')