Agentswarm/orchestrator/sandbox.py

"""Secure-within-pod code sandbox for testing specialist agents' generated code.

SECURITY MODEL — read before changing anything here:
  The OS-level isolation boundary is the **Kubernetes pod / container** this process runs in:
  non-root user, read-only root filesystem, NetworkPolicy egress-deny, CPU/memory/pids limits,
  and seccomp — all enforced at deploy time (see docs/integration/security-boundary.md and the
  Manager/release deployment manifests; NOT editable from this repo). Per Owner ruling, running
  generated test code inside that pod is acceptable.

  This module adds **in-pod defense in depth** on top of the pod boundary:
    - ephemeral per-run workdir (tempfile), always removed in `finally`
    - wall-clock timeout with process-group kill
    - POSIX resource limits (CPU time, address space, file size, subprocess count) where available
    - environment scrub: no API keys / tokens / cloud creds / proxy vars leak into the child
    - path-traversal guard on every written file (no abs paths, no `..` escape)
    - stdout/stderr size caps; counts are read from a JSON file, never parsed from stdout

  It is NOT a standalone security boundary. It MUST run only inside the isolated pod and is gated
  by ENABLE_QUALITY_EVAL upstream (orchestrator/quality.py). Never enable code execution outside
  that pod.
"""
from __future__ import annotations

import json
import os
import shutil
import subprocess
import sys
import tempfile
from dataclasses import dataclass, field
from pathlib import Path
from typing import List, Optional

try:  # POSIX only; absent on Windows dev boxes (prod is a Linux pod)
    import resource  # type: ignore
except Exception:  # pragma: no cover - platform dependent
    resource = None  # type: ignore

_RUNNER = Path(__file__).resolve().parent / "sandbox_runner.py"

# Env vars always allowed through to the child (everything else is dropped).
_ENV_ALLOWLIST = {"PATH", "SYSTEMROOT", "SystemRoot", "WINDIR", "LANG", "LC_ALL", "LC_CTYPE", "TZ", "COMSPEC", "PATHEXT"}

DEFAULT_TIMEOUT_SECONDS = float(os.getenv("SANDBOX_TIMEOUT_SECONDS", "30"))
DEFAULT_CPU_SECONDS = int(os.getenv("SANDBOX_CPU_SECONDS", "20"))
DEFAULT_MEM_BYTES = int(os.getenv("SANDBOX_MEM_BYTES", str(512 * 1024 * 1024)))
DEFAULT_FSIZE_BYTES = int(os.getenv("SANDBOX_FSIZE_BYTES", str(32 * 1024 * 1024)))
DEFAULT_NPROC = int(os.getenv("SANDBOX_NPROC", "64"))
_OUTPUT_CAP = 16 * 1024

# FAIL-CLOSED isolation gate. Executing model-generated code is refused UNLESS the operator has
# explicitly asserted, at the process level, that this runtime is the isolated pod described in
# the security model above. This is a SECOND, mandatory confirmation independent of any feature
# flag (e.g. ENABLE_QUALITY_EVAL) — code execution is not a default-on capability that can be
# turned on by an env toggle alone. Set ONLY inside the isolated pod (its manifest) or in an
# ephemeral CI/test runner. Absent → run_tests() raises SandboxIsolationError (no code runs).
SANDBOX_ISOLATION_ENV = "HEICODE_SANDBOX_ISOLATED"


class SandboxIsolationError(RuntimeError):
    """Raised when code execution is attempted without an explicit isolation confirmation."""


def isolation_confirmed() -> bool:
    return os.getenv(SANDBOX_ISOLATION_ENV, "").lower() in {"1", "true", "yes"}


def assert_isolated() -> None:
    """Fail-closed guard — call before any code execution. Raises if isolation isn't confirmed."""
    if not isolation_confirmed():
        raise SandboxIsolationError(
            f"refusing to execute generated code: {SANDBOX_ISOLATION_ENV} is not set. "
            "This capability must run ONLY inside the isolated pod (see "
            "docs/integration/security-boundary.md §8.1). Set "
            f"{SANDBOX_ISOLATION_ENV}=1 in that pod's manifest / CI runner to confirm isolation."
        )


@dataclass
class SandboxFile:
    path: str
    content: str


@dataclass
class SandboxResult:
    total: int = 0
    passed: int = 0
    failed: int = 0
    errored: int = 0
    timed_out: bool = False
    exit_code: Optional[int] = None
    stdout: str = ""
    stderr: str = ""
    error: Optional[str] = None
    details: List[dict] = field(default_factory=list)

    @property
    def pass_rate(self) -> Optional[float]:
        """passed / total * 100, or None when nothing ran (caller decides NaN/coverage)."""
        if self.total <= 0:
            return None
        return 100.0 * self.passed / self.total


def _safe_join(root: Path, rel: str) -> Path:
    # Reject absolute paths and any `..` escape; resolve strictly under root.
    candidate = (root / rel).resolve()
    if not str(candidate).startswith(str(root.resolve())):
        raise ValueError(f"path escapes sandbox: {rel!r}")
    return candidate


def _write_files(root: Path, files: List[SandboxFile]) -> None:
    for f in files:
        if not f.path or os.path.isabs(f.path):
            raise ValueError(f"unsafe file path: {f.path!r}")
        dest = _safe_join(root, f.path)
        dest.parent.mkdir(parents=True, exist_ok=True)
        dest.write_text(f.content, encoding="utf-8")


def _child_env(workdir: Path) -> dict:
    env = {k: v for k, v in os.environ.items() if k in _ENV_ALLOWLIST}
    # Confine writes and bytecode; isolate HOME/TMP to the ephemeral dir.
    env["HOME"] = str(workdir)
    env["TMPDIR"] = str(workdir)
    env["TEMP"] = str(workdir)
    env["TMP"] = str(workdir)
    env["PYTHONDONTWRITEBYTECODE"] = "1"
    env["PYTHONNOUSERSITE"] = "1"
    env["PYTHONPATH"] = ""
    return env


def _preexec():  # pragma: no cover - POSIX only, runs in the child before exec
    if resource is not None:
        resource.setrlimit(resource.RLIMIT_CPU, (DEFAULT_CPU_SECONDS, DEFAULT_CPU_SECONDS))
        try:
            resource.setrlimit(resource.RLIMIT_AS, (DEFAULT_MEM_BYTES, DEFAULT_MEM_BYTES))
        except (ValueError, OSError):
            pass
        resource.setrlimit(resource.RLIMIT_FSIZE, (DEFAULT_FSIZE_BYTES, DEFAULT_FSIZE_BYTES))
        try:
            resource.setrlimit(resource.RLIMIT_NPROC, (DEFAULT_NPROC, DEFAULT_NPROC))
        except (ValueError, OSError):
            pass
    os.setsid()  # own process group, so a timeout can kill the whole tree


def run_tests(
    source_files: List[SandboxFile],
    test_files: List[SandboxFile],
    *,
    timeout_seconds: float = DEFAULT_TIMEOUT_SECONDS,
) -> SandboxResult:
    """Write source + test files into an ephemeral workdir and run the tests in a child process.

    Returns a SandboxResult with pass/fail counts (read from the runner's JSON output). The
    workdir is always removed. Blocking/CPU-bound — callers should offload via asyncio.to_thread.

    FAIL-CLOSED: raises SandboxIsolationError before doing anything if the runtime has not
    explicitly confirmed pod isolation (HEICODE_SANDBOX_ISOLATED). No file is written, no child
    is spawned.
    """
    assert_isolated()
    workdir = Path(tempfile.mkdtemp(prefix="swarm-sbx-"))
    is_posix = os.name == "posix"
    try:
        _write_files(workdir, list(source_files) + list(test_files))
        shutil.copyfile(_RUNNER, workdir / "_runner.py")

        popen_kwargs = dict(
            cwd=str(workdir),
            env=_child_env(workdir),
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
        )
        if is_posix:
            popen_kwargs["preexec_fn"] = _preexec
        else:  # Windows dev: new process group so we can signal the tree
            popen_kwargs["creationflags"] = getattr(subprocess, "CREATE_NEW_PROCESS_GROUP", 0)

        proc = subprocess.Popen([sys.executable, "-I", "_runner.py"], **popen_kwargs)
        result = SandboxResult()
        try:
            out, err = proc.communicate(timeout=timeout_seconds)
            result.exit_code = proc.returncode
            result.stdout = (out or "")[:_OUTPUT_CAP]
            result.stderr = (err or "")[:_OUTPUT_CAP]
        except subprocess.TimeoutExpired:
            result.timed_out = True
            _kill(proc, is_posix)
            out, err = proc.communicate()
            result.stdout = (out or "")[:_OUTPUT_CAP]
            result.stderr = (err or "")[:_OUTPUT_CAP]
            result.error = f"timeout after {timeout_seconds}s"
            return result

        report = workdir / "_result.json"
        if report.exists():
            try:
                data = json.loads(report.read_text(encoding="utf-8"))
                result.total = int(data.get("total", 0))
                result.passed = int(data.get("passed", 0))
                result.failed = int(data.get("failed", 0))
                result.errored = int(data.get("errored", 0))
                result.details = data.get("details", []) or []
                if data.get("fatal"):
                    result.error = str(data["fatal"])
            except Exception as exc:
                result.error = f"unparseable result: {exc!r}"
        else:
            result.error = "no result produced by sandbox runner"
        return result
    finally:
        shutil.rmtree(workdir, ignore_errors=True)


def _kill(proc: "subprocess.Popen", is_posix: bool) -> None:  # pragma: no cover - timing dependent
    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