atlas-iac/scripts/ops/hermes_handoff_exec.py

#!/usr/bin/env python3
"""Bounded, screened, provenance-attested command execution.

It bounds pipes, times out and reaps the process group, and discards truncated output.
"""

from __future__ import annotations

import hashlib
import math
import os
import selectors
import shlex
import shutil
import signal
import stat
import subprocess
import time
from collections.abc import Callable, Mapping, Sequence
from contextlib import suppress
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

from hermes_handoff_policy import GITEA_CLIENT, READ_ONLY, PolicyError, check_argv
from hermes_handoff_redaction import (
    DEFAULT_MAX_BYTES,
    TRUNCATION_NOTE,
    redact,
    safe_text,
    scrub,
)

DEFAULT_COMMAND_TIMEOUT = 25.0
DEFAULT_RUN_DEADLINE = 900.0
DEFAULT_CLEANUP_RESERVE = 90.0
MAX_COMMAND_TIMEOUT = 600.0
MAX_RUN_DEADLINE = 7200.0
MAX_OUTPUT_BYTES = 8 * 1024 * 1024
MAX_CONCURRENCY = 16
READ_CHUNK = 64 * 1024
PIPE_DRAIN_SECONDS = 0.25

SAFE_PATH = "/opt/data/tools/bin:/opt/hermes/.venv/bin:/opt/coordinator:/usr/local/bin:/usr/bin:/bin"
# These exact allowlists are intentionally dense: reviewers compare them as a
# single provenance boundary rather than as extensible application config.
# fmt: off
ENV_PASSTHROUGH = ("KUBECONFIG", "LANG", "LC_ALL", "REQUESTS_CA_BUNDLE", "SSL_CERT_FILE")
FIXED_ENVIRONMENT = {"GIT_ASKPASS": "/opt/coordinator/gitea_askpass.sh",
                     "GIT_CONFIG_GLOBAL": "/dev/null", "GIT_CONFIG_NOSYSTEM": "1",
                     "GIT_TERMINAL_PROMPT": "0", "GITEA_BASE_URL": "https://scm.bstein.dev",
                     "LC_ALL": "C", "PATH": SAFE_PATH}
EXPECTED_PATHS = {
    "kubectl": {"/opt/data/tools/bin/kubectl", "/usr/local/bin/kubectl", "/usr/bin/kubectl"},
    "flux": {"/opt/data/tools/bin/flux", "/usr/local/bin/flux", "/usr/bin/flux"},
    "helm": {"/opt/data/tools/bin/helm", "/usr/local/bin/helm", "/usr/bin/helm"},
    "git": {"/usr/bin/git", "/bin/git"},
    "hermes": {"/opt/hermes/.venv/bin/hermes"},
    "sh": {"/usr/bin/sh", "/bin/sh"},
    GITEA_CLIENT: {GITEA_CLIENT},
}
EXPECTED_SHA256 = {
    "kubectl": {"3d514dbae5dc8c09f773df0ef0f5d449dfad05b3aca5c96b13565f886df345fd"},
    "git": {"a0e562e4bd3c4c79379e91d8c07a10104b2cefe8fac966dc6bd4874a57a807f3"},
    "sh": {"367967c823a0c391e5049b15a67c6a0a629c88b9b6dcdca75ef13ac9d65334b1"},
    "hermes": {"c8419290ef7f1a95f59eadcea9f58f446b1e6535a7f4ffee9f1c0cb1172335e4"},
    GITEA_CLIENT: {"f0943db4e81ce96908e94b0f0e78aa4b5ef95d6c4f7803f9f74a43e761576d85"},
}

DEADLINE_ERROR = "deadline-exceeded"
TIMEOUT_ERROR = "timeout"
OPERATOR = "operator"
SELF = "self"
CONTEXT_AWARE_BINARIES = ("kubectl", "flux", "helm")
POD_COMMAND_PATHS = {"kubectl": "/usr/local/bin/kubectl", "hermes": "/opt/hermes/.venv/bin/hermes",
                     "git": "/usr/bin/git", "sh": "/bin/sh", GITEA_CLIENT: GITEA_CLIENT}
# fmt: on


@dataclass(frozen=True)
class Attestation:
    """The exact outer executable used for a probe."""

    path: str
    sha256: str


@dataclass(frozen=True)
class Capture:
    """Raw bounded process result returned by an executor implementation."""

    returncode: int
    stdout: bytes = b""
    stderr: bytes = b""
    stdout_truncated: bool = False
    stderr_truncated: bool = False
    timed_out: bool = False


@dataclass(frozen=True)
class Outcome:
    """The screened, bounded result of one probe command."""

    argv: tuple[str, ...]
    vantage: str
    returncode: int | None = None
    stdout: str = ""
    stderr: str = ""
    truncated: bool = False
    duration_ms: int = 0
    error: str | None = None
    executable_path: str = ""
    executable_sha256: str = ""

    @property
    def ok(self) -> bool:
        return self.error is None and self.returncode == 0 and not self.truncated

    @property
    def ran(self) -> bool:
        return self.error is None

    @property
    def combined(self) -> str:
        return "\n".join(part for part in (self.stdout, self.stderr) if part)

    def as_dict(self) -> dict[str, object]:
        return scrub(
            {
                "command": redact(shlex.join(self.argv)),
                "vantage": self.vantage,
                "returncode": self.returncode,
                "stdout": self.stdout,
                "stderr": self.stderr,
                "truncated": self.truncated,
                "duration_ms": self.duration_ms,
                "error": self.error,
                "executable_path": self.executable_path,
                "executable_sha256": self.executable_sha256,
            }
        )


@dataclass(frozen=True)
class Vantage:
    """A named identity and transport used to observe a probe."""

    name: str
    prefix: tuple[str, ...] = ()
    inserts: Mapping[str, tuple[str, ...]] = field(default_factory=dict)
    env: Mapping[str, str] = field(default_factory=dict)
    command_paths: Mapping[str, str] = field(default_factory=dict)
    description: str = ""

    def wrap(self, argv: Sequence[str]) -> tuple[str, ...]:
        command = tuple(argv)
        if command:
            replacement = self.command_paths.get(command[0])
            if replacement:
                command = (replacement, *command[1:])
            extra = self.inserts.get(command[0].rsplit("/", 1)[-1])
            if extra:
                command = (command[0], *extra, *command[1:])
        return (*self.prefix, *command)


def _finite_number(name: str, value: Any, minimum: float, maximum: float) -> float:
    if isinstance(value, bool) or not isinstance(value, (int, float)):
        raise ValueError(f"{name} must be numeric")
    converted = float(value)
    if not math.isfinite(converted) or not minimum <= converted <= maximum:
        raise ValueError(
            f"{name} must be finite and between {minimum:g} and {maximum:g}"
        )
    return converted


def _byte_budget(value: Any) -> int:
    if (
        isinstance(value, bool)
        or not isinstance(value, int)
        or not 1 <= value <= MAX_OUTPUT_BYTES
    ):
        raise ValueError(
            f"max_bytes must be an integer between 1 and {MAX_OUTPUT_BYTES}"
        )
    return value


def _concurrency(value: Any) -> int:
    if (
        isinstance(value, bool)
        or not isinstance(value, int)
        or not 1 <= value <= MAX_CONCURRENCY
    ):
        raise ValueError(
            f"concurrency must be an integer between 1 and {MAX_CONCURRENCY}"
        )
    return value


def build_environment(
    vantage: Vantage, source: Mapping[str, str] | None = None
) -> dict[str, str]:
    """Build a deterministic child environment without caller-controlled PATH."""
    origin = os.environ if source is None else source
    environment = {name: origin[name] for name in ENV_PASSTHROUGH if name in origin}
    environment.update(FIXED_ENVIRONMENT)
    environment.update(vantage.env)
    environment["PATH"] = SAFE_PATH
    environment["GITEA_BASE_URL"] = FIXED_ENVIRONMENT["GITEA_BASE_URL"]
    environment["GIT_ASKPASS"] = FIXED_ENVIRONMENT["GIT_ASKPASS"]
    return environment


# fmt: off
def operator_vantage(kubeconfig: str | None = None, context: str | None = None) -> Vantage:
    inserts = dict.fromkeys(CONTEXT_AWARE_BINARIES, ("--context", context)) if context else {}
    detail = " ".join(part for part in (f"context={context}" if context else "",
                      "kubeconfig-pinned" if kubeconfig else "") if part)
    return Vantage(name=OPERATOR, inserts=inserts,
                    env={"KUBECONFIG": kubeconfig} if kubeconfig else {},
                    description=f"operator {detail}".strip())


def pod_vantage(namespace: str, pod: str, container: str, operator: Vantage | None = None) -> Vantage:
    reach = operator.inserts.get("kubectl", ()) if operator else ()
    return Vantage(name=SELF, prefix=("kubectl", *reach, "exec", "--namespace", namespace,
                    pod, "--container", container, "--"),
                    env=dict(operator.env) if operator else {}, command_paths=POD_COMMAND_PATHS,
                    description=f"in-pod {namespace}/{pod}[{container}]")
# fmt: on


def attest_executable(command: str, environment: Mapping[str, str]) -> Attestation:
    """Resolve and hash an executable from a fixed, trusted path set."""
    name = command if command == GITEA_CLIENT else command.rsplit("/", 1)[-1]
    candidate = (
        command if "/" in command else shutil.which(command, path=environment["PATH"])
    )
    if not candidate:
        raise OSError(f"trusted executable is unavailable: {name}")
    resolved = str(Path(candidate).resolve(strict=True))
    allowed = {
        str(Path(item).resolve(strict=False))
        for item in EXPECTED_PATHS.get(name, set())
    }
    if resolved not in allowed:
        raise OSError(f"executable resolved outside its trusted paths: {name}")
    details = os.stat(resolved)
    if (
        not stat.S_ISREG(details.st_mode)
        or details.st_mode & 0o022
        or not details.st_mode & 0o111
    ):
        raise OSError(f"trusted executable has unsafe ownership mode: {name}")
    digest = hashlib.sha256()
    with open(resolved, "rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    observed = digest.hexdigest()
    expected = EXPECTED_SHA256.get(name)
    if not expected or observed not in expected:
        raise OSError(f"executable digest is not release-attested: {name}")
    return Attestation(resolved, observed)


def _append_bounded(buffer: bytearray, chunk: bytes, limit: int) -> bool:
    remaining = limit + 1 - len(buffer)
    if remaining > 0:
        buffer.extend(chunk[:remaining])
    return len(buffer) > limit or len(chunk) > remaining


def _signal_group(process: subprocess.Popen[bytes], signum: int) -> None:
    with suppress(ProcessLookupError):
        os.killpg(process.pid, signum)


def execute_bounded(
    argv: Sequence[str],
    environment: Mapping[str, str],
    timeout: float,
    max_bytes: int,
    clock: Callable[[], float] = time.monotonic,
    popen: Callable[..., subprocess.Popen[bytes]] = subprocess.Popen,
) -> Capture:
    """Execute one process with bounded live pipe capture and group cleanup."""
    process = popen(
        list(argv),
        stdin=subprocess.DEVNULL,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        env=dict(environment),
        text=False,
        close_fds=True,
        start_new_session=True,
    )
    streams = {"stdout": process.stdout, "stderr": process.stderr}
    buffers = {"stdout": bytearray(), "stderr": bytearray()}
    truncated = {"stdout": False, "stderr": False}
    selector = selectors.DefaultSelector()
    for name, stream in streams.items():
        if (
            stream is not None
        ):  # pragma: no branch - PIPE was requested for both streams
            os.set_blocking(stream.fileno(), False)
            selector.register(stream, selectors.EVENT_READ, name)
    expires = clock() + timeout
    timed_out = False
    try:
        while process.poll() is None:
            remaining = expires - clock()
            if remaining <= 0:
                timed_out = True
                break
            for key, _ in selector.select(min(remaining, 0.1)):
                try:
                    chunk = os.read(key.fileobj.fileno(), READ_CHUNK)
                except BlockingIOError:  # pragma: no cover - readiness race, retry only
                    continue
                if not chunk:
                    selector.unregister(key.fileobj)
                    continue
                name = key.data
                truncated[name] |= _append_bounded(buffers[name], chunk, max_bytes)
        _signal_group(process, signal.SIGKILL if timed_out else signal.SIGTERM)
        drain_until = clock() + PIPE_DRAIN_SECONDS
        while selector.get_map() and clock() < drain_until:
            for key, _ in selector.select(max(0.0, drain_until - clock())):
                try:
                    chunk = os.read(key.fileobj.fileno(), READ_CHUNK)
                except BlockingIOError:  # pragma: no cover - readiness race, retry only
                    continue
                if not chunk:
                    selector.unregister(key.fileobj)
                    continue
                name = key.data
                truncated[name] |= _append_bounded(buffers[name], chunk, max_bytes)
        _signal_group(process, signal.SIGKILL)
        returncode = process.wait(timeout=PIPE_DRAIN_SECONDS)
    finally:
        selector.close()
        for stream in streams.values():
            if (
                stream is not None
            ):  # pragma: no branch - PIPE was requested for both streams
                stream.close()
        if (
            process.poll() is None
        ):  # pragma: no branch - fail-safe for wait/selector exceptions
            _signal_group(process, signal.SIGKILL)
            process.wait()
    return Capture(
        returncode=returncode,
        stdout=bytes(buffers["stdout"][:max_bytes]),
        stderr=bytes(buffers["stderr"][:max_bytes]),
        stdout_truncated=truncated["stdout"],
        stderr_truncated=truncated["stderr"],
        timed_out=timed_out,
    )


def _injected_executor(
    spawn: Callable[..., subprocess.CompletedProcess[Any]],
) -> Callable[..., Capture]:
    """Adapt deterministic test doubles without weakening the production path."""

    def execute(
        argv: Sequence[str],
        environment: Mapping[str, str],
        timeout: float,
        max_bytes: int,
    ) -> Capture:
        completed = spawn(
            list(argv),
            capture_output=True,
            text=True,
            timeout=timeout,
            check=False,
            env=dict(environment),
        )
        stdout = str(completed.stdout or "").encode()
        stderr = str(completed.stderr or "").encode()
        return Capture(
            completed.returncode,
            stdout[:max_bytes],
            stderr[:max_bytes],
            len(stdout) > max_bytes,
            len(stderr) > max_bytes,
        )

    return execute


def _screen_capture(data: bytes, truncated: bool, budget: int) -> str:
    if truncated:
        return TRUNCATION_NOTE.encode()[:budget].decode()
    return safe_text(data.decode("utf-8", "replace"), budget)[0]


class Runner:
    """Run probes under policy, provenance, timeout, deadline, and byte budgets."""

    def __init__(
        self,
        mode: str = READ_ONLY,
        command_timeout: float = DEFAULT_COMMAND_TIMEOUT,
        deadline_seconds: float = DEFAULT_RUN_DEADLINE,
        max_bytes: int = DEFAULT_MAX_BYTES,
        concurrency: int = 1,
        clock: Callable[[], float] = time.monotonic,
        executor: Callable[..., Capture] | None = None,
        spawn: Callable[..., subprocess.CompletedProcess[Any]] | None = None,
        environ: Mapping[str, str] | None = None,
        attestor: Callable[[str, Mapping[str, str]], Attestation] = attest_executable,
    ) -> None:
        self.mode = mode
        self.command_timeout = _finite_number(
            "command_timeout", command_timeout, 0.05, MAX_COMMAND_TIMEOUT
        )
        deadline = _finite_number(
            "deadline_seconds", deadline_seconds, 0.05, MAX_RUN_DEADLINE
        )
        self.max_bytes = _byte_budget(max_bytes)
        self.concurrency = _concurrency(concurrency)
        self._clock = clock
        self._executor = executor or (
            _injected_executor(spawn) if spawn else execute_bounded
        )
        self._environ = environ
        self._attestor = attestor
        self._deadline_at = clock() + deadline
        self.commands_run = 0

    @property
    def remaining_seconds(self) -> float:
        return self._deadline_at - self._clock()

    def run(
        self,
        argv: Sequence[str],
        vantage: Vantage,
        max_bytes: int | None = None,
        reserve_seconds: float = 0.0,
    ) -> Outcome:
        budget_bytes = self.max_bytes if max_bytes is None else _byte_budget(max_bytes)
        reserve = _finite_number(
            "reserve_seconds", reserve_seconds, 0.0, MAX_RUN_DEADLINE
        )
        addressed = vantage.wrap(argv)
        try:
            check_argv(addressed, self.mode)
        except PolicyError as exc:
            return Outcome(argv=addressed, vantage=vantage.name, error=f"policy: {exc}")
        budget = min(self.command_timeout, self.remaining_seconds - reserve)
        if budget <= 0:
            return Outcome(argv=addressed, vantage=vantage.name, error=DEADLINE_ERROR)
        environment = build_environment(vantage, self._environ)
        try:
            attestation = self._attestor(addressed[0], environment)
        except (OSError, ValueError) as exc:
            return Outcome(
                argv=addressed, vantage=vantage.name, error=f"attestation: {exc}"
            )
        executed = (attestation.path, *addressed[1:])
        try:
            check_argv(executed, self.mode)
        except PolicyError as exc:
            return Outcome(argv=addressed, vantage=vantage.name, error=f"policy: {exc}")
        started = self._clock()
        self.commands_run += 1
        try:
            capture = self._executor(executed, environment, budget, budget_bytes)
        except subprocess.TimeoutExpired:
            capture = Capture(-signal.SIGKILL, timed_out=True)
        except OSError as exc:
            return Outcome(
                argv=addressed,
                vantage=vantage.name,
                error=f"spawn: {exc.strerror or type(exc).__name__}",
                executable_path=attestation.path,
                executable_sha256=attestation.sha256,
            )
        was_truncated = capture.stdout_truncated or capture.stderr_truncated
        stdout = _screen_capture(capture.stdout, capture.stdout_truncated, budget_bytes)
        stderr = _screen_capture(capture.stderr, capture.stderr_truncated, budget_bytes)
        return Outcome(
            argv=addressed,
            vantage=vantage.name,
            returncode=capture.returncode,
            stdout=stdout,
            stderr=stderr,
            truncated=was_truncated,
            duration_ms=int((self._clock() - started) * 1000),
            error=TIMEOUT_ERROR if capture.timed_out else None,
            executable_path=attestation.path,
            executable_sha256=attestation.sha256,
        )