nisse/nisse-dev/src/nisse_dev/execution.py

"""Running commands on a tool VM over SSH (architecture L13).

``SSHHost.execute`` is the only way the agent runs anything: there is no
local execution path, so the host the agent lives on never runs a command a
model chose. The harness binds the host; a model never names one.

How a command travels. The remote side runs ``bash -s`` and reads a small
script from SSH's stdin: the secrets handed to this one call as ``export``
lines, a ``cd``, then the command under ``timeout``. Nothing secret is ever on
a command line (``ps`` and ``/proc/*/cmdline`` would show it); the values exist
only in that shell and the command's environment. The command's own stdin is
``/dev/null``, or a here-document carrying ``stdin``, so it can never read the
rest of the script.

Every value handed out is redacted from what comes back (L12): a test that
prints its environment must not put a key into the transcript.

The connection is reused across calls (OpenSSH's ControlMaster), so a tool
call costs one round trip, not a handshake.
"""

import os
import secrets as _random
import shlex
import subprocess
import tempfile
from collections.abc import Iterable, Mapping
from dataclasses import dataclass
from pathlib import Path

# The remote command's exit code is passed back as is, except that 255 is
# what ssh itself uses for "could not connect"; a command exiting 255 is
# reported as 254 so the two can't be confused.
SSH_FAILED = 255
_TIMED_OUT = 124  # coreutils `timeout`
_KILL_GRACE_S = 5
# The local backstop, beyond the remote timeout, for a connection that hangs.
_LOCAL_SLACK_S = 30

REDACTED = "[redacted]"


class HostUnavailable(RuntimeError):
    """The host could not be reached: not something the command did."""


@dataclass(frozen=True)
class Result:
    exit_code: int
    stdout: str
    stderr: str
    timed_out: bool = False

    @property
    def ok(self) -> bool:
        return self.exit_code == 0 and not self.timed_out


def redact(text: str, values: Iterable[str]) -> str:
    """Mask every non-trivial secret value in ``text``. Longest first, so a
    value containing another is masked whole."""
    for value in sorted((v for v in values if len(v) >= 4), key=len, reverse=True):
        text = text.replace(value, REDACTED)
    return text


def _heredoc(body: str) -> tuple[str, str]:
    """A here-document redirection whose delimiter can't occur in ``body``."""
    while True:
        tag = f"NISSE_EOF_{_random.token_hex(8)}"
        if tag not in body:
            return f"<<'{tag}'", f"\n{body}\n{tag}"


def _script(
    command: str,
    cwd: str | None,
    secrets: Mapping[str, str],
    stdin: str | None,
    timeout_s: int,
) -> str:
    lines = [f"export {name}={shlex.quote(value)}" for name, value in secrets.items()]
    if cwd is not None:
        lines.append(f"cd -- {shlex.quote(cwd)} || exit 125")
    redirect, body = ("< /dev/null", "") if stdin is None else _heredoc(stdin)
    # `timeout` runs the command in its own process group and signals the
    # whole group, so a backgrounded grandchild doesn't outlive it.
    lines.append(
        f"timeout --kill-after={_KILL_GRACE_S} {timeout_s} "
        f"bash -c {shlex.quote(command)} {redirect}{body}"
    )
    lines.append("rc=$?; [ $rc -eq 255 ] && rc=254; exit $rc")
    return "\n".join(lines) + "\n"


class SSHHost:
    """One tool VM, reached as ``user@address`` with a private key."""

    def __init__(
        self,
        address: str,
        *,
        key: Path,
        state_dir: Path,
        user: str = "root",
        port: int = 22,
    ) -> None:
        self.address = address
        self.user = user
        self.port = port
        self._key = Path(key)
        state_dir = Path(state_dir)
        state_dir.mkdir(parents=True, exist_ok=True)
        self._known_hosts = state_dir / "known_hosts"
        # A Unix socket's path is limited to about 100 bytes, so the shared
        # connection's socket lives in a short private directory rather than
        # beside the state, named by %C, ssh's hash of the connection.
        sockets = Path(tempfile.gettempdir()) / f"nisse-ssh-{os.getuid()}"
        sockets.mkdir(mode=0o700, exist_ok=True)
        self._control = sockets / "%C"

    def describe(self) -> str:
        return f"{self.user}@{self.address}"

    def _ssh(self) -> list[str]:
        return [
            "ssh",
            "-i",
            str(self._key),
            "-p",
            str(self.port),
            "-o",
            "BatchMode=yes",
            "-o",
            "IdentitiesOnly=yes",
            "-o",
            f"UserKnownHostsFile={self._known_hosts}",
            "-o",
            "StrictHostKeyChecking=accept-new",
            "-o",
            "ControlMaster=auto",
            "-o",
            f"ControlPath={self._control}",
            "-o",
            "ControlPersist=120",
            "-o",
            "ConnectTimeout=15",
            "-o",
            "ServerAliveInterval=15",
            "-o",
            "LogLevel=ERROR",
            f"{self.user}@{self.address}",
            "bash -s",
        ]

    def execute(
        self,
        command: str,
        *,
        cwd: str | None = None,
        secrets: Mapping[str, str] | None = None,
        stdin: str | None = None,
        timeout_s: int = 120,
    ) -> Result:
        """Run ``command`` with bash in ``cwd``. ``secrets`` become environment
        variables of this one command and are redacted from its output.
        Raises HostUnavailable when the host can't be reached."""
        handed_out = dict(secrets or {})
        script = _script(command, cwd, handed_out, stdin, timeout_s)
        try:
            completed = subprocess.run(
                self._ssh(),
                input=script.encode(),
                capture_output=True,
                timeout=timeout_s + _KILL_GRACE_S + _LOCAL_SLACK_S,
                check=False,
            )
        except subprocess.TimeoutExpired as exc:
            raise HostUnavailable(
                f"{self.describe()} stopped responding during a command"
            ) from exc
        except OSError as exc:
            raise HostUnavailable(f"could not run ssh: {exc}") from exc
        stdout = completed.stdout.decode(errors="replace")
        stderr = completed.stderr.decode(errors="replace")
        if completed.returncode == SSH_FAILED:
            raise HostUnavailable(
                f"could not reach {self.describe()}: "
                + redact(stderr.strip() or "ssh failed", handed_out.values())
            )
        values = handed_out.values()
        return Result(
            exit_code=completed.returncode,
            stdout=redact(stdout, values),
            stderr=redact(stderr, values),
            timed_out=completed.returncode in (_TIMED_OUT, 137),
        )

    def close(self) -> None:
        """End the shared connection, if one is open."""
        args = self._ssh()[:-2]  # without the destination and the command
        subprocess.run(
            [*args, "-O", "exit", f"{self.user}@{self.address}"],
            capture_output=True,
            timeout=10,
            check=False,
        )