149 lines
5.3 KiB
Python
149 lines
5.3 KiB
Python
"""
|
||
tools/code_executor.py
|
||
代码执行工具 —— 在沙箱中执行 Python 代码片段
|
||
配置通过 settings.tools['code_executor'] 读取
|
||
"""
|
||
|
||
import io
|
||
import sys
|
||
import textwrap
|
||
import time
|
||
import traceback
|
||
from contextlib import redirect_stderr, redirect_stdout
|
||
|
||
from config.settings import settings
|
||
from tools.base_tool import BaseTool
|
||
from utils.logger import get_logger
|
||
|
||
logger = get_logger("TOOL.CodeExecutor")
|
||
|
||
|
||
def _cfg(key: str, fallback=None):
|
||
return settings.tools['code_executor'].get(key, fallback)
|
||
|
||
|
||
class Tool(BaseTool):
|
||
name = "code_executor"
|
||
description = (
|
||
"在安全沙箱中执行 Python 代码片段，返回标准输出和执行结果。"
|
||
"适用于数据处理、计算、格式转换等任务。"
|
||
"注意：沙箱模式下禁止文件系统写入、网络访问和系统调用。"
|
||
)
|
||
parameters = {
|
||
"type": "object",
|
||
"properties": {
|
||
"code": {
|
||
"type": "string",
|
||
"description": "要执行的 Python 代码字符串",
|
||
},
|
||
"timeout": {
|
||
"type": "integer",
|
||
"description": "执行超时秒数（默认来自 config.yaml code_executor.timeout）",
|
||
},
|
||
},
|
||
"required": ["code"],
|
||
}
|
||
|
||
# 沙箱模式下禁止的模块和函数
|
||
_FORBIDDEN_SANDBOX = [
|
||
"import os", "import sys", "import subprocess",
|
||
"import socket", "import requests", "import httpx",
|
||
"import shutil", "open(", "__import__",
|
||
"exec(", "eval(", "compile(",
|
||
]
|
||
|
||
def execute(self, code: str = "", timeout: int | None = None, **_) -> str:
|
||
if not code or not code.strip():
|
||
return "❌ 参数错误: code 不能为空"
|
||
|
||
sandbox = _cfg('sandbox', True)
|
||
t = timeout or _cfg('timeout', 5)
|
||
code = textwrap.dedent(code)
|
||
|
||
logger.info(
|
||
f"🐍 执行代码 sandbox={sandbox} timeout={t}s "
|
||
f"[config timeout={_cfg('timeout')}s sandbox={_cfg('sandbox')}]\n"
|
||
f" 代码预览: {code[:100]}"
|
||
)
|
||
|
||
# 沙箱安全检查
|
||
if sandbox:
|
||
err = self._sandbox_check(code)
|
||
if err:
|
||
return err
|
||
|
||
# 使用线程超时执行
|
||
return self._run_with_timeout(code, t)
|
||
|
||
def _run_with_timeout(self, code: str, timeout: int) -> str:
|
||
"""在独立线程中执行代码，超时则终止"""
|
||
import threading
|
||
|
||
result_box: list[str] = []
|
||
error_box: list[str] = []
|
||
|
||
def _run():
|
||
stdout_buf = io.StringIO()
|
||
stderr_buf = io.StringIO()
|
||
local_ns: dict = {}
|
||
start = time.time()
|
||
try:
|
||
with redirect_stdout(stdout_buf), redirect_stderr(stderr_buf):
|
||
exec(code, {"__builtins__": __builtins__}, local_ns) # noqa: S102
|
||
elapsed = time.time() - start
|
||
stdout = stdout_buf.getvalue()
|
||
stderr = stderr_buf.getvalue()
|
||
# 尝试获取最后一个表达式的值
|
||
last_val = ""
|
||
lines = [l.strip() for l in code.strip().splitlines() if l.strip()]
|
||
if lines:
|
||
last_line = lines[-1]
|
||
if not last_line.startswith(("#", "print", "import", "from",
|
||
"def ", "class ", "if ", "for ",
|
||
"while ", "try:", "with ")):
|
||
try:
|
||
val = eval(last_line, {"__builtins__": __builtins__}, local_ns) # noqa: S307
|
||
if val is not None:
|
||
last_val = f"\n返回值: {repr(val)}"
|
||
except Exception:
|
||
pass
|
||
|
||
output = stdout + (f"\n[stderr]\n{stderr}" if stderr else "") + last_val
|
||
result_box.append(
|
||
f"✅ 执行成功 耗时={elapsed:.3f}s\n"
|
||
f"{'─' * 40}\n"
|
||
f"{output.strip() or '(无输出)'}"
|
||
)
|
||
except Exception:
|
||
elapsed = time.time() - start
|
||
tb = traceback.format_exc()
|
||
error_box.append(
|
||
f"❌ 执行错误 耗时={elapsed:.3f}s\n"
|
||
f"{'─' * 40}\n{tb}"
|
||
)
|
||
|
||
thread = threading.Thread(target=_run, daemon=True)
|
||
thread.start()
|
||
thread.join(timeout=timeout)
|
||
|
||
if thread.is_alive():
|
||
return (
|
||
f"⏰ 执行超时（>{timeout}s）\n"
|
||
f" 请增大 config.yaml → tools.code_executor.timeout\n"
|
||
f" 或优化代码逻辑"
|
||
)
|
||
|
||
if error_box:
|
||
return error_box[0]
|
||
return result_box[0] if result_box else "❌ 执行失败（未知错误）"
|
||
|
||
def _sandbox_check(self, code: str) -> str | None:
|
||
"""沙箱模式下的静态安全检查"""
|
||
for forbidden in self._FORBIDDEN_SANDBOX:
|
||
if forbidden in code:
|
||
return (
|
||
f"❌ 沙箱限制: 代码包含禁止操作 '{forbidden}'\n"
|
||
f" 如需完整权限请在 config.yaml → "
|
||
f"tools.code_executor.sandbox 设置为 false"
|
||
)
|
||
return None |