main.py

import os
import subprocess
import shutil
from fastapi import FastAPI, Header, HTTPException
from pydantic import BaseModel
from typing import Optional

app = FastAPI()

# 从环境变量获取 Token
SERVICE_TOKEN = os.getenv("HF_SERVICE_TOKEN")

class CommandRequest(BaseModel):
    command: str
    session_id: str

@app.get("/")
def health_check():
    return {"status": "Sandbox is online"}

@app.post("/execute")
async def execute(req: CommandRequest, authorization: Optional[str] = Header(None)):
    if not SERVICE_TOKEN or authorization != f"Bearer {SERVICE_TOKEN}":
        raise HTTPException(status_code=401, detail="Unauthorized")

    # 为每个用户创建隔离的工作文件夹
    user_dir = os.path.join("/tmp/user_sandboxes", req.session_id)
    os.makedirs(user_dir, exist_ok=True)

    # 构造执行环境，重定向 HOME 目录
    custom_env = os.environ.copy()
    custom_env["HOME"] = user_dir
    custom_env["TMPDIR"] = user_dir

    try:
        # 在指定的用户目录下运行命令
        process = subprocess.run(
            req.command,
            shell=True,
            cwd=user_dir,
            capture_output=True,
            text=True,
            timeout=30,
            env=custom_env
        )
        
        return {
            "stdout": process.stdout,
            "stderr": process.stderr,
            "returncode": process.returncode
        }
    except subprocess.TimeoutExpired:
        return {"stdout": "", "stderr": "Error: Timeout (30s)", "returncode": 124}
    except Exception as e:
        return {"stdout": "", "stderr": f"System Error: {str(e)}", "returncode": 1}