env/server/repl.py

"""FATHOM Python REPL sandbox.

Security design (per STACK.md §7, CLAUDE.md §7, PITFALLS.md H3):
  1. RestrictedPython AST-level filter (compile_restricted) blocks attribute-access escapes
     (__class__, __subclasses__, __globals__, traceback walks).
  2. Subprocess isolation with RLIMIT_AS=512MB, RLIMIT_CPU=5s, RLIMIT_FSIZE=0 (no file writes),
     cwd=/tmp/episode-{uuid}, start_new_session=True, network denied via unshare -n or seccomp.
  3. Whitelisted builtins + whitelisted import hook.
  4. 30s wall-clock hard kill; on timeout, caller's globals_dict is NOT overwritten (D-03).

Threat model: see Plan 00-02's <threat_model> block.

Platform notes:
  - `resource` module is POSIX-only. On Windows (dev laptops only) the rlimits are a no-op;
    RestrictedPython + the import hook + the wall-clock timeout remain active. The venue
    is a Linux A100 container, so the primary hardening path runs there.
  - `os.setsid` / `start_new_session` is POSIX-only. Used when `os.name == "posix"`.
  - `unshare -n` / seccomp are Linux-only stretch goals; the blocklist already removes the
    `socket`/`urllib`/`http` surface at the import layer.
"""
from __future__ import annotations

import base64
import json
import os
import pickle
import subprocess
import sys
import tempfile
import time
import uuid
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Callable

from RestrictedPython import compile_restricted, safe_builtins as _rp_safe_builtins

# `resource` is POSIX-only — guarded import so this module loads on Windows.
try:
    import resource  # type: ignore[import-not-found]
    _HAS_RESOURCE = True
except ImportError:  # pragma: no cover - Windows path
    resource = None  # type: ignore[assignment]
    _HAS_RESOURCE = False


# ───────────────────────────────────────────────────────────────────────────
# Whitelists / blocklists (STACK.md §7 + CLAUDE.md §7 + CONTEXT.md D-12 verbatim)
# ───────────────────────────────────────────────────────────────────────────

# Allowed imports (9 modules) — STACK.md §7 + CLAUDE.md §7 + CONTEXT.md D-12
ALLOWED_IMPORTS: frozenset[str] = frozenset({
    "re", "json", "math", "statistics", "collections",
    "itertools", "string", "functools", "operator",
})

# Explicitly blocked imports (15 modules) — STACK.md §7 + CLAUDE.md §7
BLOCKED_IMPORTS: frozenset[str] = frozenset({
    "os", "sys", "subprocess", "socket", "urllib", "http", "requests",
    "pathlib", "io", "ctypes", "threading", "multiprocessing",
    "pickle", "marshal", "importlib",
})

# Builtins whitelist (CONTEXT.md D-12 + STACK.md §7)
_SAFE_BUILTIN_NAMES: tuple[str, ...] = (
    "len", "range", "enumerate", "zip", "map", "filter",
    "sorted", "min", "max", "sum", "any", "all", "print",
    "str", "int", "float", "bool",
    "list", "dict", "tuple", "set", "frozenset",
    "abs", "round", "divmod", "pow",
    "repr", "reversed", "iter", "next",
    "isinstance", "issubclass",  # required by RestrictedPython runtime
)

# Names explicitly removed even if RestrictedPython's safe_builtins includes them.
# Attribute-access primitives (getattr/setattr/delattr/hasattr) are the #1 escape
# surface — STACK.md §7 adversarial test #5 (`getattr(__builtins__, 'open')`).
_FORBIDDEN_NAMES: frozenset[str] = frozenset({
    "open", "exec", "eval", "compile", "__import__",
    "input", "breakpoint", "help", "exit", "quit",
    "getattr", "setattr", "delattr", "hasattr",
    "__builtins__", "globals", "locals", "vars", "dir",
})


def _build_safe_builtins() -> dict[str, Any]:
    """Build a minimal builtins dict.

    Starts from RestrictedPython.safe_builtins (pre-vetted) then narrows to our whitelist.
    Whitelisted names missing from RP's ``safe_builtins`` (e.g. ``print``, ``list``,
    ``dict``, ``set``, ``enumerate``, ``map``, ``filter``, ``min``, ``max``, ``sum``,
    ``any``, ``all``, ``reversed``, ``iter``, ``next``) are fetched from Python's
    real ``builtins`` module — they are safe and RP simply doesn't include them
    by default.

    Returns a plain dict (NOT the builtins module) so user code cannot import
    through it via ``__builtins__.open(...)`` — ``_FORBIDDEN_NAMES`` is the hard
    barrier against the ``getattr(__builtins__, 'open')`` escape named in
    STACK.md §7 adversarial test #5.
    """
    import builtins as _py_builtins

    b: dict[str, Any] = {}
    # Prefer RP's safe_builtins where it defines a whitelisted name (its versions may
    # be subtly hardened); fall back to real builtins for the rest.
    for name in _SAFE_BUILTIN_NAMES:
        if name in _rp_safe_builtins:
            b[name] = _rp_safe_builtins[name]
        elif hasattr(_py_builtins, name):
            b[name] = getattr(_py_builtins, name)
        # silently skip names that don't exist on either source
    # Literals RestrictedPython expects
    b["True"] = True
    b["False"] = False
    b["None"] = None
    # Explicitly remove anything forbidden even if it slipped in via safe_builtins.
    # safe_builtins notably includes 'setattr' and 'delattr' — both are in
    # _FORBIDDEN_NAMES, so they get popped here.
    for k in _FORBIDDEN_NAMES:
        b.pop(k, None)
    return b


def _safe_import(
    name: str,
    globals: dict[str, Any] | None = None,  # noqa: A002 - Python __import__ signature
    locals: dict[str, Any] | None = None,    # noqa: A002
    fromlist: tuple[str, ...] = (),
    level: int = 0,
):
    """Whitelisted ``__import__`` hook.

    Rejects anything not in :data:`ALLOWED_IMPORTS`. Relative imports (``level != 0``)
    are always rejected. Delegates to :mod:`importlib` only for whitelisted roots —
    `importlib` itself is never exposed to user code (it's imported locally here).
    """
    if level != 0:
        raise ImportError(
            f"Relative imports not allowed in FATHOM REPL sandbox (requested level={level})"
        )
    root = name.split(".")[0]
    if root in BLOCKED_IMPORTS:
        raise ImportError(
            f"Import of {name!r} is explicitly blocked in FATHOM REPL sandbox"
        )
    if root not in ALLOWED_IMPORTS:
        raise ImportError(
            f"Import of {name!r} not in whitelist {sorted(ALLOWED_IMPORTS)}"
        )
    import importlib  # local import; NOT exposed to user code
    return importlib.import_module(name)


# ───────────────────────────────────────────────────────────────────────────
# Compile wrapper
# ───────────────────────────────────────────────────────────────────────────


class ReplCompileError(Exception):
    """RestrictedPython rejected the source (AST-level policy violation)."""


def _compile_user_code(code: str) -> Any:
    """Run RestrictedPython's AST filter.

    Raises :class:`ReplCompileError` on policy violation, :class:`SyntaxError`
    on malformed source.

    STACK §7 adversarial test coverage:
      - ``().__class__.__base__.__subclasses__()`` → compile_restricted flags ``__class__`` access.
      - ``import os`` → Python compile succeeds; ``_safe_import`` rejects at exec time.
      - ``__builtins__['open']`` → compile_restricted flags subscript/attr on ``__builtins__``.
    """
    compiled = compile_restricted(code, filename="<fathom_repl>", mode="exec")
    if compiled is None:
        raise ReplCompileError(
            "RestrictedPython returned None (policy violation in submitted source)"
        )
    return compiled


# ───────────────────────────────────────────────────────────────────────────
# Result envelope
# ───────────────────────────────────────────────────────────────────────────


@dataclass
class ReplResult:
    """Return value of :func:`run_repl`.

    Attributes:
        stdout: captured stdout from the child.
        stderr: captured stderr from the child (or diagnostic from the parent).
        exception: ``"ExcType: message"`` if user code raised, else ``None``.
        globals_dict: updated globals after successful exec. **On timeout or pre-exec
            parent-side failure, this equals the original input dict unchanged**
            (D-03 — prior state preserved).
        timed_out: ``True`` iff the 30s wall-clock kill fired.
        wall_time_s: monotonic elapsed time from Popen to result assembly.
    """

    stdout: str = ""
    stderr: str = ""
    exception: str | None = None
    globals_dict: dict[str, Any] = field(default_factory=dict)
    timed_out: bool = False
    wall_time_s: float = 0.0


# ───────────────────────────────────────────────────────────────────────────
# Subprocess-isolated executor
# ───────────────────────────────────────────────────────────────────────────
# The child subprocess receives (pickled globals, code string) as a JSON envelope
# on stdin, runs compile_restricted + exec in a locked-down namespace, and writes
# a JSON envelope {stdout, stderr, exception, globals_b64} to stdout. Parent
# enforces the 30s wall-clock via subprocess.communicate(timeout=...) and kills
# on expiry (D-03 — globals NOT overwritten on timeout).
#
# Defense-in-depth: RestrictedPython compile_restricted is ALSO re-run inside the
# child, so the parent's Task-1 pre-flight is advisory — the child is the true
# security boundary (STACK.md §7 "AST-level + OS-level").

# Placeholders (NAME_WRAPPED_IN_TRIPLE_UNDERSCORES) are replaced at runtime by
# _render_child_runner. Using unique sentinel tokens avoids clashes with literal
# curly braces that appear inside f-strings in the child source.
_CHILD_RUNNER_TEMPLATE = r'''
import base64, json, pickle, sys, traceback, io, contextlib

# Re-derive the sandbox inside the child. Importing anything beyond
# RestrictedPython + stdlib essentials would widen the attack surface.
from RestrictedPython import compile_restricted, safe_builtins as _rp_safe_builtins

ALLOWED_IMPORTS = set(___ALLOWED___)
BLOCKED_IMPORTS = set(___BLOCKED___)
_SAFE_NAMES     = tuple(___SAFE_NAMES___)
_FORBIDDEN      = set(___FORBIDDEN___)


def _safe_import(name, globals=None, locals=None, fromlist=(), level=0):
    if level != 0:
        raise ImportError("Relative imports not allowed")
    root = name.split(".")[0]
    if root in BLOCKED_IMPORTS:
        raise ImportError("Import of " + repr(name) + " is blocked in FATHOM REPL sandbox")
    if root not in ALLOWED_IMPORTS:
        raise ImportError("Import of " + repr(name) + " not in whitelist")
    import importlib
    return importlib.import_module(name)


def _stub_llm(prompt, chunk, depth=1):
    """Deterministic regex-echo stub (Phase 0 only; Plan 04 replaces via RPC).

    Extracts the first keyword (>=3 alphanumeric chars) from `prompt`, returns
    an 80-char window of `chunk` containing it, falling back to `chunk[:200]`.
    """
    import re as _re
    m = _re.search(r"\b\w{3,}\b", prompt)
    kw = m.group(0) if m else ""
    if kw and kw in chunk:
        idx = chunk.index(kw)
        return chunk[idx:idx + 80]
    return chunk[:200]


def _build_builtins():
    import builtins as _py_builtins
    b = {}
    for name in _SAFE_NAMES:
        if name in _rp_safe_builtins:
            b[name] = _rp_safe_builtins[name]
        elif hasattr(_py_builtins, name):
            b[name] = getattr(_py_builtins, name)
    b["True"] = True
    b["False"] = False
    b["None"] = None
    for k in _FORBIDDEN:
        b.pop(k, None)
    # Install the whitelisted import hook so `import re` works but `import os` fails.
    b["__import__"] = _safe_import
    return b


def main():
    payload = json.loads(sys.stdin.read())
    code = payload["code"]
    try:
        user_globals = pickle.loads(base64.b64decode(payload["globals_b64"]))
        if not isinstance(user_globals, dict):
            user_globals = {}
    except Exception as e:
        sys.stdout.write(json.dumps({
            "stdout": "",
            "stderr": "",
            "exception": "globals unpickle failed: " + repr(e),
            "globals_b64": payload["globals_b64"],
        }))
        return

    # Install safe builtins + llm stub into the user's namespace.
    user_globals["__builtins__"] = _build_builtins()
    user_globals.setdefault("llm", _stub_llm)
    # RestrictedPython expects _getattr_ / _getitem_ helpers if attribute access
    # is ever emitted by legitimate user code. Provide minimal safe versions.
    user_globals.setdefault("_getattr_", getattr)
    user_globals.setdefault("_getitem_", lambda o, k: o[k])
    user_globals.setdefault("_getiter_", iter)
    user_globals.setdefault("_iter_unpack_sequence_", lambda it, count, _obj: tuple(it))
    # RestrictedPython rewrites `print(...)` to `_print = _print_(); _print._call_print(...)`.
    # Our collector forwards to the real builtins.print so the surrounding
    # redirect_stdout captures output into stdout_buf.
    class _StdoutPrintCollector:
        def __init__(self, _getattr_=None):
            self._getattr_ = _getattr_
        def _call_print(self, *objects, **kwargs):
            import builtins as _b
            _b.print(*objects, **kwargs)
        def __call__(self):
            return ""
        def write(self, text):
            sys.stdout.write(text)
    user_globals.setdefault("_print_", _StdoutPrintCollector)

    stdout_buf = io.StringIO()
    stderr_buf = io.StringIO()
    exception_str = None
    try:
        compiled = compile_restricted(code, filename="<fathom_repl_child>", mode="exec")
        if compiled is None:
            raise SyntaxError("RestrictedPython rejected source (compile returned None)")
        with contextlib.redirect_stdout(stdout_buf), contextlib.redirect_stderr(stderr_buf):
            exec(compiled, user_globals, user_globals)
    except BaseException as e:
        exception_str = type(e).__name__ + ": " + str(e)

    # Strip non-picklable injected values before returning globals to parent.
    for k in ("__builtins__", "llm", "_getattr_", "_getitem_", "_getiter_",
              "_iter_unpack_sequence_", "_print_", "_print"):
        user_globals.pop(k, None)

    try:
        updated_b64 = base64.b64encode(pickle.dumps(user_globals)).decode("ascii")
    except Exception as e:
        updated_b64 = payload["globals_b64"]  # fall back to original on pickle failure
        exception_str = (exception_str or "") + " | globals pickle failed: " + repr(e)

    sys.stdout.write(json.dumps({
        "stdout": stdout_buf.getvalue(),
        "stderr": stderr_buf.getvalue(),
        "exception": exception_str,
        "globals_b64": updated_b64,
    }))


if __name__ == "__main__":
    main()
'''


def _render_child_runner() -> str:
    """Fill the child-runner template with our verbatim whitelist/blocklist data."""
    return (_CHILD_RUNNER_TEMPLATE
            .replace("___ALLOWED___", json.dumps(sorted(ALLOWED_IMPORTS)))
            .replace("___BLOCKED___", json.dumps(sorted(BLOCKED_IMPORTS)))
            .replace("___SAFE_NAMES___", json.dumps(list(_SAFE_BUILTIN_NAMES)))
            .replace("___FORBIDDEN___", json.dumps(sorted(_FORBIDDEN_NAMES))))


def _set_rlimits_posix() -> None:  # pragma: no cover - POSIX path, exercised at venue
    """``preexec_fn`` for POSIX subprocesses — enforces the STACK §7 rlimits.

    Limits applied (STACK.md §7 subprocess-level defenses):
      * ``RLIMIT_AS``    = 512 MiB (virtual address space hard cap)
      * ``RLIMIT_CPU``   = 5 s     (CPU time; 30 s wall-clock is the outer bound)
      * ``RLIMIT_FSIZE`` = 0       (no file writes even if ``open`` leaks)
      * ``RLIMIT_NOFILE``= 64      (fd cap)

    Also detaches from the parent process group via :func:`os.setsid` so that
    :func:`os.killpg` on timeout reaches the whole tree.
    """
    if not _HAS_RESOURCE or resource is None:
        return
    _BYTES_512_MIB = 512 * 1024 * 1024
    resource.setrlimit(resource.RLIMIT_AS, (_BYTES_512_MIB, _BYTES_512_MIB))
    resource.setrlimit(resource.RLIMIT_CPU, (5, 5))
    resource.setrlimit(resource.RLIMIT_FSIZE, (0, 0))
    resource.setrlimit(resource.RLIMIT_NOFILE, (64, 64))
    # Detach from parent process group so SIGKILL on timeout kills the whole tree.
    os.setsid()


def run_repl(
    code: str,
    globals_dict: dict[str, Any],
    llm_callable: Callable[[str, str, int], str] | None = None,
    timeout_s: float = 30.0,
    episode_id: str | None = None,
) -> ReplResult:
    """Execute ``code`` in an isolated subprocess with stateful globals.

    Args:
        code: Python source submitted by the model.
        globals_dict: persistent globals from prior steps. May contain ``ctx`` and
            any variables the model defined earlier. Must be JSON/pickle-safe —
            only JSON-serialisable values survive the round-trip; functions and
            open file handles do not.
        llm_callable: accepted for forward-compat with Plan 04; ignored in Phase 0
            because the child uses an embedded deterministic regex-echo stub
            (D-05). Plan 04 will inject the real Qwen-backed callable via an
            RPC side channel.
        timeout_s: wall-clock limit (default 30 s per STACK §7).
        episode_id: used to build the ephemeral cwd ``<tmp>/episode-<id>``;
            auto-generated if ``None``.

    Returns:
        :class:`ReplResult`. On timeout, pre-exec pickle failure, or child
        protocol error, ``globals_dict`` in the result equals the unchanged
        input dict (D-03 — prior state preserved).
    """
    eid = episode_id or uuid.uuid4().hex
    # Ephemeral cwd — on POSIX this is /tmp/episode-<id>; tempfile.gettempdir()
    # is cross-platform so the Windows dev laptop still gets an isolated cwd.
    ep_cwd = Path(tempfile.gettempdir()) / f"episode-{eid}"
    ep_cwd.mkdir(parents=True, exist_ok=True)

    started = time.monotonic()

    # Serialise globals for the child. If the caller has stuffed something
    # un-picklable into globals_dict (e.g., a live socket), bail fast —
    # DO NOT overwrite the caller's state.
    try:
        globals_b64 = base64.b64encode(pickle.dumps(globals_dict)).decode("ascii")
    except Exception as e:
        return ReplResult(
            stderr=f"Parent failed to pickle globals_dict: {e!r}",
            exception=f"PickleError: {e!r}",
            globals_dict=dict(globals_dict),
            wall_time_s=time.monotonic() - started,
        )

    payload = json.dumps({"code": code, "globals_b64": globals_b64})
    runner_src = _render_child_runner()

    # POSIX-only hardening: preexec_fn sets rlimits + new session. On Windows,
    # preexec_fn is unsupported (Popen would raise); RestrictedPython + the
    # wall-clock timeout + the restricted builtins remain the primary defense.
    is_posix = os.name == "posix"
    preexec = _set_rlimits_posix if is_posix else None

    # Minimal env — strip secrets (HF_TOKEN, WANDB_API_KEY) from the child.
    # PYTHONPATH: the child uses -E (ignores PYTHON* env vars at startup) but
    # its bootstrap code imports RestrictedPython before any user code runs,
    # so the package must be locatable. We pass the parent's site-packages
    # explicitly. The child's _safe_import hook still restricts user-code imports.
    # On Windows: SystemRoot is required for BCryptGenRandom (Python hash-seed init).
    import site as _site
    _site_dirs = _site.getsitepackages() if hasattr(_site, "getsitepackages") else []
    _user_site = _site.getusersitepackages() if hasattr(_site, "getusersitepackages") else ""
    _pythonpath_parts = [p for p in _site_dirs + ([_user_site] if _user_site else []) if p]
    child_env = {
        "PATH": os.environ.get("PATH", ""),
        "PYTHONDONTWRITEBYTECODE": "1",
        "PYTHONPATH": os.pathsep.join(_pythonpath_parts),
    }
    if os.name == "nt":
        for _k in ("SystemRoot", "SystemDrive", "TEMP", "TMP", "WINDIR"):
            if _k in os.environ:
                child_env[_k] = os.environ[_k]

    # -s: skip user site-packages (Python startup won't auto-add them).
    # NOT -I/-E: we pass PYTHONPATH explicitly in child_env so the child can find
    # RestrictedPython; secrets are already stripped from child_env above.
    args = [sys.executable, "-s", "-c", runner_src]

    try:
        proc = subprocess.Popen(
            args,
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            cwd=str(ep_cwd),
            preexec_fn=preexec,
            start_new_session=is_posix,
            env=child_env,
        )
    except Exception as e:
        return ReplResult(
            stderr=f"Popen failed: {e!r}",
            exception=f"SpawnError: {e!r}",
            globals_dict=dict(globals_dict),
            wall_time_s=time.monotonic() - started,
        )

    try:
        stdout_b, stderr_b = proc.communicate(
            input=payload.encode("utf-8"),
            timeout=timeout_s,
        )
    except subprocess.TimeoutExpired:
        # D-03: kill the child AND its descendants; DO NOT overwrite globals_dict.
        try:
            if is_posix:
                import signal
                os.killpg(os.getpgid(proc.pid), signal.SIGKILL)
            else:
                proc.kill()
        except Exception:
            try:
                proc.kill()
            except Exception:
                pass
        try:
            proc.wait(timeout=2)
        except Exception:
            pass
        _cleanup_ep_cwd(ep_cwd)
        return ReplResult(
            stderr="TimeoutError",
            exception=f"TimeoutError: wall-clock exceeded {timeout_s}s",
            globals_dict=dict(globals_dict),  # UNCHANGED — D-03
            timed_out=True,
            wall_time_s=time.monotonic() - started,
        )

    wall = time.monotonic() - started
    stdout_raw = stdout_b.decode("utf-8", errors="replace")
    stderr_raw = stderr_b.decode("utf-8", errors="replace")
    _cleanup_ep_cwd(ep_cwd)

    # Parse the JSON envelope from child stdout.
    try:
        env = json.loads(stdout_raw)
    except json.JSONDecodeError as e:
        return ReplResult(
            stdout="",
            stderr=stderr_raw or stdout_raw,
            exception=f"ChildProtocolError: {e!r}",
            globals_dict=dict(globals_dict),
            wall_time_s=wall,
        )

    # Decode updated globals; fall back to input on any deserialisation error.
    try:
        new_globals = pickle.loads(base64.b64decode(env["globals_b64"]))
        if not isinstance(new_globals, dict):
            new_globals = dict(globals_dict)
    except Exception as e:
        new_globals = dict(globals_dict)
        env["exception"] = (env.get("exception") or "") + f" | globals unpickle failed: {e!r}"

    return ReplResult(
        stdout=env.get("stdout", ""),
        stderr=env.get("stderr", "") or stderr_raw,
        exception=env.get("exception"),
        globals_dict=new_globals,
        timed_out=False,
        wall_time_s=wall,
    )


def _cleanup_ep_cwd(ep_cwd: Path) -> None:
    """Best-effort ephemeral-cwd wipe. Never raises."""
    try:
        for p in ep_cwd.iterdir():
            try:
                if p.is_dir():
                    import shutil
                    shutil.rmtree(p, ignore_errors=True)
                else:
                    p.unlink()
            except Exception:
                pass
        ep_cwd.rmdir()
    except Exception:
        pass