#!/usr/bin/env python3
"""units.py -- every process the harness starts runs in its own `systemd-run --user` unit; the engine's unit also runs
inside a no-network sandbox. Builders and readers only: nothing here decides a gate.

THE EGRESS BLOCK, AS MEASURED ON the laptop (2026-10-02, systemd 259, the dry run's receipts):
  * `IPAddressDeny=any` + `IPAddressAllow=localhost` on a USER unit is NOT enforced: the user manager logs "unit
    configures an IP firewall, but not running as root" and a probe inside such a unit connects to the box's own
    non-loopback address. The properties stay on every engine unit (the recon's recipe, step 8) and every record
    says they are configured and not enforced.
  * The block that IS enforced is the house's no-network sandbox (estate bench/kits/vllm-seats/egress.py
    `sandbox_argv`, the phone-home probe's recipe, minus its strace): `bwrap --unshare-net` (only `lo` exists), a
    tmpfs over /run (systemd-resolved's and avahi's Unix sockets and the system and user buses cannot be reached)
    and over /tmp, everything else read-only except the arm's own run directory. The engine AND its client run in
    that one sandbox, so the client reaches the server on loopback and nothing else is reachable.
  * inner.py's probe proves it inside the unit before any engine starts (only `lo`, no DNS, a non-loopback connect
    fails, loopback works); a probe that can reach anything refuses the arm (`egress_mode="ipfilter"` runs without
    the sandbox and is refused by that same probe: the dry run fires it once).
"""
import os
import shlex
import subprocess
import time

#: The systemd properties every engine unit carries (MemoryMax is per arm: common.memory_plan).
IP_PROPS = ("IPAddressDeny=any", "IPAddressAllow=localhost")
PERSISTENCED = "/run/nvidia-persistenced"


def bwrap_argv(inner, *, rw_binds, ro_binds=(), bwrap="bwrap", var_run="/var/run", persistenced=PERSISTENCED):
    """(argv, notes): egress.sandbox_argv's sandbox around `inner`, without strace (a traced engine would measure the
    tracer). `ro_binds` puts back, read-only, anything the tmpfs mounts would hide."""
    notes = []
    argv = [bwrap, "--unshare-net", "--unshare-pid", "--die-with-parent", "--ro-bind", "/", "/",
            "--dev-bind", "/dev", "/dev", "--proc", "/proc", "--tmpfs", "/run"]
    if os.path.isdir(var_run) and not os.path.islink(var_run):
        argv += ["--tmpfs", var_run]
    else:
        notes.append(f"{var_run} is a link to /run (or absent): the /run tmpfs covers it")
    argv += ["--tmpfs", "/tmp"]
    rw = _collapse(rw_binds)
    for d in rw:
        argv += ["--bind", d, d]
    for d in _collapse(ro_binds):
        if not any(d == q or d.startswith(q.rstrip("/") + "/") for q in rw):
            argv += ["--ro-bind", d, d]
            notes.append(f"{d} bound back read-only")
    if os.path.isdir(persistenced):
        argv += ["--ro-bind", persistenced, persistenced]
        notes.append(f"{persistenced} bound read-only (the driver's persistence socket)")
    argv += ["--"] + list(inner)
    return argv, notes


def _collapse(paths):
    real = sorted({os.path.abspath(p) for p in paths if p and os.path.isdir(p)})
    out = []
    for p in real:
        if not any(p == q or p.startswith(q.rstrip("/") + "/") for q in out):
            out.append(p)
    return out


def hidden_by_tmpfs(paths):
    """Paths the sandbox's /tmp or /run tmpfs would hide (an interpreter or a tree that lives there)."""
    out = []
    for p in paths:
        if p and (p == "/tmp" or p.startswith("/tmp/") or p.startswith("/run/")):
            out.append(p)
    return out


def systemd_run_argv(unit, cmd, *, memory_max_gib=None, swap_max=None, ip_props=True, workdir=None, log=None,
                     env=None, extra_props=(), runtime_max_s=None):
    """The `systemd-run --user` line for one unit (a service, so `systemctl --user stop` takes the whole cgroup)."""
    argv = ["systemd-run", "--user", "--quiet", f"--unit={unit}", "--property=KillMode=control-group",
            "--property=TimeoutStopSec=30"]
    if memory_max_gib is not None:
        argv.append(f"--property=MemoryMax={int(memory_max_gib)}G")
    if swap_max is not None:
        argv.append(f"--property=MemorySwapMax={swap_max}")
    if ip_props:
        argv += [f"--property={p}" for p in IP_PROPS]
    if runtime_max_s:
        argv.append(f"--property=RuntimeMaxSec={int(runtime_max_s)}")
    if workdir:
        argv.append(f"--working-directory={workdir}")
    if log:
        argv += [f"--property=StandardOutput=append:{log}", f"--property=StandardError=append:{log}"]
    for p in extra_props:
        argv.append(f"--property={p}")
    for k, v in (env or {}).items():
        argv.append(f"--setenv={k}={v}")
    return argv + ["--"] + list(cmd)


def redact(argv, secrets):
    """The argv as it may be written down: every secret value replaced (an API key passes by environment only)."""
    out = []
    for a in argv:
        for s in secrets:
            if s and s in a:
                a = a.replace(s, "<redacted>")
        out.append(a)
    return out


def start(argv):
    r = subprocess.run(argv, capture_output=True, text=True, timeout=60)
    return r.returncode, (r.stdout + r.stderr).strip()


def show(unit, props=("ActiveState", "SubState", "Result", "ExecMainStatus", "ExecMainPID", "ControlGroup",
                      "MemoryMax", "MemorySwapMax", "MemoryPeak", "MemoryCurrent", "IPAddressDeny", "IPAddressAllow",
                      "NRestarts")):
    r = subprocess.run(["systemctl", "--user", "show", unit, "-p", ",".join(props)], capture_output=True, text=True,
                       timeout=30)
    out = {}
    for line in r.stdout.splitlines():
        if "=" in line:
            k, v = line.split("=", 1)
            out[k] = v
    return out


def active(unit):
    return show(unit, ("ActiveState",)).get("ActiveState") in ("active", "activating", "reloading")


def stop(unit, wait_s=40):
    subprocess.run(["systemctl", "--user", "stop", unit], capture_output=True, text=True, timeout=wait_s + 10)
    end = time.monotonic() + wait_s
    while time.monotonic() < end and active(unit):
        time.sleep(0.5)
    return show(unit)


def reset_failed(unit):
    subprocess.run(["systemctl", "--user", "reset-failed", unit], capture_output=True, text=True, timeout=30)


def cgroup_dir(unit):
    cg = show(unit, ("ControlGroup",)).get("ControlGroup")
    return os.path.join("/sys/fs/cgroup", cg.lstrip("/")) if cg else None


def cgroup_pids(cgdir):
    pids = []
    if not cgdir:
        return pids
    for root, _dirs, files in os.walk(cgdir):
        if "cgroup.procs" in files:
            try:
                pids += [int(x) for x in open(os.path.join(root, "cgroup.procs")).read().split()]
            except (OSError, ValueError):
                pass
    return pids


def cgroup_memory(cgdir):
    """memory.current / memory.peak / memory.events / memory.stat (anon, file, kernel) of the unit, or {}."""
    out = {}
    if not cgdir:
        return out
    for name in ("memory.current", "memory.peak", "memory.max", "memory.swap.current"):
        try:
            out[name] = open(os.path.join(cgdir, name)).read().strip()
        except OSError:
            pass
    for name in ("memory.events", "memory.stat"):
        try:
            kv = {}
            for line in open(os.path.join(cgdir, name)):
                k, v = line.split()
                if name == "memory.events" or k in ("anon", "file", "kernel", "shmem", "file_mapped", "unevictable"):
                    kv[k] = int(v)
            out[name] = kv
        except (OSError, ValueError):
            pass
    return out


def proc_status(pid, keys=("Name", "VmLck", "VmPin", "VmRSS", "VmHWM", "RssAnon", "RssFile")):
    out = {}
    try:
        for line in open(f"/proc/{pid}/status"):
            k, _, v = line.partition(":")
            if k in keys:
                v = v.strip()
                out[k] = int(v.split()[0]) * 1024 if v.endswith("kB") else v
    except OSError:
        return None
    return out


def quote(argv):
    return " ".join(shlex.quote(a) for a in argv)