"""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 import sys import time 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 run_streaming( self, command: str, *, check: bool = True, timeout: float | None = None ) -> SshResult: """Run ``command`` streaming stdout/stderr live to this process. Unlike :meth:`run` (which blocks on a full channel read), output is forwarded chunk-by-chunk as it arrives so long builds show progress instead of one block at the end. Captured output is still returned in :class:`SshResult`. """ client = self._connect() try: chan = client.get_transport().open_session() # type: ignore[union-attr] if timeout: chan.settimeout(timeout) chan.exec_command(command) except Exception as exc: raise TransportConnectError(f"SSH exec_command failed: {exc}") from exc out_parts: list[str] = [] err_parts: list[str] = [] try: while True: progressed = False while chan.recv_ready(): d = chan.recv(65536).decode("utf-8", errors="replace") out_parts.append(d) sys.stdout.write(d) sys.stdout.flush() progressed = True while chan.recv_stderr_ready(): d = chan.recv_stderr(65536).decode("utf-8", errors="replace") err_parts.append(d) sys.stderr.write(d) sys.stderr.flush() progressed = True if ( chan.exit_status_ready() and not chan.recv_ready() and not chan.recv_stderr_ready() ): break if not progressed: time.sleep(0.05) rc = chan.recv_exit_status() except Exception as exc: raise TransportConnectError(f"SSH stream failed: {exc}") from exc finally: with contextlib.suppress(Exception): chan.close() result = SshResult( stdout="".join(out_parts), stderr="".join(err_parts), 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