手写 AI 代码执行沙箱：从零实现安全的代码运行环境

import os
import sys
import subprocess
import tempfile
import time
import signal
from pathlib import Path
from typing import Optional, Dict, Any, List
from dataclasses import dataclass, field

@dataclass
class SandboxResult:
    """沙箱执行结果"""
    stdout: str = ""
    stderr: str = ""
    return_code: int = -1
    execution_time: float = 0.0
    timed_out: bool = False
    error: Optional[str] = None
    resource_usage: Dict[str, Any] = field(default_factory=dict)

class ProcessSandbox:
    """基于进程隔离的代码沙箱"""

    def __init__(
        self,
        work_dir: Optional[str] = None,
        timeout: int = 30,
        max_memory_mb: int = 256,
        allow_network: bool = False,
        allowed_paths: Optional[List[str]] = None,
        forbidden_modules: Optional[List[str]] = None
    ):
        self.work_dir = work_dir or tempfile.mkdtemp(prefix="sandbox_")
        self.timeout = timeout
        self.max_memory_mb = max_memory_mb
        self.allow_network = allow_network
        self.allowed_paths = allowed_paths or [self.work_dir]
        self.forbidden_modules = forbidden_modules or [
            "os", "subprocess", "shutil", "ctypes",
            "multiprocessing", "threading", "socket",
            "signal", "sys", "importlib"
        ]
        # 确保工作目录存在
        Path(self.work_dir).mkdir(parents=True, exist_ok=True)

    def _build_restricted_runner(self, code: str) -> str:
        """
        构建受限运行器代码
        在子进程中执行代码前，先禁用危险模块
        """
        # 构建模块拦截
        module_block = ""
        for mod in self.forbidden_modules:
            # 构建一个 ModuleNotFoundError 或 ImportError
            module_block += f"""
import sys
if '{mod}' in sys.modules:
    del sys.modules['{mod}']
class _BlockedModule:
    def __init__(self, name):
        raise ImportError(f"Module '{{name}}' is blocked in sandbox")
sys.modules['{mod}'] = _BlockedModule('{mod}')
"""

        # 在代码执行前注入限制
        runner_code = f"""
import sys, builtins

# ── 1. 限制危险内建函数 ──────────────────
_safe_builtins = {{}}
for _name in dir(builtins):
    if _name.startswith('_') and _name not in ('__import__',):
        continue
    _safe_builtins[_name] = getattr(builtins, _name)

# 移除危险函数
_safe_builtins.pop('exec', None)
_safe_builtins.pop('eval', None)
_safe_builtins.pop('compile', None)
_safe_builtins.pop('__import__', None)
_safe_builtins.pop('open', None)
_safe_builtins.pop('input', None)

# ── 2. 模块拦截 ──────────────────────────
{module_block}

# ── 3. 替换 builtins ──────────────────────
import types
_restricted_mod = types.ModuleType('__restricted__')
_restricted_mod.__dict__.update(_safe_builtins)

# ── 4. 执行用户代码 ───────────────────────
_code = {repr(code)}

# 捕获输出
import io
_stdout_capture = io.StringIO()
_stderr_capture = io.StringIO()
_old_stdout = sys.stdout
_old_stderr = sys.stderr
sys.stdout = _stdout_capture
sys.stderr = _stderr_capture

try:
    exec(_code, _restricted_mod.__dict__)
except Exception as _e:
    import traceback
    traceback.print_exc(file=_stderr_capture)

# 恢复输出
sys.stdout = _old_stdout
sys.stderr = _old_stderr

# 输出结果
print("__stdout__" + _stdout_capture.getvalue() + "__end_stdout__")
print("__stderr__" + _stderr_capture.getvalue() + "__end_stderr__", file=sys.stderr)
"""
        return runner_code

    def run(self, code: str) -> SandboxResult:
        """
        在沙箱中执行代码
        """
        start_time = time.time()
        result = SandboxResult()

        try:
            # 创建临时文件
            runner_path = os.path.join(self.work_dir, "_runner.py")
            with open(runner_path, 'w', encoding='utf-8') as f:
                f.write(self._build_restricted_runner(code))

            # 使用 subprocess 在隔离进程中运行
            proc = subprocess.Popen(
                [sys.executable, runner_path],
                cwd=self.work_dir,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                env={
                    **os.environ,
                    # 限制 Python 环境
                    "PYTHONPATH": "",
                    "PYTHONDONTWRITEBYTECODE": "1",
                    # 删除可能敏感的环境变量
                    **{k: "" for k in os.environ if k.startswith(("AWS_", "AZURE_", "GCP_", "DB_", "SECRET", "TOKEN", "PASSWORD", "API_KEY"))}
                },
                preexec_fn=self._set_process_limits  # 进程资源限制
            )

            try:
                stdout, stderr = proc.communicate(timeout=self.timeout)
                result.execution_time = time.time() - start_time
            except subprocess.TimeoutExpired:
                proc.kill()
                stdout, stderr = proc.communicate()
                result.timed_out = True
                result.execution_time = self.timeout

            # 解析输出
            result.stdout = stdout.decode('utf-8', errors='replace') if stdout else ""
            result.stderr = stderr.decode('utf-8', errors='replace') if stderr else ""
            result.return_code = proc.returncode

            # 提取实际输出（去掉 __stdout__/__stderr__ 标记）
            if "__stdout__" in result.stdout and "__end_stdout__" in result.stdout:
                result.stdout = result.stdout.split("__stdout__")[1].split("__end_stdout__")[0]
            if "__stderr__" in result.stderr and "__end_stderr__" in result.stderr:
                result.stderr = result.stderr.split("__stderr__")[1].split("__end_stderr__")[0]

        except Exception as e:
            result.error = str(e)

        return result

    def _set_process_limits(self):
        """
        在子进程中设置资源限制
        """
        try:
            import resource
            # CPU 时间限制
            resource.setrlimit(resource.RLIMIT_CPU, (5, 5))
            # 内存限制
            mem_bytes = self.max_memory_mb * 1024 * 1024
            resource.setrlimit(resource.RLIMIT_AS, (mem_bytes, mem_bytes))
            # 文件大小限制（禁止大文件写入）
            resource.setrlimit(resource.RLIMIT_FSIZE, (1024 * 1024, 1024 * 1024))
            # 子进程数限制
            resource.setrlimit(resource.RLIMIT_NPROC, (10, 10))
        except (ImportError, ValueError):
            pass

    def cleanup(self):
        """清理沙箱工作目录"""
        import shutil
        shutil.rmtree(self.work_dir, ignore_errors=True)

import ast
from typing import Set, Tuple, List

class ASTSecurityAnalyzer:
    """基于 AST 的代码安全分析"""

    # 危险函数调用模式
    DANGEROUS_CALLS: Set[str] = {
        # 文件系统
        "os.system", "os.popen", "os.remove", "os.rmdir",
        "os.unlink", "os.chmod", "os.chown", "os.kill",
        "shutil.rmtree", "shutil.move", "shutil.copy",
        "subprocess.call", "subprocess.Popen", "subprocess.run",
        "subprocess.check_call", "subprocess.check_output",
        # 网络
        "socket.connect", "socket.send", "socket.sendto",
        "urllib.request.urlopen", "requests.get", "requests.post",
        # 代码执行
        "exec", "eval", "compile", "__import__",
        # 系统操作
        "os.execv", "os.execl", "os.execve", "os.fork",
        "ctypes.CDLL", "ctypes.c_char_p",
    }

    # 危险属性访问
    DANGEROUS_ATTRS: Set[str] = {
        "__class__", "__base__", "__subclasses__",
        "__globals__", "__builtins__", "__code__",
    }

    # 允许的安全 builtins
    SAFE_BUILTINS: Set[str] = {
        "abs", "all", "any", "ascii", "bin", "bool", "bytes",
        "callable", "chr", "complex", "dict", "dir", "divmod",
        "enumerate", "filter", "float", "format", "frozenset",
        "getattr", "hasattr", "hash", "hex", "id", "int",
        "isinstance", "issubclass", "iter", "len", "list",
        "map", "max", "min", "next", "object", "oct", "ord",
        "pow", "print", "property", "range", "repr", "reversed",
        "round", "set", "slice", "sorted", "str", "sum",
        "tuple", "type", "vars", "zip",
    }

    def analyze(self, code: str) -> List[str]:
        """
        分析代码安全性
        返回: 检测到的安全问题列表（空列表 = 安全）
        """
        issues = []

        try:
            tree = ast.parse(code)
        except SyntaxError as e:
            return [f"语法错误: {e}"]

        for node in ast.walk(tree):
            # 检查危险函数调用
            if isinstance(node, ast.Call):
                issue = self._check_call(node)
                if issue:
                    issues.append(issue)

            # 检查危险属性访问
            if isinstance(node, ast.Attribute):
                issue = self._check_attribute(node)
                if issue:
                    issues.append(issue)

            # 检查 import
            if isinstance(node, ast.Import):
                for alias in node.names:
                    if alias.name in ("os", "subprocess", "shutil", "ctypes", "socket"):
                        issues.append(f"禁止导入危险模块: {alias.name}")

            if isinstance(node, ast.ImportFrom):
                module = node.module or ""
                if module in ("os", "subprocess", "shutil", "ctypes", "socket"):
                    issues.append(f"禁止从危险模块导入: {module}")

            # 检查 with 语句中的 open 调用
            if isinstance(node, ast.With):
                for item in node.items:
                    if isinstance(item.context_expr, ast.Call):
                        call = item.context_expr
                        if isinstance(call.func, ast.Name) and call.func.id == "open":
                            issues.append("文件操作被禁止: open()")

            # 检查装饰器（隐藏恶意操作的一种常见方式）
            if isinstance(node, ast.FunctionDef) and node.decorator_list:
                for dec in node.decorator_list:
                    if isinstance(dec, ast.Attribute) and dec.attr in ("__class__", "__subclasses__"):
                        issues.append(f"禁止使用装饰器访问内省属性: {dec.attr}")

            # 检查 while True 死循环风险
            if isinstance(node, ast.While):
                if isinstance(node.test, ast.Constant) and node.test.value is True:
                    # 检查循环体中是否有 break
                    has_break = any(
                        isinstance(n, ast.Break)
                        for n in ast.walk(node)
                    )
                    if not has_break:
                        issues.append("检测到无 break 的 while True 循环（可能导致资源耗尽）")

        return issues

    def _check_call(self, node: ast.Call) -> Optional[str]:
        """检查函数调用是否危险"""
        # 解析完整的函数名
        full_name = self._get_call_name(node)
        if full_name and full_name in self.DANGEROUS_CALLS:
            return f"禁止调用危险函数: {full_name}"

        # 检查是否引用了危险的 dunder 方法
        if isinstance(node.func, ast.Attribute):
            if node.func.attr in ("__subclasses__", "__bases__", "__globals__"):
                return f"禁止访问内省属性: {node.func.attr}"

        return None

    def _check_attribute(self, node: ast.Attribute) -> Optional[str]:
        """检查属性访问是否危险"""
        if node.attr in self.DANGEROUS_ATTRS:
            return f"禁止访问危险属性: {node.attr}"
        return None

    def _get_call_name(self, node: ast.Call) -> str:
        """获取函数调用的完整名称"""
        if isinstance(node.func, ast.Name):
            return node.func.id
        elif isinstance(node.func, ast.Attribute):
            parts = []
            current = node.func
            while isinstance(current, ast.Attribute):
                parts.append(current.attr)
                current = current.value
            if isinstance(current, ast.Name):
                parts.append(current.id)
            return ".".join(reversed(parts))
        return ""

    def is_safe(self, code: str) -> Tuple[bool, List[str]]:
        """
        快速安全检测
        返回: (是否安全, 问题列表)
        """
        issues = self.analyze(code)
        return len(issues) == 0, issues

class WhitelistAnalyzer:
    """白名单模式的安全分析器"""

    ALLOWED_MODULES = {
        "math", "random", "datetime", "time", "json",
        "collections", "itertools", "functools", "operator",
        "re", "string", "typing", "enum", "decimal",
        "fractions", "statistics", "uuid", "hashlib", "base64",
        "textwrap", "pprint", "copy", "bisect", "heapq",
    }

    ALLOWED_FUNCTIONS = {
        "print", "len", "range", "list", "dict", "set", "tuple",
        "str", "int", "float", "bool", "bytes", "bytearray",
        "sum", "min", "max", "abs", "round", "sorted",
        "enumerate", "zip", "map", "filter", "reversed",
        "type", "isinstance", "hasattr", "getattr", "setattr",
        "chr", "ord", "hex", "bin", "oct", "format",
        "any", "all", "callable", "hash", "id", "repr",
        "open",  # 特殊处理：只允许读模式
    }

    def analyze_with_whitelist(self, code: str) -> List[str]:
        """白名单检查"""
        issues = []

        try:
            tree = ast.parse(code)
        except SyntaxError as e:
            return [f"语法错误: {e}"]

        for node in ast.walk(tree):
            if isinstance(node, ast.Import):
                for alias in node.names:
                    if alias.name not in self.ALLOWED_MODULES:
                        issues.append(f"模块不在白名单中: {alias.name}")

            elif isinstance(node, ast.ImportFrom):
                if node.module and node.module not in self.ALLOWED_MODULES:
                    issues.append(f"导入的模块不在白名单中: {node.module}")

            elif isinstance(node, ast.Call):
                if isinstance(node.func, ast.Name):
                    if node.func.id not in self.ALLOWED_FUNCTIONS:
                        issues.append(f"函数调用不在白名单中: {node.func.id}")

                # 对 open 的特殊检查
                if isinstance(node.func, ast.Name) and node.func.id == "open":
                    if len(node.args) >= 2:
                        mode_arg = node.args[1]
                        if isinstance(mode_arg, ast.Constant) and 'w' in str(mode_arg.value):
                            issues.append("禁止以写模式打开文件")
                        elif isinstance(mode_arg, ast.Constant) and 'a' in str(mode_arg.value):
                            issues.append("禁止以追加模式打开文件")

        return issues

import os
import tempfile
from pathlib import Path

class VirtualFileSystem:
    """
    虚拟文件系统：将代码的文件操作重定向到沙箱目录

    实现方式：
    1. 创建一个临时目录作为虚拟根目录
    2. 代码中所有文件操作路径被重写，前缀加上虚拟根目录
    3. 可以配置白名单路径，允许只读访问某些真实目录
    """

    def __init__(
        self,
        sandbox_root: Optional[str] = None,
        read_whitelist: Optional[list] = None
    ):
        self.sandbox_root = sandbox_root or tempfile.mkdtemp(prefix="vfs_")
        self.read_whitelist = read_whitelist or []

        # 在沙箱根目录下创建常用子目录
        Path(self.sandbox_root, "tmp").mkdir(parents=True, exist_ok=True)
        Path(self.sandbox_root, "home").mkdir(parents=True, exist_ok=True)
        Path(self.sandbox_root, "data").mkdir(parents=True, exist_ok=True)

    def sandbox_path(self, original_path: str) -> str:
        """
        将原始路径映射到沙箱路径

        - 绝对路径如 /tmp/x → /sandbox_root/tmp/x
        - 相对路径如 ./data → /sandbox_root/data
        - 白名单路径保留原路径
        """
        # 检查白名单
        for allowed in self.read_whitelist:
            if original_path.startswith(allowed):
                return original_path

        # 去除盘符和多余分隔符
        clean_path = os.path.normpath(original_path)

        if os.path.isabs(original_path):
            # 绝对路径：去掉开头的 /，拼接到沙箱根
            relative = clean_path.lstrip('/')
            return os.path.join(self.sandbox_root, relative)
        else:
            # 相对路径：直接拼接到沙箱根 /data 下
            return os.path.join(self.sandbox_root, "data", clean_path)

    def create_vfs_runner(self, code: str) -> str:
        """
        创建带虚拟文件系统的运行器代码
        通过重写 os.path 和内置 open 函数来实现
        """
        vfs_path = self.sandbox_root.replace('\\', '\\\\')

        return f'''
import os
import io
from pathlib import Path

# ── 虚拟文件系统 ──────────────────────
SANDBOX_ROOT = r"{vfs_path}"
READ_WHITELIST = {repr(self.read_whitelist)}

_sandbox_open = open

def _vfs_open(file, mode='r', *args, **kwargs):
    """重写 open：将路径映射到沙箱内"""
    # 如果是读模式，检查白名单
    if 'r' in mode and not ('w' in mode or 'a' in mode or '+' in mode):
        for allowed in READ_WHITELIST:
            if str(file).startswith(allowed):
                return _sandbox_open(file, mode, *args, **kwargs)

    # 映射到沙箱路径
    clean = os.path.normpath(str(file))
    if os.path.isabs(clean):
        sandbox_file = os.path.join(SANDBOX_ROOT, clean.lstrip('/'))
    else:
        sandbox_file = os.path.join(SANDBOX_ROOT, "data", clean)

    # 确保父目录存在
    os.makedirs(os.path.dirname(sandbox_file), exist_ok=True)

    return _sandbox_open(sandbox_file, mode, *args, **kwargs)

# 替换内建 open
import builtins
builtins.open = _vfs_open

# ── 用户代码 ──────────────────────────
{code}
'''

    def cleanup(self):
        """清理虚拟文件系统"""
        import shutil
        shutil.rmtree(self.sandbox_root, ignore_errors=True)

import subprocess
import platform

class NetworkController:
    """网络访问控制器"""

    def __init__(self, sandbox_uid: int = None):
        self.sandbox_uid = sandbox_uid
        self.rules_added = False

    def block_network(self, sandbox_uid: int):
        """
        阻止指定用户 ID 的网络访问

        使用 nftables 或 iptables 设置 egress 过滤
        """
        system = platform.system()

        if system == "Linux":
            try:
                # 尝试使用 iptables
                subprocess.run([
                    "iptables", "-A", "OUTPUT",
                    "-m", "owner", "--uid-owner", str(sandbox_uid),
                    "-j", "DROP"
                ], check=True, capture_output=True)
                self.rules_added = True
                return True
            except subprocess.CalledProcessError:
                pass

            try:
                # 尝试使用 nftables
                subprocess.run([
                    "nft", "add", "rule", "inet", "filter", "output",
                    "skuid", str(sandbox_uid), "drop"
                ], check=True, capture_output=True)
                self.rules_added = True
                return True
            except (subprocess.CalledProcessError, FileNotFoundError):
                pass

        return False

    def restore_network(self, sandbox_uid: int):
        """恢复网络访问"""
        if not self.rules_added:
            return

        try:
            subprocess.run([
                "iptables", "-D", "OUTPUT",
                "-m", "owner", "--uid-owner", str(sandbox_uid),
                "-j", "DROP"
            ], check=True, capture_output=True)
        except subprocess.CalledProcessError:
            pass

import socket
import select

class SocketProxy:
    """
    Socket 代理：监控并限制网络访问

    工作原理：
    1. 在沙箱中设置 http_proxy/https_proxy 环境变量指向此代理
    2. 代理检查每个请求的目标地址和端口
    3. 白名单内的请求放行，其他拒绝
    """

    def __init__(self, host: str = "127.0.0.1", port: int = 0):
        self.host = host
        self.port = port
        self.server = None
        self.running = False

        # 默认网络白名单
        self.domain_whitelist = {
            "pypi.org", "files.pythonhosted.org",
            "api.github.com", "raw.githubusercontent.com",
        }
        self.block_all = False

    def start(self):
        """启动代理服务器"""
        self.server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.server.bind((self.host, self.port))
        self.server.listen(5)
        self.port = self.server.getsockname()[1]
        self.running = True
        return self.host, self.port

    def stop(self):
        """停止代理"""
        self.running = False
        if self.server:
            self.server.close()

    def get_env(self) -> dict:
        """获取代理环境变量"""
        proxy_url = f"http://{self.host}:{self.port}"
        return {
            "http_proxy": proxy_url,
            "https_proxy": proxy_url,
            "no_proxy": ""
        }

    def is_allowed(self, host: str, port: int) -> bool:
        """检查是否允许访问目标"""
        if self.block_all:
            return False
        return host in self.domain_whitelist and port == 443

import os
import pwd
import grp
import tempfile
from enum import Enum
from typing import Optional, List, Callable

class SandboxLevel(Enum):
    """沙箱安全等级"""
    LOW = "low"           # 基础：进程隔离 + 超时
    MEDIUM = "medium"     # 中等：+ AST 预检 + 文件系统虚拟化
    HIGH = "high"         # 高级：+ 资源限制 + 网络封锁
    PARANOID = "paranoid" # 极限：+ 独立用户 + 全部封锁

class UnifiedSandbox:
    """统一沙箱接口"""

    def __init__(
        self,
        level: SandboxLevel = SandboxLevel.MEDIUM,
        timeout: int = 30,
        max_memory_mb: int = 256,
        work_dir: Optional[str] = None
    ):
        self.level = level
        self.timeout = timeout
        self.max_memory_mb = max_memory_mb
        self.work_dir = work_dir
        self.dedicated_user = None
        self.vfs = None
        self.strict = level in (SandboxLevel.HIGH, SandboxLevel.PARANOID)

        # 根据等级初始化
        if level == SandboxLevel.PARANOID:
            self.dedicated_user = self._create_dedicated_user()

        if level in (SandboxLevel.MEDIUM, SandboxLevel.HIGH, SandboxLevel.PARANOID):
            self.vfs = VirtualFileSystem()

    def _create_dedicated_user(self) -> Optional[str]:
        """创建专用系统用户（PARANOID 模式）"""
        import random
        import string
        username = f"sandbox_{''.join(random.choices(string.ascii_lowercase, k=8))}"
        try:
            subprocess.run(
                ["useradd", "-m", "-s", "/usr/sbin/nologin", username],
                check=True, capture_output=True
            )
            return username
        except (subprocess.CalledProcessError, FileNotFoundError):
            return None

    def execute(self, code: str) -> SandboxResult:
        """
        统一执行入口

        根据安全等级执行不同级别的安全检查
        """
        # ── 1. AST 预检（MEDIUM+） ──────────
        if self.level in (SandboxLevel.MEDIUM, SandboxLevel.HIGH, SandboxLevel.PARANOID):
            analyzer = ASTSecurityAnalyzer()
            safe, issues = analyzer.is_safe(code)
            if not safe:
                result = SandboxResult()
                result.stderr = "AST 安全检查未通过：\n" + "\n".join(issues)
                result.return_code = -2
                result.error = "SECURITY_BLOCKED"
                return result

        # ── 2. 构建运行代码 ────────────────
        final_code = code
        if self.vfs:
            final_code = self.vfs.create_vfs_runner(code)

        # ── 3. 执行 ────────────────────────
        sandbox = ProcessSandbox(
            work_dir=self.work_dir,
            timeout=self.timeout,
            max_memory_mb=self.max_memory_mb,
            allow_network=self.level != SandboxLevel.PARANOID
        )

        result = sandbox.run(final_code)

        # ── 4. 清理 ────────────────────────
        if self.vfs:
            self.vfs.cleanup()

        return result

    def cleanup(self):
        """清理沙箱资源"""
        if self.dedicated_user:
            try:
                subprocess.run(
                    ["userdel", "-r", self.dedicated_user],
                    capture_output=True, check=True
                )
            except subprocess.CalledProcessError:
                pass

# 工厂函数
def create_sandbox(level: str = "medium", **kwargs) -> UnifiedSandbox:
    """便捷工厂函数"""
    return UnifiedSandbox(SandboxLevel(level), **kwargs)

# ── 安全代码示例 ──
safe_code = """
import math
import json

def fibonacci(n):
    a, b = 0, 1
    for _ in range(n):
        a, b = b, a + b
    return a

result = {
    "fib_10": fibonacci(10),
    "fib_20": fibonacci(20),
    "pi": math.pi,
    "e": math.e
}
print(json.dumps(result, indent=2))
"""

# ├── 执行（MEDIUM 级别）
sandbox = create_sandbox("medium", timeout=10)
result = sandbox.execute(safe_code)
print(result.stdout)
# {
#   "fib_10": 55,
#   "fib_20": 6765,
#   "pi": 3.141592653589793,
#   "e": 2.718281828459045
# }


# ── 恶意代码示例 ──
malicious_code = """
import os
os.system("rm -rf /")
"""

result = sandbox.execute(malicious_code)
print(result.stderr)
# AST 安全检查未通过：
# 禁止导入危险模块: os


# ── 文件操作示例 ──
file_code = """
# 写操作被重定向到虚拟文件系统
with open("output.txt", "w") as f:
    f.write("Hello from sandbox!")

# 读操作也被重定向
with open("output.txt", "r") as f:
    print(f.read())
"""

result = sandbox.execute(file_code)
print(result.stdout)
# Hello from sandbox!

# sandbox_api.py
from flask import Flask, request, jsonify

app = Flask(__name__)

@app.route("/execute", methods=["POST"])
def execute_code():
    """
    执行代码 API

    请求体:
    {
        "code": "print('hello')",
        "level": "medium",
        "timeout": 30
    }
    """
    data = request.get_json(force=True)
    code = data.get("code", "")
    level = data.get("level", "medium")
    timeout = data.get("timeout", 30)

    if not code.strip():
        return jsonify({"error": "代码不能为空"}), 400

    try:
        sandbox = create_sandbox(level, timeout=timeout)
        result = sandbox.execute(code)

        return jsonify({
            "stdout": result.stdout,
            "stderr": result.stderr,
            "return_code": result.return_code,
            "execution_time": result.execution_time,
            "timed_out": result.timed_out,
            "error": result.error
        })
    except Exception as e:
        return jsonify({"error": str(e)}), 500

@app.route("/health", methods=["GET"])
def health_check():
    return jsonify({"status": "ok"})

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=5000)

# agent_integration.py
import requests
import json

class SandboxAgentTool:
    """
    AI Agent 的代码执行工具

    让 AI Agent 可以在安全的沙箱中执行代码
    """

    def __init__(self, api_url: str = "http://localhost:5000"):
        self.api_url = api_url

    def execute(self, code: str, level: str = "medium") -> dict:
        """Agent 调用此方法执行代码"""
        resp = requests.post(
            f"{self.api_url}/execute",
            json={"code": code, "level": level, "timeout": 30},
            timeout=35
        )
        return resp.json()

    def get_tool_spec(self) -> dict:
        """
        获取 Agent 工具规范（兼容 MCP/Function Calling 格式）

        返回 OpenAI Function Calling 格式的 tool spec
        """
        return {
            "type": "function",
            "function": {
                "name": "execute_python",
                "description": "在安全的沙箱环境中执行 Python 代码。适用于运行 AI 生成的代码、计算、数据分析等。代码中禁止访问文件系统（虚拟文件系统除外）、网络和系统命令。",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "code": {
                            "type": "string",
                            "description": "要执行的 Python 代码"
                        },
                        "level": {
                            "type": "string",
                            "enum": ["low", "medium", "high"],
                            "description": "安全等级：low（仅进程隔离）、medium（+AST检查+虚拟文件系统）、high（+网络封锁）"
                        }
                    },
                    "required": ["code"]
                }
            }
        }

from queue import Queue
import threading

class SandboxPool:
    """沙箱预热池"""

    def __init__(self, min_size: int = 3, max_size: int = 10, **sandbox_kwargs):
        self.min_size = min_size
        self.max_size = max_size
        self.sandbox_kwargs = sandbox_kwargs
        self._pool = Queue()
        self._lock = threading.Lock()
        self._active_count = 0

        # 预先创建沙箱
        for _ in range(min_size):
            self._pool.put(create_sandbox(**sandbox_kwargs))

    def acquire(self) -> UnifiedSandbox:
        """从池中获取沙箱"""
        try:
            return self._pool.get_nowait()
        except:
            with self._lock:
                if self._active_count < self.max_size:
                    sandbox = create_sandbox(**self.sandbox_kwargs)
                    self._active_count += 1
                    return sandbox
            return self._pool.get()  # 阻塞等待

    def release(self, sandbox: UnifiedSandbox):
        """归还沙箱到池中"""
        self._pool.put(sandbox)