Source code for remote_compression.ssh

"""SSH transport: one reusable session per batch, resolved through ``~/.ssh/config``."""

import ipaddress
import logging
import socket
import time
from collections import deque
from dataclasses import dataclass
from pathlib import Path

import paramiko

logger = logging.getLogger("rcomp")

#: Bytes of output kept per channel when reporting a command result.
TAIL_BYTES = 4096


[docs] class SSHError(Exception): """Base class for errors of the ssh layer."""
[docs] class SSHConnectionError(SSHError): """Host unreachable, authentication failure, host key mismatch, or connection lost for good."""
[docs] class RemoteExecError(SSHError): """A remote command stopped producing output (inactivity timeout)."""
[docs] @dataclass(frozen=True, slots=True) class ExecResult: """ Outcome of a remote command. Attributes ---------- exit_status: :class:`int` Remote exit status; ``-1`` when the server never sent one (treated as failure). stdout_tail: :class:`str` Last :data:`TAIL_BYTES` bytes of stdout, decoded with replacement. stderr_tail: :class:`str` Last :data:`TAIL_BYTES` bytes of stderr, decoded with replacement. """ exit_status: int stdout_tail: str = "" stderr_tail: str = "" @property def ok(self): """:class:`bool`: True iff the exit status is 0.""" return self.exit_status == 0
[docs] def get_config(hostname, config_path=None): """ Parameters ---------- hostname: :class:`str` Destination alias. config_path: :class:`str`, optional Location of the openSSH config file if different from ``~/.ssh/config``. Returns ------- :class:`dict` Parameters for the destination (``hostname``, ``user``, ``port``, ``identityfile``, ``proxyjump``...). """ if config_path is None: config_path = str(Path.home() / ".ssh" / "config") try: config = paramiko.SSHConfig.from_path(config_path) except FileNotFoundError: config = paramiko.SSHConfig() return config.lookup(hostname)
[docs] def split_user_host(destination): """ Split an ``[user@]alias`` destination. Parameters ---------- destination: :class:`str` e.g. ``nas`` or ``admin@nas``. Returns ------- :class:`str` or None Explicit user, when given (it beats the ssh config ``User``). :class:`str` The alias to look up in the ssh config. """ user, _, alias = destination.rpartition("@") return (user or None), alias
def _local_ip_set(): """IP addresses of this machine (best effort, never raises).""" ips = set() try: for info in socket.getaddrinfo(socket.gethostname(), None): ips.add(str(info[4][0]).split("%")[0]) except OSError: pass # UDP-connect trick: no packet is sent, but the OS picks the outbound address. for family, probe_addr in ( (socket.AF_INET, ("8.8.8.8", 80)), (socket.AF_INET6, ("2001:4860:4860::8888", 80)), ): try: with socket.socket(family, socket.SOCK_DGRAM) as s: s.connect(probe_addr) ips.add(s.getsockname()[0].split("%")[0]) except OSError: pass return ips
[docs] def is_localhost(hostname, config_path=None): """ Tell whether an SSH destination is in fact this very machine. Parameters ---------- hostname: :class:`str` Destination alias (``local`` is the explicit sentinel). config_path: :class:`str`, optional Alternative openSSH config file (mostly for tests). Returns ------- :class:`bool` True when the effective HostName designates this machine. Guards: a ProxyJump or a non-standard port always means "remote" (a tunnel like ``HostName localhost / Port 2222`` usually leads to another machine). Notes ----- Known false negative, by design: a DNS name resolving to the public IP of a NAT that loops back to this machine (hairpin) is treated as remote — the public address matches no local interface. Everything still works, only through a pointless SSH round-trip. When it matters, use ``-D local`` or an alias whose HostName is the LAN name of the machine. Detecting this reliably would require comparing SSH host keys after connecting, which is not worth the machinery. """ if hostname.lower() == "local": return True _, alias = split_user_host(hostname) cfg = get_config(alias, config_path) if cfg.get("proxyjump"): return False if int(cfg.get("port", 22)) != 22: return False target = str(cfg.get("hostname", alias)).lower().rstrip(".") if target in {"localhost", "127.0.0.1", "::1"}: return True if target in {socket.gethostname().lower(), socket.getfqdn().lower()}: return True try: infos = socket.getaddrinfo(target, None) except OSError: return False target_ips = {str(info[4][0]).split("%")[0] for info in infos} try: if any(ipaddress.ip_address(ip).is_loopback for ip in target_ips): return True except ValueError: pass return bool(target_ips & _local_ip_set())
#: Permanent failures: retrying cannot help. Order matters — these are #: subclasses of the retryable ones and must be caught first. _FATAL_EXC = ( paramiko.ssh_exception.AuthenticationException, paramiko.ssh_exception.BadHostKeyException, socket.gaierror, ) _RETRYABLE_EXC = (paramiko.ssh_exception.SSHException, ConnectionError, TimeoutError, OSError, EOFError) def _connect_with_retry(client, *, attempts=3, base_delay=2.0, max_delay=30.0, **connect_kwargs): """ Call ``client.connect(**connect_kwargs)`` with bounded retries. Retries transient errors with exponential backoff; authentication, host-key and DNS failures raise immediately. Raises ------ SSHConnectionError On permanent failure or when all attempts are exhausted. """ host = connect_kwargs.get("hostname") delay = base_delay last = None for attempt in range(attempts): try: client.connect(**connect_kwargs) return except _FATAL_EXC as e: msg = f"Cannot connect to {host}: {e}" raise SSHConnectionError(msg) from e except _RETRYABLE_EXC as e: last = e if attempt + 1 < attempts: logger.warning("Connection to %s failed (%s), retrying in %.0f s.", host, e, delay) time.sleep(delay) delay = min(delay * 2, max_delay) msg = f"Cannot connect to {host} after {attempts} attempts: {last}" raise SSHConnectionError(msg) from last def _parse_jump(jump, default_user): """Parse a ProxyJump value ``[user@]host[:port]`` -> (user, host, port or None).""" if "," in jump: msg = "Multi-hop ProxyJump is not supported." raise SSHConnectionError(msg) user, _, host = jump.rpartition("@") port = None if ":" in host and not host.startswith("["): host, _, port_text = host.rpartition(":") port = int(port_text) return user or default_user, host, port def _tail(chunks): return b"".join(chunks)[-TAIL_BYTES:].decode("utf-8", errors="replace")
[docs] class SSHSession: """ A reusable SSH connection (plus SFTP) to one destination. Composition over inheritance: the session owns up to two :class:`paramiko.SSHClient` (an optional ProxyJump gateway and the destination) and a lazy SFTP channel. Use as a context manager; it returns itself. Parameters ---------- hostname: :class:`str` Destination alias; extra parameters (HostName, User, Port, IdentityFile, ProxyJump) come from the openSSH config file. config_path: :class:`str`, optional Alternative openSSH config file. attempts: :class:`int` Connection attempts before giving up. connect_timeout: :class:`float` Timeout (seconds) applied to TCP, banner and auth phases. keepalive: :class:`int` Keepalive interval (seconds), applied to both transports. """ def __init__(self, hostname, config_path=None, attempts=3, connect_timeout=20.0, keepalive=30): self.hostname = hostname self.config_path = config_path self.attempts = attempts self.connect_timeout = connect_timeout self.keepalive = keepalive self._gw = None self._ssh = None self._sftp = None # -- lifecycle ----------------------------------------------------------
[docs] def open(self): """Connect (idempotent). Returns the session itself.""" if self.is_alive(): return self self.close() explicit_user, alias = split_user_host(self.hostname) cfg = get_config(alias, self.config_path) dest = cfg.get("hostname", alias) user = explicit_user or cfg.get("user") port = int(cfg.get("port", 22)) timeouts = { "timeout": self.connect_timeout, "banner_timeout": self.connect_timeout, "auth_timeout": self.connect_timeout, } sock = None if cfg.get("proxyjump"): gw_user, gw_alias, gw_port = _parse_jump(cfg["proxyjump"], default_user=None) gw_cfg = get_config(gw_alias, self.config_path) self._gw = paramiko.SSHClient() self._gw.load_system_host_keys() self._gw.set_missing_host_key_policy(paramiko.AutoAddPolicy()) _connect_with_retry( self._gw, attempts=self.attempts, hostname=gw_cfg.get("hostname", gw_alias), port=gw_port if gw_port is not None else int(gw_cfg.get("port", 22)), # explicit user@ in ProxyJump beats the jump host's own config username=gw_user or gw_cfg.get("user"), key_filename=gw_cfg.get("identityfile"), **timeouts, ) transport = self._gw.get_transport() transport.set_keepalive(self.keepalive) try: sock = transport.open_channel("direct-tcpip", (dest, port), ("127.0.0.1", 0)) except paramiko.ssh_exception.SSHException as e: self.close() msg = f"ProxyJump channel to {dest}:{port} failed: {e}" raise SSHConnectionError(msg) from e self._ssh = paramiko.SSHClient() self._ssh.load_system_host_keys() self._ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy()) try: _connect_with_retry( self._ssh, attempts=self.attempts, hostname=dest, port=port, username=user, key_filename=cfg.get("identityfile"), sock=sock, **timeouts, ) except SSHConnectionError: self.close() raise self._ssh.get_transport().set_keepalive(self.keepalive) logger.debug("Connected to %s.", self.hostname) return self
[docs] def close(self): """Close everything (sftp, destination, gateway), never raising.""" for client in (self._sftp, self._ssh, self._gw): if client is not None: try: client.close() except Exception: # noqa: BLE001, S110 - closing is best-effort by design pass self._sftp = self._ssh = self._gw = None
def __enter__(self): return self.open() def __exit__(self, exc_type, exc_val, exc_tb): self.close() # -- state --------------------------------------------------------------
[docs] def is_alive(self): """:class:`bool`: True iff the destination transport is up.""" if self._ssh is None: return False transport = self._ssh.get_transport() return bool(transport and transport.is_active())
[docs] def ensure_connected(self): """Reconnect if the transport died.""" if not self.is_alive(): self.reconnect()
[docs] def reconnect(self): """Tear everything down and connect again.""" logger.info("Reconnecting to %s.", self.hostname) self.close() self.open()
@property def sftp(self): """ :class:`paramiko.SFTPClient`: lazy SFTP channel, reopened when dead. Raises ------ SSHError When the server refuses the SFTP subsystem (e.g. Synology DSM ships with the SFTP service disabled) — with the remedy spelled out. SSHConnectionError When the connection died. """ if self._ssh is None: msg = "Session is not open." raise SSHError(msg) if self._sftp is None or (self._sftp.sock is not None and self._sftp.sock.closed): try: self._sftp = self._ssh.open_sftp() except (paramiko.ssh_exception.SSHException, EOFError, OSError) as e: if not self.is_alive(): msg = f"Connection to {self.hostname} lost while opening SFTP: {e}" raise SSHConnectionError(msg) from e msg = ( f"The server {self.hostname} refuses SFTP ({e}). rcomp does all its file " "transfers over SFTP: enable the SFTP service on the server (on Synology " "DSM: Control Panel > File Services > FTP > SFTP)." ) raise SSHError(msg) from e return self._sftp # -- execution ----------------------------------------------------------
[docs] def run(self, command, inactivity_timeout=300.0, on_stdout=None): """ Execute a remote command, draining its output continuously. Continuous draining is not cosmetic: paramiko's flow-control window (~2 MB) fills up otherwise and the remote process blocks on write. The timeout is an *inactivity* timeout, not a total duration: any output (e.g. the ffmpeg ``-progress`` stream) resets it. Parameters ---------- command: :class:`str` Command line (ASCII; syntax must suit both sh and cmd.exe). inactivity_timeout: :class:`float` Seconds without any output before giving up. on_stdout: callable, optional Called with each decoded stdout chunk (progress feed). Returns ------- :class:`ExecResult` Exit status and bounded output tails. Raises ------ RemoteExecError If the command stops producing output for `inactivity_timeout` seconds. SSHConnectionError If the connection drops during execution. """ if not self.is_alive(): msg = f"Not connected to {self.hostname}." raise SSHConnectionError(msg) transport = self._ssh.get_transport() try: chan = transport.open_session() chan.exec_command(command) chan.shutdown_write() except (paramiko.ssh_exception.SSHException, EOFError, OSError) as e: msg = f"Connection to {self.hostname} lost while starting a command: {e}" raise SSHConnectionError(msg) from e out_chunks, err_chunks = deque(maxlen=64), deque(maxlen=64) last_activity = time.monotonic() try: while True: drained = False while chan.recv_ready(): chunk = chan.recv(32768) if chunk: out_chunks.append(chunk) drained = True if on_stdout is not None: on_stdout(chunk.decode("utf-8", errors="replace")) while chan.recv_stderr_ready(): chunk = chan.recv_stderr(32768) if chunk: err_chunks.append(chunk) drained = True if chan.exit_status_ready() and not chan.recv_ready() and not chan.recv_stderr_ready(): break if not transport.is_active(): msg = f"Connection to {self.hostname} lost during command execution." raise SSHConnectionError(msg) if drained: last_activity = time.monotonic() elif time.monotonic() - last_activity > inactivity_timeout: msg = f"No output for {inactivity_timeout:.0f} s: {command[:80]}" raise RemoteExecError(msg) else: time.sleep(0.05) exit_status = chan.recv_exit_status() finally: try: chan.close() except Exception: # noqa: BLE001, S110 - closing is best-effort by design pass return ExecResult(exit_status, _tail(out_chunks), _tail(err_chunks))