admin / Synapse-Cortex

"""SSH execution engine, behind a pluggable interface.

Two executors exist, each tied to a specific operator action:

* ``SimulatedSSHExecutor`` powers the **Simulate** button. It never opens a
  network connection and never inspects/logs the secret - it logs exactly what
  it *would* have run and returns a synthetic success. This is the safe
  rehearsal an operator reviews before executing live.

* ``ParamikoSSHExecutor`` powers **Execute Live on Client**. It opens a real
  SSH connection and actually runs the command on the client, capturing
  stdout/stderr and the exit status - this is how Cortex actively fixes a host.
  It auto-detects, per host, whether the login account is root: root accounts
  run the command directly, non-root accounts run it under ``sudo`` (password
  supplied on stdin), so a mixed fleet of root and non-root hosts just works.

Live execution is available on every deployment; nothing reaches a host by
accident because it is always gated behind an explicit human confirmation AND a
prior simulation of the same run (see app/routers/remediation.py).
"""
import os
import shlex
import time
from dataclasses import dataclass
from typing import Optional, Protocol

# Cap stored per-command output so a chatty command can't bloat the DB row / UI.
_MAX_OUTPUT_CHARS = 20000


@dataclass
class CommandResult:
    command: str
    output: str
    exit_status: int
    duration_ms: int


def sudo_enabled() -> bool:
    """Master switch for privilege elevation. Default ON: on a host where the
    login account is **not** root, Cortex runs the command through sudo. On a
    host where the account **is** root, no sudo is used regardless (detected at
    run time). Set SSH_USE_SUDO=false to never use sudo (e.g. a locked-down
    account that must run commands exactly as given)."""
    return os.getenv("SSH_USE_SUDO", "true").strip().lower() in ("1", "true", "yes", "on")


def build_execution_command(command: str) -> str:
    """Wrap a command so the WHOLE command runs with root privileges via sudo:
    it runs inside a shell under ``sudo -S`` (so pipes and redirects are
    elevated too), and the password is read from stdin with an empty prompt -
    the secret is never placed in the command text."""
    shell = os.getenv("SSH_SUDO_SHELL", "bash")
    return f"sudo -S -p '' -- {shell} -c {shlex.quote(command)}"


def sudo_password(secret: str) -> str:
    """Password fed to ``sudo -S`` on stdin: the connection secret by default
    (a user's login password is normally their own sudo password), overridable
    with SSH_SUDO_PASSWORD."""
    return os.getenv("SSH_SUDO_PASSWORD") or secret


class SSHExecutor(Protocol):
    def run(self, *, host: str, port: int, username: str, secret: str, command: str) -> CommandResult: ...


class SimulatedSSHExecutor:
    """Powers the Simulate button. Never opens a network connection and never
    inspects/logs the secret - it logs exactly what it *would* run. It can't
    know the target's root-ness without connecting, so it notes that sudo is
    applied automatically for a non-root account."""

    def run(self, *, host: str, port: int, username: str, secret: str, command: str) -> CommandResult:
        start = time.monotonic()
        output = (
            f"[SIMULATED] would run: {command!r} on {username}@{host}:{port} "
            f"(elevated with sudo automatically if that account is non-root)"
        )
        duration_ms = int((time.monotonic() - start) * 1000)
        return CommandResult(command=command, output=output, exit_status=0, duration_ms=duration_ms)


class ParamikoSSHExecutor:
    """Real SSH execution over paramiko. Opens a connection to
    ``username@host:port``, authenticates with the vaulted secret as a
    password, and runs the (already guardrail-checked) command.

    Privilege handling: on the first command for a given ``(host, username)``
    it probes ``id -u`` to learn whether the account is root. A root account
    runs commands directly; a non-root account runs them under ``sudo`` (the
    whole command in a shell, password on stdin) - so a fleet mixing root and
    non-root hosts is handled automatically, with no per-host configuration.
    ``SSH_USE_SUDO=false`` forces direct execution everywhere.

    Safety choices:

    * Host-key policy defaults to **trust-on-first-use** (add the key on first
      contact) so a new host just connects; set ``SSH_AUTO_ADD_HOST_KEYS=false``
      to instead **reject** hosts absent from ``known_hosts``.
    * Connect and per-command timeouts are bounded (``SSH_CONNECT_TIMEOUT`` /
      ``SSH_COMMAND_TIMEOUT`` seconds).
    * Local SSH agent and on-disk private keys are **not** consulted
      (``allow_agent=False``, ``look_for_keys=False``) - only the explicit
      vaulted credential is ever used.
    * A connection/auth failure is returned as a non-zero result (exit 255,
      the SSH convention), never raised, so the run records the failure
      cleanly instead of 500-ing. The secret is never placed in the output.
    """

    def __init__(self):
        # Cache the root-ness probe per (host, username) so a multi-step run
        # only runs `id -u` once against a given host.
        self._is_root_cache: dict[tuple[str, str], bool] = {}

    def _detect_is_root(self, client, host: str, username: str, timeout: float) -> Optional[bool]:
        key = (host, username)
        if key in self._is_root_cache:
            return self._is_root_cache[key]
        try:
            _in, out, _err = client.exec_command("id -u", timeout=timeout)
            uid = out.read().decode("utf-8", "replace").strip()
            is_root = uid == "0"
        except Exception:
            return None  # couldn't tell - caller assumes non-root (safer: use sudo)
        self._is_root_cache[key] = is_root
        return is_root

    def run(self, *, host: str, port: int, username: str, secret: str, command: str) -> CommandResult:
        start = time.monotonic()
        connect_timeout = float(os.getenv("SSH_CONNECT_TIMEOUT", "10"))
        command_timeout = float(os.getenv("SSH_COMMAND_TIMEOUT", "60"))
        # Default: trust-on-first-use so a new host connects; opt back into
        # strict verification with SSH_AUTO_ADD_HOST_KEYS=false.
        auto_add = os.getenv("SSH_AUTO_ADD_HOST_KEYS", "true").strip().lower() != "false"

        client = None
        try:
            # Imported lazily and inside the try so a deployment missing the
            # paramiko dependency records a clean failed run (exit 255) instead
            # of a 500 - constructing the executor never raises.
            import paramiko

            client = paramiko.SSHClient()
            client.load_system_host_keys()
            try:
                client.load_host_keys(os.path.expanduser("~/.ssh/known_hosts"))
            except (FileNotFoundError, OSError):
                pass
            client.set_missing_host_key_policy(paramiko.AutoAddPolicy() if auto_add else paramiko.RejectPolicy())

            client.connect(
                hostname=host,
                port=port,
                username=username,
                password=secret,
                timeout=connect_timeout,
                allow_agent=False,
                look_for_keys=False,
            )

            # Decide privilege elevation per host: sudo only when it's enabled
            # AND the account is not root (unknown root-ness -> assume non-root).
            use_sudo = False
            if sudo_enabled():
                is_root = self._detect_is_root(client, host, username, command_timeout)
                use_sudo = is_root is not True
            effective = build_execution_command(command) if use_sudo else command

            stdin, stdout, stderr = client.exec_command(effective, timeout=command_timeout)
            # For a sudo run, feed the password to `sudo -S` on stdin (never in
            # the command text), then signal EOF so a command that reads stdin
            # can't hang.
            if use_sudo:
                try:
                    stdin.write(sudo_password(secret) + "\n")
                    stdin.flush()
                    stdin.channel.shutdown_write()
                except OSError:
                    pass
            out = stdout.read().decode("utf-8", "replace")
            err = stderr.read().decode("utf-8", "replace")
            exit_status = stdout.channel.recv_exit_status()
            combined = out
            if err.strip():
                combined = (combined + "\n" if combined and not combined.endswith("\n") else combined) + err
            if len(combined) > _MAX_OUTPUT_CHARS:
                combined = combined[:_MAX_OUTPUT_CHARS] + "\n...[output truncated]"
            duration_ms = int((time.monotonic() - start) * 1000)
            return CommandResult(command=effective, output=combined, exit_status=exit_status, duration_ms=duration_ms)
        except Exception as exc:  # connect/auth/timeout/channel errors
            duration_ms = int((time.monotonic() - start) * 1000)
            # Never echo the secret; report the failure type only.
            return CommandResult(
                command=command,
                output=f"[SSH ERROR] {type(exc).__name__}: {exc}",
                exit_status=255,
                duration_ms=duration_ms,
            )
        finally:
            if client is not None:
                try:
                    client.close()
                except Exception:
                    pass


def get_simulated_executor() -> SSHExecutor:
    """Always the simulated executor - powers the **Simulate** button (the
    dry-run rehearsal an operator reviews before a live execution)."""
    return SimulatedSSHExecutor()


def get_live_executor() -> SSHExecutor:
    """The executor for the operator's **Execute Live on Client** action:
    ALWAYS the real paramiko SSH executor, on every deployment. There is no
    'demo live' - the simulated/rehearsal path is the separate Simulate button.
    Live execution is always gated behind an explicit human confirmation AND a
    prior simulation of the same run, so it never runs unattended by accident.
    (Constructing this never raises even if paramiko is missing; the run then
    records a clean failed result - see ParamikoSSHExecutor.run.)"""
    return ParamikoSSHExecutor()