UmbraAudio/infra/remote_credential_broker/executor.py

from __future__ import annotations

import os
import shlex
import shutil
import socket
import subprocess
import tempfile
import threading
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Callable, Sequence

from .config import ConfigError, RemoteProfile
from .keychain import CredentialError, CredentialVault

Runner = Callable[..., subprocess.CompletedProcess[bytes]]


class ExecutionError(RuntimeError):
    """Raised when a remote operation cannot be safely started."""


@dataclass(frozen=True)
class ExecutionResult:
    exit_code: int
    stdout: str
    stderr: str
    duration_ms: int


class SecretRedactor:
    """Redact known secrets before output leaves the broker."""

    def __init__(self) -> None:
        self._secrets: set[str] = set()

    def add(self, secret: str | None) -> None:
        if secret:
            self._secrets.add(secret)

    def redact(self, value: str) -> str:
        result = value
        for secret in sorted(self._secrets, key=len, reverse=True):
            if secret:
                result = result.replace(secret, "<redacted>")
        return result


class _AskpassServer:
    """One-shot local channel used instead of exposing a Keychain item to ssh."""

    def __init__(self, secret: str) -> None:
        self._directory = Path(tempfile.mkdtemp(prefix="remote-broker-askpass-"))
        self._directory.chmod(0o700)
        self.path = self._directory / "secret.sock"
        self._secret = secret.encode("utf-8")
        self._stop = threading.Event()
        self._socket = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        self._socket.bind(str(self.path))
        self._socket.listen(8)
        self._socket.settimeout(0.2)
        self._thread = threading.Thread(target=self._serve, name="remote-broker-askpass", daemon=True)
        self._thread.start()

    def _serve(self) -> None:
        while not self._stop.is_set():
            try:
                connection, _ = self._socket.accept()
            except socket.timeout:
                continue
            except OSError:
                return
            with connection:
                try:
                    connection.sendall(self._secret)
                except OSError:
                    pass

    def close(self) -> None:
        self._stop.set()
        try:
            self._socket.close()
        except OSError:
            pass
        self._thread.join(timeout=1)
        shutil.rmtree(self._directory, ignore_errors=True)


class RemoteExecutor:
    """Run bounded SSH operations; sudo commands are exact allowlist entries."""

    def __init__(
        self,
        store: CredentialVault,
        *,
        runner: Runner = subprocess.run,
        askpass_path: Path | None = None,
        audit_path: Path | None = None,
    ) -> None:
        self._store = store
        self._runner = runner
        self._askpass_path = askpass_path or Path(__file__).with_name("askpass_helper.py")
        if self._askpass_path.exists():
            self._askpass_path.chmod(0o700)
        self._audit_path = audit_path

    @staticmethod
    def _validate_command(command: Sequence[str]) -> tuple[str, ...]:
        if not command or not all(isinstance(part, str) and part for part in command):
            raise ExecutionError("Remote command must be a non-empty list of strings")
        if any("\x00" in part for part in command):
            raise ExecutionError("Remote command contains a NUL byte")
        return tuple(command)

    def _audit(self, profile: RemoteProfile, command: tuple[str, ...], sudo: bool, result: int) -> None:
        if self._audit_path is None:
            return
        self._audit_path.parent.mkdir(parents=True, exist_ok=True)
        # Store command metadata only. Never store passwords or command stdin.
        line = (
            f"profile={profile.name} sudo={str(sudo).lower()} "
            f"command={shlex.join(command)} exit={result}\n"
        )
        fd = os.open(self._audit_path, os.O_WRONLY | os.O_CREAT | os.O_APPEND, 0o600)
        try:
            os.write(fd, line.encode("utf-8"))
        finally:
            os.close(fd)

    def _base_ssh_command(self, profile: RemoteProfile, remote_command: str, use_password: bool) -> list[str]:
        known_hosts = str(Path(profile.known_hosts).expanduser())
        command = [
            "ssh",
            "-p",
            str(profile.port),
            "-o",
            "ConnectTimeout=15",
            "-o",
            "ServerAliveInterval=15",
            "-o",
            "ServerAliveCountMax=2",
            "-o",
            "StrictHostKeyChecking=yes",
            "-o",
            f"UserKnownHostsFile={known_hosts}",
        ]
        if use_password:
            command.extend(
                [
                    "-o",
                    "BatchMode=no",
                    "-o",
                    "PreferredAuthentications=publickey,password",
                    "-o",
                    "PubkeyAuthentication=yes",
                ]
            )
        command.extend([f"{profile.username}@{profile.host}", remote_command])
        return command

    def run(
        self,
        profile: RemoteProfile,
        command: Sequence[str],
        *,
        sudo: bool = False,
        timeout_seconds: int = 120,
    ) -> ExecutionResult:
        command_tuple = self._validate_command(command)
        if not 5 <= timeout_seconds <= 3600:
            raise ExecutionError("timeout_seconds must be between 5 and 3600")
        if sudo and not profile.sudo_allowed(command_tuple):
            raise ExecutionError(
                "The requested sudo command is not an exact allowlist entry for this profile"
            )

        ssh_secret: str | None = None
        sudo_secret: str | None = None
        redactor = SecretRedactor()
        if profile.ssh_credential:
            try:
                ssh_secret = self._store.get(profile.ssh_credential)
            except CredentialError as exc:
                raise ExecutionError(
                    f"SSH credential is not configured for profile {profile.name!r}"
                ) from exc
            redactor.add(ssh_secret)
        if sudo:
            if not profile.sudo_credential:
                raise ExecutionError(f"No sudo credential is configured for profile {profile.name!r}")
            try:
                sudo_secret = self._store.get(profile.sudo_credential)
            except CredentialError as exc:
                raise ExecutionError(
                    f"Sudo credential is not configured for profile {profile.name!r}"
                ) from exc
            redactor.add(sudo_secret)

        remote_argv = list(command_tuple)
        if sudo:
            remote_command = "sudo -S -p '' -- " + shlex.join(remote_argv)
        else:
            remote_command = shlex.join(remote_argv)
        ssh_command = self._base_ssh_command(profile, remote_command, use_password=ssh_secret is not None)
        environment = os.environ.copy()
        askpass_server: _AskpassServer | None = None
        if ssh_secret is not None:
            askpass_server = _AskpassServer(ssh_secret)
            environment.update(
                {
                    "SSH_ASKPASS": str(self._askpass_path),
                    "SSH_ASKPASS_REQUIRE": "force",
                    "DISPLAY": environment.get("DISPLAY", ":0"),
                    "UMBRA_BROKER_ASKPASS_SOCKET": str(askpass_server.path),
                }
            )
        stdin_data = (sudo_secret + "\n").encode("utf-8") if sudo_secret is not None else None
        started = time.monotonic()
        try:
            completed = self._runner(
                ssh_command,
                input=stdin_data,
                capture_output=True,
                timeout=timeout_seconds,
                env=environment,
                check=False,
            )
        except subprocess.TimeoutExpired as exc:
            raise ExecutionError("Remote operation timed out") from exc
        except OSError as exc:
            raise ExecutionError("Unable to start ssh") from exc
        finally:
            if askpass_server is not None:
                askpass_server.close()
        duration_ms = int((time.monotonic() - started) * 1000)
        stdout = redactor.redact(completed.stdout.decode("utf-8", errors="replace"))
        stderr = redactor.redact(completed.stderr.decode("utf-8", errors="replace"))
        result = ExecutionResult(completed.returncode, stdout, stderr, duration_ms)
        self._audit(profile, command_tuple, sudo, result.exit_code)
        return result