local-ci-cd-system/src/ci_orchestrator/transport/ssh.py

"""SSH transport via ``paramiko``.

Replaces the ``ssh.exe`` / ``scp.exe`` calls used in
``scripts/_Transport.psm1``. By default the host-key policy is
``AutoAddPolicy`` with an empty ``known_hosts`` file (i.e. permissive),
which mirrors the ``StrictHostKeyChecking=no UserKnownHostsFile=NUL``
default used for ephemeral CI build VMs.

See ``AGENTS.md`` error #12: in PowerShell with ``$ErrorActionPreference =
'Stop'`` the native ``ssh-keygen -R`` call could raise terminating errors
on stderr output. With paramiko there is no equivalent issue.
"""

from __future__ import annotations

import contextlib
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING

from ci_orchestrator.transport.errors import TransportCommandError, TransportConnectError

if TYPE_CHECKING:  # pragma: no cover
    import paramiko


@dataclass(frozen=True)
class SshResult:
    """Result of a remote SSH command."""

    stdout: str
    stderr: str
    returncode: int

    @property
    def ok(self) -> bool:
        return self.returncode == 0


class SshTransport:
    """Thin wrapper around :class:`paramiko.SSHClient`."""

    def __init__(
        self,
        host: str,
        *,
        username: str = "ci_build",
        key_path: str | Path | None = None,
        port: int = 22,
        connect_timeout: float = 10.0,
        known_hosts: str | Path | None = None,
    ) -> None:
        self._host = host
        self._username = username
        self._key_path = Path(key_path) if key_path else None
        self._port = port
        self._connect_timeout = connect_timeout
        self._known_hosts = Path(known_hosts) if known_hosts else None
        self._client: paramiko.SSHClient | None = None

    # ----------------------------------------------------------- connection
    def _connect(self) -> paramiko.SSHClient:
        if self._client is not None:
            return self._client
        try:
            import paramiko
        except ImportError as exc:  # pragma: no cover - dep declared
            raise TransportConnectError(
                "paramiko is not installed; run `pip install paramiko`."
            ) from exc

        client = paramiko.SSHClient()
        if self._known_hosts is not None and self._known_hosts.is_file():
            client.load_host_keys(str(self._known_hosts))
            client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        else:
            # Permissive mode for ephemeral CI VMs (see _Transport.psm1).
            client.set_missing_host_key_policy(paramiko.AutoAddPolicy())

        try:
            client.connect(
                hostname=self._host,
                port=self._port,
                username=self._username,
                key_filename=str(self._key_path) if self._key_path else None,
                timeout=self._connect_timeout,
                allow_agent=False,
                look_for_keys=False,
            )
        except Exception as exc:
            raise TransportConnectError(
                f"SSH connection to {self._username}@{self._host}:{self._port} failed: {exc}"
            ) from exc
        self._client = client
        return client

    def close(self) -> None:
        if self._client is not None:
            with contextlib.suppress(Exception):
                self._client.close()
            self._client = None

    def __enter__(self) -> SshTransport:
        self._connect()
        return self

    def __exit__(self, *_exc: object) -> None:
        self.close()

    # --------------------------------------------------------------- ops
    def run(self, command: str, *, check: bool = True, timeout: float | None = None) -> SshResult:
        """Run ``command`` in the guest shell."""
        client = self._connect()
        try:
            _stdin, stdout, stderr = client.exec_command(command, timeout=timeout)
            out = stdout.read().decode("utf-8", errors="replace")
            err = stderr.read().decode("utf-8", errors="replace")
            rc = stdout.channel.recv_exit_status()
        except Exception as exc:
            raise TransportConnectError(f"SSH exec_command failed: {exc}") from exc
        result = SshResult(stdout=out, stderr=err, returncode=rc)
        if check and not result.ok:
            raise TransportCommandError(result.returncode, result.stdout, result.stderr)
        return result

    def copy(self, local_path: str | Path, remote_path: str) -> None:
        """Upload via SFTP."""
        client = self._connect()
        try:
            with client.open_sftp() as sftp:
                sftp.put(str(local_path), remote_path)
        except Exception as exc:
            raise TransportConnectError(f"SFTP put failed: {exc}") from exc

    def fetch(self, remote_path: str, local_path: str | Path) -> None:
        """Download via SFTP."""
        client = self._connect()
        try:
            with client.open_sftp() as sftp:
                sftp.get(remote_path, str(local_path))
        except Exception as exc:
            raise TransportConnectError(f"SFTP get failed: {exc}") from exc

    def is_ready(self) -> bool:
        """Lightweight readiness probe."""
        try:
            result = self.run("true", check=False, timeout=5.0)
        except (TransportConnectError, TransportCommandError):
            return False
        return result.ok