spb/uqo-chat

1"""Isolated Python execution with hard limits.23Isolation layers (best available on the host):4  * macOS: `sandbox-exec` seatbelt profile — no network, writes only inside the run dir.5  * Linux (container): expected to run inside a hardened pod (NetworkPolicy deny-all,6    read-only rootfs); we still apply rlimits.7  * Always: `python -I -B`, rlimits (CPU, RAM, files, processes), wall-clock timeout,8    output truncation, auto-save of open matplotlib figures.9"""1011from __future__ import annotations1213import base6414import os15import platform16import resource17import shutil18import subprocess19import sys20import tempfile21import time22import uuid23from pathlib import Path2425MAX_OUTPUT = 50 * 102426MAX_FILES = 1027MAX_FILES_BYTES = 20 * 1024 * 102428MEM_LIMIT = 512 * 1024 * 102429PYTHON = os.environ.get("SANDBOX_PYTHON", sys.executable)3031PRELUDE = '''32import os as _os, sys as _sys33_os.chdir(_os.environ["SANDBOX_WORKDIR"])34_sys.path.insert(0, _os.getcwd())35try:36    import matplotlib as _mpl37    _mpl.use("Agg")38    import matplotlib.pyplot as _plt39    _plt.rcParams["figure.dpi"] = 13040    _plt.rcParams["axes.grid"] = True41    _plt.rcParams["grid.alpha"] = 0.342    _plt.rcParams["axes.spines.top"] = False43    _plt.rcParams["axes.spines.right"] = False44except Exception:45    _plt = None46import atexit as _atexit47def _save_figs():48    if _plt is None:49        return50    for i, num in enumerate(_plt.get_fignums(), 1):51        try:52            _plt.figure(num).savefig(f"outputs/figure_{i}.png", bbox_inches="tight")53        except Exception as e:54            print("figure save failed:", e, file=_sys.stderr)55_atexit.register(_save_figs)56'''575859def _seatbelt_profile(workdir: str) -> str:60    return f"""(version 1)61(deny default)62(allow process-exec)63(allow process-fork)64(allow sysctl-read)65(allow mach-lookup)66(allow file-read*)67(deny file-read* (subpath "/Users") (with no-log))68(allow file-read* (subpath "{workdir}"))69(allow file-read* (subpath "{os.path.dirname(os.path.realpath(PYTHON))}"))70(allow file-read* (subpath "{sys.prefix}"))71(allow file-read* (subpath "{sys.base_prefix}"))72(allow file-read* (subpath "{os.path.expanduser('~')}/.local/share/uv"))73(allow file-read* (subpath "/opt/homebrew"))74(allow file-read* (subpath "/private/tmp"))75(allow file-write* (subpath "{workdir}"))76(allow file-write* (subpath "/private/var/folders"))77(allow file-write* (subpath "/private/tmp"))78(deny network*)79"""808182def _limits() -> None:83    # Runs in the child before exec.84    try:85        resource.setrlimit(resource.RLIMIT_CPU, (25, 30))86        resource.setrlimit(resource.RLIMIT_FSIZE, (MAX_FILES_BYTES, MAX_FILES_BYTES))87        resource.setrlimit(resource.RLIMIT_NOFILE, (256, 256))88        if platform.system() != "Darwin":  # RLIMIT_AS breaks numpy on macOS89            resource.setrlimit(resource.RLIMIT_AS, (MEM_LIMIT, MEM_LIMIT))90        resource.setrlimit(resource.RLIMIT_NPROC, (32, 32))91    except (ValueError, OSError):92        pass939495def run(code: str, files_in: list[dict], timeout_s: int = 30) -> dict:96    run_id = uuid.uuid4().hex97    base = Path(tempfile.gettempdir()) / "uqo-sandbox" / run_id98    (base / "inputs").mkdir(parents=True)99    (base / "outputs").mkdir()100    for f in files_in[:MAX_FILES]:101        name = os.path.basename(str(f.get("name", "input")))102        data = base64.b64decode(f.get("content_b64", "") or "")103        (base / "inputs" / name).write_bytes(data[:MAX_FILES_BYTES])104    main = base / "main.py"105    main.write_text(PRELUDE + "\n" + code, encoding="utf-8")106107    env = {108        "PATH": "/usr/bin:/bin",109        "HOME": str(base),110        "TMPDIR": str(base),111        "SANDBOX_WORKDIR": str(base),112        "MPLCONFIGDIR": str(base / ".mpl"),113        "PYTHONHASHSEED": "0",114        "LANG": "fr_CA.UTF-8",115        "LC_ALL": "fr_CA.UTF-8",116        "OMP_NUM_THREADS": "2",117        "OPENBLAS_NUM_THREADS": "2",118    }119    cmd = [PYTHON, "-I", "-B", str(main)]120    if platform.system() == "Darwin" and shutil.which("sandbox-exec"):121        profile = base / "profile.sb"122        profile.write_text(_seatbelt_profile(str(base)))123        cmd = ["/usr/bin/sandbox-exec", "-f", str(profile)] + cmd124125    t0 = time.perf_counter()126    truncated = False127    try:128        proc = subprocess.run(129            cmd, cwd=base, env=env, capture_output=True, timeout=timeout_s,130            preexec_fn=_limits, check=False,131        )132        stdout, stderr, code_ = proc.stdout, proc.stderr, proc.returncode133    except subprocess.TimeoutExpired as exc:134        stdout = exc.stdout or b""135        stderr = (exc.stderr or b"") + f"\nTemps d'exécution dépassé ({timeout_s} s).".encode()136        code_ = 124137    duration = int((time.perf_counter() - t0) * 1000)138139    def clip(b: bytes) -> str:140        nonlocal truncated141        s = b.decode("utf-8", errors="replace")142        if len(s) > MAX_OUTPUT:143            truncated = True144            return s[:MAX_OUTPUT] + "\n… (tronqué)"145        return s146147    files_out: list[dict] = []148    total = 0149    candidates = sorted(list((base / "outputs").iterdir()) +150                        [p for p in base.iterdir() if p.is_file() and p.name not in151                         {"main.py", "profile.sb"}])152    for p in candidates:153        if not p.is_file() or p.name.startswith("."):154            continue155        size = p.stat().st_size156        if size == 0 or total + size > MAX_FILES_BYTES or len(files_out) >= MAX_FILES:157            continue158        total += size159        files_out.append({"name": p.name, "size": size,160                          "content_b64": base64.b64encode(p.read_bytes()).decode()})161    shutil.rmtree(base, ignore_errors=True)162    return {"stdout": clip(stdout), "stderr": clip(stderr), "exit_code": code_,163            "duration_ms": duration, "files_out": files_out, "truncated": truncated}164