"""SSH client wrapper for the dms-ui. Provides a single class, :class:`DMSSSH`, that executes a shell command on a remote DMS host and returns ``(stdout, stderr, exit_code)``. We deliberately do **not** keep a long-lived connection between requests — every API call opens a fresh connection, runs the command, and closes it. This keeps the layer simple and avoids stuck/half-dead sessions. Two backends are supported, tried in order: 1. ``paramiko`` if it is importable. 2. ``subprocess`` + ``ssh -o BatchMode=yes`` with ``SSH_ASKPASS`` pointed at ``/home/hermes/.ssh-askpass.sh`` — this is the same trick used by ``mcp_server.py`` for the sharekit project. The active backend is recorded in :data:`BACKEND` for diagnostic logging. """ from __future__ import annotations import logging import os import re import shlex import subprocess import threading import time from dataclasses import dataclass from pathlib import Path from typing import List, Optional, Tuple log = logging.getLogger("dms-ui.ssh") ASKPASS_PATH = Path("/home/hermes/.ssh-askpass.sh") # SSH client config: point ssh at /home/hermes/.ssh/config so that # aliases (e.g. "smail-server" → smail.icu) resolve correctly when dms-ui # is launched by the root user (e.g. via systemd or /etc/nginx reverse proxy). SSH_CONFIG_PATH = Path(os.environ.get("DMS_UI_SSH_CONFIG", "/home/hermes/.ssh/config")) ANSI_RE = re.compile(rb"\x1b\[[0-9;?]*[A-Za-z]") BACKEND: str = "unknown" _PARAMIKO_IMPORT_ERROR: Optional[Exception] = None try: import paramiko # type: ignore BACKEND = "paramiko" except Exception as _exc: # pragma: no cover - import-time probe paramiko = None # type: ignore _PARAMIKO_IMPORT_ERROR = _exc BACKEND = "subprocess" @dataclass class CommandResult: """The result of a remote command.""" stdout: str stderr: str exit_code: int backend: str = BACKEND @property def ok(self) -> bool: return self.exit_code == 0 def to_dict(self) -> dict: return { "ok": self.ok, "stdout": self.stdout, "stderr": self.stderr, "exit_code": self.exit_code, "backend": self.backend, } def _strip_ansi(data: bytes) -> str: """Return ``data`` decoded as utf-8 with ANSI escape sequences removed.""" try: text = data.decode("utf-8", errors="replace") except Exception: # pragma: no cover - extremely defensive text = data.decode("latin-1", errors="replace") # Combine byte-level and string-level strip for safety. text = ANSI_RE.sub(b"", data).decode("utf-8", errors="replace") # Also drop stray CSI sequences that snuck in as unicode. text = re.sub(r"\x1b\[[0-9;?]*[A-Za-z]", "", text) return text # A trivial lock so we don't fork-bomb a remote host if 50 requests land at # the same time. The lock is per-process; good enough for a single-tenant UI. _CONCURRENCY = threading.Semaphore(int(os.environ.get("DMS_UI_SSH_CONCURRENCY", "4"))) class DMSSSH: """Run commands on a single DMS host. Parameters ---------- host_alias: The SSH alias as configured in ``~/.ssh/config`` (for example ``smail-server``). password: Optional password override. When ``None`` we fall back to the ``SSH_ASKPASS`` script (``/home/hermes/.ssh-askpass.sh``) by setting ``SSH_ASKPASS_REQUIRE=force``. timeout: Per-command timeout in seconds. Default 60s; log tail streams override this with their own deadline. """ def __init__(self, host_alias: str, password: Optional[str] = None, timeout: int = 60): if not host_alias: raise ValueError("host_alias is required") self.host_alias = host_alias self.password = password self.timeout = timeout # ---- public API ------------------------------------------------------- def run(self, command: str, timeout: Optional[int] = None) -> CommandResult: """Run ``command`` on the remote host and return its output.""" deadline = timeout if timeout is not None else self.timeout with _CONCURRENCY: t0 = time.monotonic() if BACKEND == "paramiko": result = self._run_paramiko(command, deadline) else: result = self._run_subprocess(command, deadline) log.info( "ssh[%s] %r -> rc=%d in %.2fs", self.host_alias, command[:80], result.exit_code, time.monotonic() - t0, ) return result # ---- paramiko backend ------------------------------------------------- def _run_paramiko(self, command: str, timeout: int) -> CommandResult: assert paramiko is not None client = paramiko.SSHClient() client.set_missing_host_key_policy(paramiko.AutoAddPolicy()) # noqa: S507 try: connect_kwargs: dict = { "hostname": self.host_alias, "timeout": min(15, timeout), "allow_agent": True, "look_for_keys": True, } if self.password: connect_kwargs["password"] = self.password client.connect(**connect_kwargs) stdin, stdout, stderr = client.exec_command(command, timeout=timeout) out_b = stdout.read() # blocks until command exits err_b = stderr.read() rc = stdout.channel.recv_exit_status() except Exception as exc: # paramiko raises a grab-bag of exceptions log.exception("paramiko run failed: %s", exc) return CommandResult(stdout="", stderr=str(exc), exit_code=255) finally: try: client.close() except Exception: # pragma: no cover pass return CommandResult( stdout=_strip_ansi(out_b), stderr=_strip_ansi(err_b), exit_code=rc, backend="paramiko", ) # ---- subprocess backend ---------------------------------------------- def _run_subprocess(self, command: str, timeout: int) -> CommandResult: if not ASKPASS_PATH.exists(): return CommandResult( stdout="", stderr=f"askpass script missing: {ASKPASS_PATH}", exit_code=254, ) env = os.environ.copy() env["SSH_ASKPASS"] = str(ASKPASS_PATH) env["SSH_ASKPASS_REQUIRE"] = "force" env["DISPLAY"] = env.get("DISPLAY", ":0") # BatchMode=yes would *suppress* the askpass helper; we want it to be # invoked, so we deliberately omit it. ssh_argv = [ "ssh", "-T", # no pty: avoid interactive prompt side-effects "-F", str(SSH_CONFIG_PATH), # use the hermes ssh config so aliases resolve "-o", "StrictHostKeyChecking=accept-new", "-o", "UserKnownHostsFile=/dev/null", "-o", f"ConnectTimeout={min(15, timeout)}", self.host_alias, command, ] if self.password: env["SSHPASS"] = self.password ssh_argv = ["sshpass", "-e"] + ssh_argv # type: ignore[list-item] try: completed = subprocess.run( # noqa: S603 - inputs controlled ssh_argv, capture_output=True, timeout=timeout, env=env, check=False, ) except subprocess.TimeoutExpired as exc: return CommandResult( stdout=_strip_ansi(exc.stdout or b""), stderr=(exc.stderr or b"").decode("utf-8", errors="replace") + "\n[timeout]", exit_code=124, ) except FileNotFoundError as exc: return CommandResult( stdout="", stderr=f"ssh binary not found: {exc}", exit_code=127, ) except Exception as exc: # pragma: no cover - defensive return CommandResult(stdout="", stderr=str(exc), exit_code=255) return CommandResult( stdout=_strip_ansi(completed.stdout), stderr=_strip_ansi(completed.stderr), exit_code=completed.returncode, backend="subprocess", ) # ---- streaming helpers used by the SSE log endpoint ---------------------- class StreamingSSH: """Open an SSH channel and yield stdout lines until the channel closes. Used by the log tail endpoint. The caller (FastAPI) iterates the lines and forwards them as Server-Sent Events. Cancellation is cooperative: closing the iterator terminates the underlying channel. """ def __init__(self, host_alias: str, command: str, password: Optional[str] = None): self.host_alias = host_alias self.command = command self.password = password self._client: Optional["paramiko.SSHClient"] = None # type: ignore[name-defined] self._channel = None def __iter__(self) -> "StreamingSSH": if BACKEND != "paramiko" or paramiko is None: raise RuntimeError( "streaming tail requires paramiko backend; " f"current backend={BACKEND} import_error={_PARAMIKO_IMPORT_ERROR}" ) client = paramiko.SSHClient() client.set_missing_host_key_policy(paramiko.AutoAddPolicy()) # noqa: S507 kwargs = { "hostname": self.host_alias, "timeout": 15, "allow_agent": True, "look_for_keys": True, } if self.password: kwargs["password"] = self.password client.connect(**kwargs) transport = client.get_transport() assert transport is not None channel = transport.open_session() channel.settimeout(1.0) channel.exec_command(self.command) self._client = client self._channel = channel return self def __next__(self) -> str: assert self._channel is not None # recv_ready + read so we can periodically check for close + cancellation. if self._channel.closed or self._channel.exit_status_ready() and not self._channel.recv_ready(): raise StopIteration try: data = self._channel.recv(4096) except Exception: raise StopIteration if not data: raise StopIteration return _strip_ansi(data) def close(self) -> None: try: if self._channel is not None: self._channel.close() except Exception: # pragma: no cover pass try: if self._client is not None: self._client.close() except Exception: # pragma: no cover pass def run_command(host_alias: str, command: str, password: Optional[str] = None) -> CommandResult: """Functional helper used throughout the FastAPI layer.""" return DMSSSH(host_alias=host_alias, password=password).run(command) def quote_args(args: List[str]) -> str: """POSIX-shell-quote a list of arguments; safe for ``docker exec`` commands.""" return " ".join(shlex.quote(a) for a in args)