Create app.py

app.py

@@ -0,0 +1,136 @@
+# app.py
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel, Field
+from typing import Optional
+import subprocess
+import textwrap
+import ast
+import resource
+import os
+
+app = FastAPI(title="Python Exec Service (PES) on HuggingFace")
+
+# ----------------- 配置 -----------------
+MAX_CODE_LENGTH = 20_000           # 代码最大长度
+MAX_TIMEOUT = 20                   # 固定最长 20 秒（由 PES 控制）
+MAX_MEMORY_MB = 256                # 子进程最大内存
+MAX_OUTPUT_SIZE = 4 * 1024 * 1024  # 4MB 输出限制
+
+# ----------------- 请求 / 响应模型 -----------------
+class ExecRequest(BaseModel):
+    code: str = Field(..., description="要执行的 Python 代码")
+    stdin: Optional[str] = Field("", description="传给程序的标准输入")
+
+class ExecResponse(BaseModel):
+    stdout: str
+    stderr: str
+    return_code: Optional[int]
+    timeout: bool
+
+# ----------------- 安全检查：AST 审计 -----------------
+DANGEROUS_NAMES = {
+    "__import__", "eval", "exec", "open", "compile",
+    "input", "globals", "locals", "vars",
+    "os", "sys", "subprocess", "socket",
+    "shutil", "pathlib"
+}
+
+DANGEROUS_MODULES = {
+    "os", "sys", "subprocess", "socket", "shutil",
+    "pathlib", "threading", "multiprocessing"
+}
+
+class SafeVisitor(ast.NodeVisitor):
+    def visit_Import(self, node):
+        for alias in node.names:
+            if alias.name.split('.')[0] in DANGEROUS_MODULES:
+                raise ValueError(f"禁止导入模块: {alias.name}")
+        self.generic_visit(node)
+
+    def visit_ImportFrom(self, node):
+        if node.module and node.module.split('.')[0] in DANGEROUS_MODULES:
+            raise ValueError(f"禁止 from 导入模块: {node.module}")
+        self.generic_visit(node)
+
+    def visit_Attribute(self, node):
+        # 禁止访问 __dict__ / __class__ / __globals__ 等双下划线属性
+        if isinstance(node.attr, str) and node.attr.startswith("__"):
+            raise ValueError(f"禁止访问特殊属性: {node.attr}")
+        self.generic_visit(node)
+
+    def visit_Name(self, node):
+        if node.id in DANGEROUS_NAMES:
+            raise ValueError(f"禁止使用名称: {node.id}")
+        self.generic_visit(node)
+
+def static_security_check(code: str):
+    try:
+        tree = ast.parse(code, mode="exec")
+    except SyntaxError as e:
+        raise ValueError(f"语法错误: {e}")
+    SafeVisitor().visit(tree)
+
+# ----------------- 资源限制（子进程 preexec_fn） -----------------
+def set_limits():
+    # CPU 时间限制（秒）—— hard limit MAX_TIMEOUT
+    resource.setrlimit(resource.RLIMIT_CPU, (MAX_TIMEOUT, MAX_TIMEOUT))
+    # 地址空间限制（内存）
+    max_bytes = MAX_MEMORY_MB * 1024 * 1024
+    resource.setrlimit(resource.RLIMIT_AS, (max_bytes, max_bytes))
+    # 文件大小限制（防止写超大文件）
+    resource.setrlimit(resource.RLIMIT_FSIZE, (MAX_OUTPUT_SIZE, MAX_OUTPUT_SIZE))
+
+# ----------------- 核心执行 API -----------------
+@app.post("/execute", response_model=ExecResponse)
+def execute(req: ExecRequest):
+    # 1. 限制代码长度
+    if len(req.code) > MAX_CODE_LENGTH:
+        raise HTTPException(status_code=400, detail="代码过长，超过限制")
+
+    # 2. 去掉多余缩进
+    code = textwrap.dedent(req.code)
+
+    # 3. 静态安全检查（AST）
+    try:
+        static_security_check(code)
+    except ValueError as e:
+        raise HTTPException(status_code=400, detail=f"安全检查不通过: {e}")
+
+    # 4. 执行：PES 仅提供命令参数，不允许外部控制 timeout
+    try:
+        proc = subprocess.run(
+            ["python", "-c", code],
+            input=req.stdin.encode("utf-8") if req.stdin else None,
+            capture_output=True,
+            timeout=MAX_TIMEOUT,   # 固定 20 秒
+            preexec_fn=set_limits  # 资源限制
+        )
+        stdout = proc.stdout[:MAX_OUTPUT_SIZE].decode("utf-8", errors="replace")
+        stderr = proc.stderr[:MAX_OUTPUT_SIZE].decode("utf-8", errors="replace")
+
+        return ExecResponse(
+            stdout=stdout,
+            stderr=stderr,
+            return_code=proc.returncode,
+            timeout=False
+        )
+    except subprocess.TimeoutExpired as e:
+        # 子进程超时
+        stdout = (e.stdout or b"")[:MAX_OUTPUT_SIZE].decode("utf-8", errors="replace")
+        stderr = (e.stderr or b"")[:MAX_OUTPUT_SIZE].decode("utf-8", errors="replace")
+        if not stderr:
+            stderr = "Process timed out"
+        return ExecResponse(
+            stdout=stdout,
+            stderr=stderr,
+            return_code=None,
+            timeout=True
+        )
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"执行失败: {e}")
+
+# 本地调试用，HF Space 也可以用这个启动命令
+if __name__ == "__main__":
+    import uvicorn
+    port = int(os.environ.get("PORT", 7860))
+    uvicorn.run("app:app", host="0.0.0.0", port=port)