ForcePilot/backend/package/yuxi/channel/plugins/sandbox.py

from __future__ import annotations

import ast
import builtins
import importlib.util
import logging
import types
from pathlib import Path
from typing import Any

logger = logging.getLogger(__name__)


class SandboxViolationError(Exception):
    """插件代码违反沙箱安全策略时抛出。"""


class PluginSandbox:
    """插件代码沙箱 — 限制动态加载的插件代码只能访问安全的内置函数和模块。

    通过 AST 静态分析和运行时全局变量限制，防止插件执行危险操作：
    - 禁止文件系统访问（open, file 等）
    - 禁止系统命令执行（os.system, subprocess 等）
    - 禁止网络访问（socket, urllib 等）
    - 禁止导入未授权模块
    - 禁止访问私有属性（__import__ 等）
    """

    # 允许的内置函数白名单
    ALLOWED_BUILTINS: set[str] = {
        # 基础类型
        "bool",
        "bytearray",
        "bytes",
        "complex",
        "dict",
        "enumerate",
        "filter",
        "float",
        "frozenset",
        "int",
        "list",
        "map",
        "memoryview",
        "object",
        "range",
        "set",
        "slice",
        "str",
        "tuple",
        "zip",
        # 常用函数
        "abs",
        "all",
        "any",
        "ascii",
        "bin",
        "callable",
        "chr",
        "classmethod",
        "compile",
        "delattr",
        "dir",
        "divmod",
        "format",
        "getattr",
        "hasattr",
        "hash",
        "hex",
        "id",
        "isinstance",
        "issubclass",
        "iter",
        "len",
        "max",
        "min",
        "next",
        "oct",
        "ord",
        "pow",
        "property",
        "repr",
        "reversed",
        "round",
        "sorted",
        "staticmethod",
        "sum",
        "vars",
        # 异常处理
        "BaseException",
        "Exception",
        "TypeError",
        "ValueError",
        "AttributeError",
        "KeyError",
        "IndexError",
        "RuntimeError",
        "NotImplementedError",
        "StopIteration",
        "GeneratorExit",
        "SystemExit",
        "ArithmeticError",
        "LookupError",
        "AssertionError",
        "BufferError",
        "EOFError",
        "ImportError",
        "ModuleNotFoundError",
        "NameError",
        "OSError",
        "OverflowError",
        "RecursionError",
        "ReferenceError",
        "SyntaxError",
        "TimeoutError",
        "ZeroDivisionError",
        # 其他安全函数
        "help",
        "print",
        "input",
        "open",  # 允许但会在模块级别重写为安全版本
        "__build_class__",
        "__name__",
        "__doc__",
        "__package__",
        "__spec__",
        "__annotations__",
        "__dict__",
        "__class__",
        "__bases__",
        "__mro__",
        "__subclasses__",
        "__init__",
        "__new__",
        "__setattr__",
        "__getattribute__",
        "__delattr__",
        "__str__",
        "__repr__",
        "__eq__",
        "__ne__",
        "__lt__",
        "__le__",
        "__gt__",
        "__ge__",
        "__hash__",
        "__bool__",
        "__len__",
        "__iter__",
        "__next__",
        "__contains__",
        "__getitem__",
        "__setitem__",
        "__delitem__",
        "__enter__",
        "__exit__",
        "__await__",
        "__aiter__",
        "__anext__",
        "__call__",
    }

    # 禁止导入的模块黑名单
    FORBIDDEN_MODULES: set[str] = {
        "os",
        "sys",
        "subprocess",
        "socket",
        "urllib",
        "urllib.request",
        "urllib.parse",
        "http",
        "http.client",
        "ftplib",
        "smtplib",
        "poplib",
        "imaplib",
        "telnetlib",
        "ctypes",
        "mmap",
        "resource",
        "signal",
        "pty",
        "pickle",
        "cPickle",
        "marshal",
        "imp",
        "builtins",
        "importlib",
        "importlib.util",
        "importlib.machinery",
        "types",
        "inspect",
        "code",
        "codeop",
        "compileall",
        "py_compile",
        "shutil",
        "pathlib",
        "tempfile",
        "glob",
        "fnmatch",
        "linecache",
        "traceback",
        "warnings",
        "contextlib",
        "atexit",
        "gc",
        "sysconfig",
        "platform",
        "pwd",
        "grp",
        "spwd",
        "crypt",
        "termios",
        "tty",
        "fcntl",
        "msvcrt",
        "winreg",
        "winsound",
        "posix",
        "nt",
        "java",
        "org",
        "sun",
        "com",
        "net",
    }

    # 允许导入的模块白名单（如果为空则使用黑名单模式）
    ALLOWED_MODULES: set[str] = {
        "abc",
        "array",
        "base64",
        "binascii",
        "bisect",
        "calendar",
        "collections",
        "collections.abc",
        "copy",
        "csv",
        "dataclasses",
        "datetime",
        "decimal",
        "enum",
        "fractions",
        "functools",
        "heapq",
        "html",
        "html.entities",
        "html.parser",
        "io",
        "itertools",
        "json",
        "keyword",
        "math",
        "numbers",
        "operator",
        "random",
        "re",
        "string",
        "struct",
        "textwrap",
        "time",
        "typing",
        "unicodedata",
        "uuid",
        "weakref",
        "xml.etree.ElementTree",
        "xml.dom.minidom",
        "xml.sax",
        "xml.parsers.expat",
        "zoneinfo",
        "hashlib",
        "hmac",
        "secrets",
        "statistics",
        "stringprep",
        "tokenize",
        "trace",
        "types",
        "warnings",
        "contextlib",
        "copyreg",
        "dis",
        "opcode",
        "pickletools",
        "quopri",
        "reprlib",
        "shelve",
        "symtable",
        "tabnanny",
        "timeit",
        "token",
        "turtle",
        "webbrowser",
        "xxlimited",
        "xxlimited_35",
        "_thread",
        "_warnings",
        "_weakref",
        "_weakrefset",
        "_io",
        "_collections",
        "_collections_abc",
        "_functools",
        "_operator",
        "_stat",
        "_string",
        "_symtable",
        "_typing",
    }

    def __init__(self) -> None:
        self._violation_log: list[str] = []

    @property
    def violations(self) -> list[str]:
        return list(self._violation_log)

    def _safe_import(self, name: str, globals_: dict | None = None, locals_: dict | None = None, fromlist: tuple = (), level: int = 0) -> Any:
        """安全的 __import__ 替代函数，限制模块导入。"""
        # 解析完整模块名
        if level > 0 and globals_ is not None:
            package = globals_.get("__package__", "")
            if package:
                if name:
                    full_name = f"{package}.{name}"
                else:
                    full_name = package
            else:
                full_name = name
        else:
            full_name = name

        # 检查是否命中黑名单
        for forbidden in self.FORBIDDEN_MODULES:
            if full_name == forbidden or full_name.startswith(forbidden + "."):
                msg = f"Import of forbidden module '{full_name}' is not allowed in sandbox"
                self._violation_log.append(msg)
                raise SandboxViolationError(msg)

        # 检查是否在白名单中（如果白名单非空）
        if self.ALLOWED_MODULES:
            allowed = any(
                full_name == mod or full_name.startswith(mod + ".")
                for mod in self.ALLOWED_MODULES
            )
            if not allowed:
                msg = f"Import of module '{full_name}' is not allowed in sandbox"
                self._violation_log.append(msg)
                raise SandboxViolationError(msg)

        # 使用原始 import 但限制在允许范围内
        return __import__(name, globals_, locals_, fromlist, level)

    def _safe_open(self, *args: Any, **kwargs: Any) -> Any:
        """安全的 open 替代函数，禁止文件系统访问。"""
        msg = "File operations (open) are not allowed in sandbox"
        self._violation_log.append(msg)
        raise SandboxViolationError(msg)

    def _safe_eval(self, *args: Any, **kwargs: Any) -> Any:
        """禁止 eval。"""
        msg = "eval() is not allowed in sandbox"
        self._violation_log.append(msg)
        raise SandboxViolationError(msg)

    def _safe_exec(self, *args: Any, **kwargs: Any) -> Any:
        """禁止 exec。"""
        msg = "exec() is not allowed in sandbox"
        self._violation_log.append(msg)
        raise SandboxViolationError(msg)

    def _safe_compile(self, *args: Any, **kwargs: Any) -> Any:
        """禁止 compile。"""
        msg = "compile() is not allowed in sandbox"
        self._violation_log.append(msg)
        raise SandboxViolationError(msg)

    @staticmethod
    def _is_private_attr_name(name: str) -> bool:
        """检查属性名是否为私有属性（双下划线前缀或 Python name mangling 模式）。"""
        if name.startswith("__") and not name.endswith("__"):
            return True
        if "__" in name and not name.startswith("__") and not name.endswith("__"):
            return True
        return False

    def _safe_getattr(self, obj: Any, name: str, default: Any = None) -> Any:
        """安全的 getattr，禁止访问双下划线私有属性及 name mangling 绕过。"""
        if self._is_private_attr_name(name):
            msg = f"Access to private attribute '{name}' is not allowed in sandbox"
            self._violation_log.append(msg)
            raise SandboxViolationError(msg)
        return getattr(obj, name, default)

    def _safe_setattr(self, obj: Any, name: str, value: Any) -> None:
        """安全的 setattr，禁止修改双下划线私有属性及 name mangling 绕过。"""
        if self._is_private_attr_name(name):
            msg = f"Modification of private attribute '{name}' is not allowed in sandbox"
            self._violation_log.append(msg)
            raise SandboxViolationError(msg)
        setattr(obj, name, value)

    def _safe_delattr(self, obj: Any, name: str) -> None:
        """安全的 delattr，禁止删除双下划线私有属性及 name mangling 绕过。"""
        if self._is_private_attr_name(name):
            msg = f"Deletion of private attribute '{name}' is not allowed in sandbox"
            self._violation_log.append(msg)
            raise SandboxViolationError(msg)
        delattr(obj, name)

    def _create_restricted_globals(self) -> dict[str, Any]:
        """创建受限的全局命名空间。"""
        restricted_builtins = {}
        for name in self.ALLOWED_BUILTINS:
            if hasattr(builtins, name):
                restricted_builtins[name] = getattr(builtins, name)

        # 替换危险函数
        restricted_builtins["__import__"] = self._safe_import
        restricted_builtins["open"] = self._safe_open
        restricted_builtins["eval"] = self._safe_eval
        restricted_builtins["exec"] = self._safe_exec
        restricted_builtins["compile"] = self._safe_compile
        restricted_builtins["getattr"] = self._safe_getattr
        restricted_builtins["setattr"] = self._safe_setattr
        restricted_builtins["delattr"] = self._safe_delattr

        return {"__builtins__": restricted_builtins}

    @staticmethod
    def validate_source(source: str, filename: str = "<sandbox>") -> ast.AST:
        """静态分析源代码 AST，检测危险语法结构（不依赖沙箱实例）。

        可用于在任何动态代码执行前进行安全预检，无需创建沙箱实例。
        仅执行 AST 级别的静态检查，不涉及运行时全局变量限制。

        Args:
            source: 要验证的 Python 源代码
            filename: 用于错误报告的源文件名

        Returns:
            解析后的 AST 树

        Raises:
            SandboxViolationError: 代码包含危险语法结构时
        """
        try:
            tree = ast.parse(source, filename=filename)
        except SyntaxError as e:
            raise SandboxViolationError(f"Syntax error in source code: {e}") from e

        for node in ast.walk(tree):
            if isinstance(node, ast.Call):
                if isinstance(node.func, ast.Name) and node.func.id == "__import__":
                    raise SandboxViolationError(
                        "Direct __import__() calls are not allowed in sandbox"
                    )

            if isinstance(node, ast.ImportFrom):
                if node.names[0].name == "*":
                    raise SandboxViolationError(
                        "'from ... import *' is not allowed in sandbox"
                    )

        return tree

    def _validate_ast(self, source: str, filename: str = "<sandbox>") -> ast.AST:
        """静态分析 AST，检测危险语法结构（实例方法，会记录违规日志）。"""
        try:
            tree = self.validate_source(source, filename)
        except SandboxViolationError as e:
            self._violation_log.append(str(e))
            raise
        return tree

    def execute(self, source: str, filename: str = "<sandbox>", extra_globals: dict[str, Any] | None = None) -> dict[str, Any]:
        """在沙箱中执行 Python 源代码，返回执行后的全局命名空间。

        Args:
            source: 要执行的 Python 源代码
            filename: 用于错误报告的源文件名
            extra_globals: 额外注入到全局命名空间的安全对象

        Returns:
            执行后的全局命名空间字典

        Raises:
            SandboxViolationError: 代码违反安全策略时
        """
        # AST 静态分析
        tree = self._validate_ast(source, filename)

        # 编译为代码对象
        code = compile(tree, filename, "exec")

        # 创建受限全局环境
        restricted_globals = self._create_restricted_globals()
        if extra_globals:
            # 只允许注入安全的额外对象
            for key, value in extra_globals.items():
                if not key.startswith("__") or key in ("__name__", "__file__", "__doc__", "__package__", "__spec__"):
                    restricted_globals[key] = value

        # 执行代码
        exec(code, restricted_globals)  # noqa: S102

        return restricted_globals

    def load_module_from_source(self, source: str, module_name: str, filename: str = "<sandbox>") -> types.ModuleType:
        """在沙箱中从源代码创建模块。

        Args:
            source: 模块源代码
            module_name: 模块名称
            filename: 用于错误报告的源文件名

        Returns:
            创建的模块对象
        """
        tree = self._validate_ast(source, filename)
        code = compile(tree, filename, "exec")

        module = types.ModuleType(module_name)
        module.__dict__.update(self._create_restricted_globals())
        module.__name__ = module_name
        module.__file__ = filename

        exec(code, module.__dict__)  # noqa: S102

        return module

    def load_module_from_file(self, file_path: str | Path, module_name: str | None = None) -> types.ModuleType:
        """在沙箱中从文件加载模块。

        Args:
            file_path: 文件路径
            module_name: 模块名称（默认为文件名）

        Returns:
            创建的模块对象
        """
        file_path = Path(file_path)
        if not file_path.exists():
            raise FileNotFoundError(f"Plugin file not found: {file_path}")

        source = file_path.read_text(encoding="utf-8")
        name = module_name or file_path.stem

        return self.load_module_from_source(source, name, str(file_path))