from __future__ import annotations import base64 import codecs from collections import deque from dataclasses import dataclass import fcntl import json import os from pathlib import Path import re import shlex import subprocess import tempfile import time from typing import Callable, Iterable, Sequence from urllib.parse import SplitResult, unquote, urlsplit from .domain import is_valid_session_name REGISTRY_VERSION = 1 # Non-interactive, non-replayed SSH: fail fast instead of prompting. SSH_BATCH_OPTIONS = ( "-T", "-o", "BatchMode=yes", "-o", "StrictHostKeyChecking=yes", "-o", "ConnectionAttempts=1", ) def posix_herdr_script(arguments: str, *, prefix: str = "") -> str: """POSIX sh snippet running `herdr ARGUMENTS` with the per-user fallback. A non-interactive SSH PATH often lacks ~/.local/bin; ``arguments`` must already be shell-quoted. Exit 127 lets callers tell "no herdr" apart from a failed command. """ return ( prefix + 'if [ -x "$HOME/.local/bin/herdr" ]; then ' f'exec "$HOME/.local/bin/herdr" {arguments}; ' "elif command -v herdr >/dev/null 2>&1; then " f"exec herdr {arguments}; " "else echo 'herdr is not installed on the remote host' >&2; exit 127; fi" ) REMOTE_DIRECT_COMMAND = "herdr session list --json" REMOTE_LIST_SCRIPT = posix_herdr_script("session list --json") class RegistryError(RuntimeError): """A remote registry error which is safe to show to the user.""" class RemoteTargetError(ValueError): """A remote target cannot be represented safely as an SSH command.""" class RemoteSessionDeleteError(RuntimeError): """A stopped remote Session could not be deleted.""" class RemoteSessionStopError(RuntimeError): """A running remote Session could not be stopped.""" @dataclass(frozen=True) class RemoteEntry: target: str remote_keybindings: str = "local" remote_transport: str = "native" def __post_init__(self) -> None: validate_target(self.target) if self.remote_keybindings not in {"local", "server"}: raise ValueError("remote keybindings must be 'local' or 'server'") if self.remote_transport not in {"native", "ssh"}: raise ValueError("remote transport must be 'native' or 'ssh'") @dataclass(frozen=True) class RemoteSession: name: str running: bool @dataclass(frozen=True) class RemoteQueryResult: entry: RemoteEntry sessions: tuple[RemoteSession, ...] = () error: str | None = None platform: str = "unknown" def validate_target(target: str) -> None: if not target: raise RemoteTargetError("remote target cannot be empty") if target.startswith("-"): raise RemoteTargetError("remote target must not start with '-'") if any(ord(character) < 32 or ord(character) == 127 for character in target): raise RemoteTargetError("remote target must not contain control characters") if target.startswith("ssh://"): _parse_ssh_url(target) def _parse_ssh_url(target: str) -> SplitResult: parsed = urlsplit(target) if parsed.scheme != "ssh" or parsed.hostname is None: raise RemoteTargetError("invalid ssh:// remote target") if parsed.password is not None: raise RemoteTargetError("ssh:// targets must not contain passwords") if parsed.path not in {"", "/"} or parsed.query or parsed.fragment: raise RemoteTargetError("ssh:// targets must not contain a path, query or fragment") try: parsed.port except ValueError as error: raise RemoteTargetError(f"ssh:// target has an invalid port: {error}") from error host = parsed.hostname if host.startswith("-") or any(character.isspace() for character in host): raise RemoteTargetError("ssh:// target contains an invalid host") if parsed.username is not None: username = unquote(parsed.username) if ( not username or username.startswith("-") or any( character.isspace() or ord(character) < 32 or ord(character) == 127 for character in username ) ): raise RemoteTargetError("ssh:// target contains an invalid username") return parsed def default_registry_path() -> Path: if override := os.environ.get("HERDR_HUB_REMOTES_FILE"): return Path(override).expanduser() config_home = os.environ.get("XDG_CONFIG_HOME") root = Path(config_home).expanduser() if config_home else Path.home() / ".config" return root / "herdr-hub" / "remotes.json" class RemoteRegistry: def __init__(self, path: Path | None = None) -> None: self.path = path or default_registry_path() @property def lock_path(self) -> Path: return self.path.with_name(f".{self.path.name}.lock") def load(self) -> list[RemoteEntry]: if not self.path.exists(): return [] try: payload = json.loads(self.path.read_text(encoding="utf-8")) except (OSError, json.JSONDecodeError) as error: raise RegistryError(f"cannot read {self.path}: {error}") from error return _parse_registry(payload, self.path) def upsert(self, entry: RemoteEntry) -> bool: def update(entries: list[RemoteEntry]) -> tuple[list[RemoteEntry], bool]: by_target = {item.target: item for item in entries} changed = by_target.get(entry.target) != entry by_target[entry.target] = entry return sorted(by_target.values(), key=lambda item: item.target), changed return self._update(update) def upsert_many(self, entries: Iterable[RemoteEntry]) -> int: incoming = list(entries) if not incoming: return 0 def update(current: list[RemoteEntry]) -> tuple[list[RemoteEntry], int]: by_target = {item.target: item for item in current} changed = 0 for entry in incoming: if by_target.get(entry.target) != entry: changed += 1 by_target[entry.target] = entry return sorted(by_target.values(), key=lambda item: item.target), changed return self._update(update) def remember_many(self, entries: Iterable[RemoteEntry]) -> int: """Add discovered targets without overwriting explicit preferences. SSH transport is sticky so legacy Windows windows can teach the registry that a target must not be reopened through native ``--remote`` mode. """ incoming = list(entries) if not incoming: return 0 def update(current: list[RemoteEntry]) -> tuple[list[RemoteEntry], int]: by_target = {item.target: item for item in current} changed = 0 for entry in incoming: existing = by_target.get(entry.target) if existing is None: by_target[entry.target] = entry changed += 1 elif ( existing.remote_transport != "ssh" and entry.remote_transport == "ssh" ): by_target[entry.target] = RemoteEntry( existing.target, existing.remote_keybindings, "ssh", ) changed += 1 return sorted(by_target.values(), key=lambda item: item.target), changed return self._update(update) def remove(self, target: str) -> bool: validate_target(target) def update(entries: list[RemoteEntry]) -> tuple[list[RemoteEntry], bool]: remaining = [entry for entry in entries if entry.target != target] return remaining, len(remaining) != len(entries) return self._update(update) def _update(self, callback): # type: ignore[no-untyped-def] try: self.path.parent.mkdir(mode=0o700, parents=True, exist_ok=True) with self.lock_path.open("a+", encoding="utf-8") as lock: fcntl.flock(lock.fileno(), fcntl.LOCK_EX) entries = self.load() updated, outcome = callback(entries) if updated != entries: self._write(updated) return outcome except RegistryError: raise except OSError as error: raise RegistryError(f"cannot update {self.path}: {error}") from error def _write(self, entries: Sequence[RemoteEntry]) -> None: payload = { "version": REGISTRY_VERSION, "remotes": [ { "target": entry.target, "remote_keybindings": entry.remote_keybindings, "remote_transport": entry.remote_transport, } for entry in entries ], } temporary_name: str | None = None try: with tempfile.NamedTemporaryFile( mode="w", encoding="utf-8", dir=self.path.parent, prefix=f".{self.path.name}.", delete=False, ) as temporary: temporary_name = temporary.name json.dump(payload, temporary, ensure_ascii=False, indent=2) temporary.write("\n") temporary.flush() os.fsync(temporary.fileno()) os.chmod(temporary_name, 0o600) os.replace(temporary_name, self.path) directory_fd = os.open(self.path.parent, os.O_RDONLY | os.O_DIRECTORY) try: os.fsync(directory_fd) finally: os.close(directory_fd) finally: if temporary_name is not None: try: Path(temporary_name).unlink(missing_ok=True) except OSError: pass def _parse_registry(payload: object, path: Path) -> list[RemoteEntry]: if not isinstance(payload, dict) or payload.get("version") != REGISTRY_VERSION: raise RegistryError(f"{path} has an unsupported or missing version") raw_entries = payload.get("remotes") if not isinstance(raw_entries, list): raise RegistryError(f"{path} does not contain a remotes list") entries: list[RemoteEntry] = [] seen: set[str] = set() for index, raw in enumerate(raw_entries): if not isinstance(raw, dict): raise RegistryError(f"{path}: remotes[{index}] is not an object") target = raw.get("target") keybindings = raw.get("remote_keybindings", "local") transport = raw.get("remote_transport", "native") if ( not isinstance(target, str) or not isinstance(keybindings, str) or not isinstance(transport, str) ): raise RegistryError(f"{path}: remotes[{index}] has invalid fields") try: entry = RemoteEntry(target, keybindings, transport) except ValueError as error: raise RegistryError(f"{path}: remotes[{index}]: {error}") from error if entry.target in seen: raise RegistryError(f"{path}: duplicate remote target {entry.target!r}") seen.add(entry.target) entries.append(entry) return sorted(entries, key=lambda entry: entry.target) def ssh_destination_arguments(target: str) -> list[str]: validate_target(target) if not target.startswith("ssh://"): return [target] parsed = _parse_ssh_url(target) assert parsed.hostname is not None host = parsed.hostname arguments: list[str] = [] if parsed.port is not None: arguments.extend(["-p", str(parsed.port)]) if parsed.username is not None: username = unquote(parsed.username) arguments.extend(["-l", username]) arguments.append(host) return arguments def ssh_command( ssh_bin: str, entry: RemoteEntry, remote_command: str, *, connect_timeout: int = 4, ) -> list[str]: """One batch-mode SSH invocation running ``remote_command`` on the target.""" return [ ssh_bin, *SSH_BATCH_OPTIONS, "-o", f"ConnectTimeout={connect_timeout}", *ssh_destination_arguments(entry.target), remote_command, ] def remote_herdr_command( ssh_bin: str, entry: RemoteEntry, arguments: str, *, connect_timeout: int = 4, posix_fallback: bool = False, ) -> list[str]: """Run `herdr ARGUMENTS` remotely, via PATH or the POSIX ~/.local/bin fallback.""" remote_command = ( posix_herdr_script(arguments) if posix_fallback else f"herdr {arguments}" ) return ssh_command(ssh_bin, entry, remote_command, connect_timeout=connect_timeout) def remote_session_command( ssh_bin: str, entry: RemoteEntry, *, connect_timeout: int = 4, posix_fallback: bool = False, ) -> list[str]: return remote_herdr_command( ssh_bin, entry, "session list --json", connect_timeout=connect_timeout, posix_fallback=posix_fallback, ) def _session_arguments(action: str, session: str) -> str: if not is_valid_session_name(session): raise ValueError("invalid remote Session name") return f"session {action} -- {session}" def remote_delete_session_command( ssh_bin: str, entry: RemoteEntry, session: str, *, connect_timeout: int = 4, posix_fallback: bool = False, ) -> list[str]: return remote_herdr_command( ssh_bin, entry, _session_arguments("delete --json", session), connect_timeout=connect_timeout, posix_fallback=posix_fallback, ) def remote_stop_session_command( ssh_bin: str, entry: RemoteEntry, session: str, *, connect_timeout: int = 4, posix_fallback: bool = False, ) -> list[str]: return remote_herdr_command( ssh_bin, entry, _session_arguments("stop", session), connect_timeout=connect_timeout, posix_fallback=posix_fallback, ) def _run_remote_session_command( build: Callable[..., list[str]], error_type: type[RuntimeError], action: str, ssh_bin: str, entry: RemoteEntry, session: str, timeout: float, ) -> tuple[str, str]: """Run a Session mutation, retrying once through the POSIX PATH fallback. Returns decoded (stdout, stderr); any non-zero exit becomes ``error_type``. """ result: subprocess.CompletedProcess[bytes] | None = None for posix_fallback in (False, True): command = build(ssh_bin, entry, session, posix_fallback=posix_fallback) try: result = subprocess.run( command, capture_output=True, check=False, timeout=timeout, ) except subprocess.TimeoutExpired as error: raise error_type( f"remote Session {action} timed out after {timeout:g}s" ) from error except OSError as error: raise error_type(f"could not start remote Session {action}: {error}") from error if result.returncode != 127 or posix_fallback: break assert result is not None stdout = decode_remote_output(result.stdout) stderr = decode_remote_output(result.stderr) if result.returncode != 0: detail = (stderr or stdout or f"ssh exited with status {result.returncode}").strip() raise error_type(" ".join(detail.splitlines())) return stdout, stderr def delete_remote_session( ssh_bin: str, entry: RemoteEntry, session: str, *, timeout: float = 8.0, ) -> None: """Delete one stopped Herdr Session through a non-interactive SSH call.""" stdout, _ = _run_remote_session_command( remote_delete_session_command, RemoteSessionDeleteError, "delete", ssh_bin, entry, session, timeout, ) try: payload = json.loads(stdout) except json.JSONDecodeError as error: raise RemoteSessionDeleteError( f"invalid Herdr delete response: {error}" ) from error if not isinstance(payload, dict) or payload.get("deleted") is not True: raise RemoteSessionDeleteError("Herdr did not confirm Session deletion") def stop_remote_session( ssh_bin: str, entry: RemoteEntry, session: str, *, timeout: float = 8.0, ) -> None: """Stop one running Herdr Session through a non-interactive SSH call.""" _run_remote_session_command( remote_stop_session_command, RemoteSessionStopError, "stop", ssh_bin, entry, session, timeout, ) def clip_utf8_tail(raw: bytes, limit: int) -> bytes: """Keep the last ``limit`` bytes without starting inside a UTF-8 sequence.""" if len(raw) <= limit: return raw clipped = raw[-limit:] skip = 0 while skip < min(3, len(clipped)) and 0x80 <= clipped[skip] <= 0xBF: skip += 1 return clipped[skip:] def decode_remote_output(raw: bytes, *, partial: bool = False) -> str: """Decode UTF-8 or Windows code page 936 output. ``partial`` marks a prefix of a stream that is still being written. An incomplete trailing multi-byte sequence is then held back instead of being treated as corruption, which would otherwise push the whole buffer through the GB18030 fallback and garble every character of a progress display. """ if not raw: return "" if partial: for encoding in ("utf-8", "gb18030"): try: return codecs.getincrementaldecoder(encoding)().decode(raw, final=False) except UnicodeDecodeError: continue try: return raw.decode("utf-8") except UnicodeDecodeError: try: return raw.decode("gb18030") except UnicodeDecodeError: return raw.decode("utf-8", errors="replace") def parse_remote_session_result( entry: RemoteEntry, returncode: int, stdout: str, stderr: str, ) -> RemoteQueryResult: if returncode != 0: detail = (stderr or stdout or f"ssh exited with status {returncode}").strip() detail = " ".join(detail.splitlines()) if len(detail) > 180: detail = detail[:179] + "…" return RemoteQueryResult(entry, error=detail) try: payload = json.loads(stdout) raw_sessions = payload["sessions"] if not isinstance(raw_sessions, list): raise TypeError("sessions is not a list") except (json.JSONDecodeError, KeyError, TypeError) as error: return RemoteQueryResult(entry, error=f"invalid Herdr session JSON: {error}") sessions: dict[str, RemoteSession] = {} for raw in raw_sessions: if not isinstance(raw, dict): continue name = raw.get("name") running = raw.get("running") if isinstance(name, str) and isinstance(running, bool): sessions.setdefault(name, RemoteSession(name, running)) ordered = sorted( sessions.values(), key=lambda session: (session.name != "default", session.name), ) return RemoteQueryResult( entry, sessions=tuple(ordered), platform=detect_remote_platform(raw_sessions), ) def detect_remote_platform(raw_sessions: list[object]) -> str: saw_posix_path = False for raw in raw_sessions: if not isinstance(raw, dict): continue for field in ("session_dir", "socket_path"): value = raw.get(field) if not isinstance(value, str): continue if ( len(value) >= 3 and value[0].isalpha() and value[1] == ":" and value[2] in {"\\", "/"} ): return "windows" if value.startswith("/"): saw_posix_path = True return "posix" if saw_posix_path else "unknown" @dataclass class _RemoteJob: entry: RemoteEntry process: subprocess.Popen[bytes] started_at: float posix_fallback: bool = False @dataclass class _StoppingProcess: process: subprocess.Popen[bytes] since: float killed: bool = False class RemoteQueryPool: """Pollable bounded pool of SSH processes for an interactive selector.""" def __init__( self, ssh_bin: str, *, max_parallel: int = 6, query_timeout: float = 6.0, connect_timeout: int = 4, ) -> None: self.ssh_bin = ssh_bin self.max_parallel = max_parallel self.query_timeout = query_timeout self.connect_timeout = connect_timeout self._pending: deque[RemoteEntry] = deque() self._active: dict[str, _RemoteJob] = {} self._completed: deque[RemoteQueryResult] = deque() self._stopping: list[_StoppingProcess] = [] def enqueue(self, entry: RemoteEntry) -> None: self.cancel(entry.target) self._pending.append(entry) self._start_pending() def enqueue_many(self, entries: Iterable[RemoteEntry]) -> None: for entry in entries: self.cancel(entry.target) self._pending.append(entry) self._start_pending() def cancel(self, target: str) -> None: self._pending = deque(entry for entry in self._pending if entry.target != target) job = self._active.pop(target, None) if job is not None: self._stop_process(job.process) self._start_pending() def is_loading(self, target: str) -> bool: return target in self._active or any( entry.target == target for entry in self._pending ) def poll(self) -> list[RemoteQueryResult]: self._reap_stopping() finished = list(self._completed) self._completed.clear() now = time.monotonic() for target, job in list(self._active.items()): returncode = job.process.poll() timed_out = now - job.started_at >= self.query_timeout if returncode is None and not timed_out: continue if timed_out and returncode is None: job.process.kill() try: stdout_bytes, stderr_bytes = job.process.communicate() except (OSError, subprocess.SubprocessError) as error: self._active.pop(target, None) finished.append( RemoteQueryResult( job.entry, error=f"remote query output could not be read: {error}", ) ) continue self._active.pop(target, None) if timed_out: finished.append( RemoteQueryResult( job.entry, error=f"remote query timed out after {self.query_timeout:g}s", ) ) elif job.process.returncode == 127 and not job.posix_fallback: fallback = self._start_job(job.entry, posix_fallback=True) if fallback is not None: self._active[target] = fallback else: finished.append( parse_remote_session_result( job.entry, job.process.returncode, decode_remote_output(stdout_bytes), decode_remote_output(stderr_bytes), ) ) self._start_pending() return finished def close(self) -> None: self._pending.clear() for job in list(self._active.values()): self._stop_process(job.process) self._active.clear() # The popup is exiting: wait briefly so no zombie or open pipe # outlives it. This is the only place a stop is allowed to block. for stopping in self._stopping: try: stopping.process.communicate(timeout=0.25) except subprocess.TimeoutExpired: stopping.process.kill() stopping.process.communicate() self._stopping.clear() def _start_pending(self) -> None: while self._pending and len(self._active) < self.max_parallel: entry = self._pending.popleft() job = self._start_job(entry) if job is not None: self._active[entry.target] = job def _start_job( self, entry: RemoteEntry, *, posix_fallback: bool = False, ) -> _RemoteJob | None: try: process = subprocess.Popen( remote_session_command( self.ssh_bin, entry, connect_timeout=self.connect_timeout, posix_fallback=posix_fallback, ), stdin=subprocess.DEVNULL, stdout=subprocess.PIPE, stderr=subprocess.PIPE, ) except (OSError, ValueError) as error: self._completed.append( RemoteQueryResult(entry, error=f"cannot start remote query: {error}") ) return None return _RemoteJob(entry, process, time.monotonic(), posix_fallback) def _stop_process(self, process: subprocess.Popen[bytes]) -> None: """Ask a query to stop without blocking the UI thread. Cancelling used to wait up to 250 ms per process; a refresh of six hosts could freeze the popup for over a second. The process is reaped on later polls and escalated to SIGKILL if it ignores SIGTERM. """ if process.poll() is None: process.terminate() self._stopping.append(_StoppingProcess(process, time.monotonic())) def _reap_stopping(self) -> None: now = time.monotonic() remaining: list[_StoppingProcess] = [] for stopping in self._stopping: if stopping.process.poll() is not None: for stream in (stopping.process.stdout, stopping.process.stderr): if stream is not None: stream.close() continue if not stopping.killed and now - stopping.since >= 0.25: stopping.process.kill() stopping.killed = True remaining.append(stopping) self._stopping = remaining def remote_control_command( ssh_bin: str, entry: RemoteEntry, arguments: Sequence[str], *, platform: str, unset_environment: Sequence[str] = (), load_proxy_environment: bool = False, ) -> list[str]: """Build one non-replayed mutation with platform-specific argv encoding. Windows SSH commonly uses cmd.exe; only a base64 PowerShell script crosses that shell. ProcessStartInfo then receives Windows-quoted native arguments. """ if any(not re.fullmatch(r"[A-Z_][A-Z_0-9]*", key) for key in unset_environment): raise ValueError("invalid environment variable name") if any("\0" in argument for argument in arguments): raise ValueError("NUL in remote argument") if platform == "posix": unset = "".join(f"unset {key}; " for key in unset_environment) script = posix_herdr_script(shlex.join(arguments), prefix=unset) elif platform == "windows": def literal(value: str) -> str: return "'" + value.replace("'", "''") + "'" unset = "".join(f"Remove-Item Env:{key} -ErrorAction SilentlyContinue; " for key in unset_environment) proxy_setup = "" if load_proxy_environment: proxy_setup = ( "foreach ($key in @('http_proxy','https_proxy','all_proxy','no_proxy')) { " "$value=[Environment]::GetEnvironmentVariable($key,'Process'); " "if ([string]::IsNullOrEmpty($value)) { " "$value=[Environment]::GetEnvironmentVariable($key,'User'); " "if ([string]::IsNullOrEmpty($value)) { $value=[Environment]::GetEnvironmentVariable($key,'Machine') }; " "if (![string]::IsNullOrEmpty($value)) { [Environment]::SetEnvironmentVariable($key,$value,'Process') } } }; " "if ($env:https_proxy -or $env:http_proxy -or $env:all_proxy) { " "Write-Output 'Using remote proxy environment (process/user/machine).' " "} else { Write-Output 'No remote proxy environment found; using remote network defaults.' }; " ) script_text = ( "$ErrorActionPreference='Stop'; " + unset + proxy_setup + "$exe=Join-Path $HOME '.local\\bin\\herdr.exe'; " "if (!(Test-Path -LiteralPath $exe)) { $exe=(Get-Command herdr -ErrorAction Stop).Source }; " "$p=New-Object System.Diagnostics.Process; " "$p.StartInfo.FileName=$exe; " "$p.StartInfo.UseShellExecute=$false; " "$p.StartInfo.Arguments=" + literal(subprocess.list2cmdline(arguments)) + "; " "$null=$p.Start(); $p.WaitForExit(); exit $p.ExitCode" ) encoded = base64.b64encode(script_text.encode("utf-16le")).decode("ascii") script = "powershell.exe -NoLogo -NoProfile -NonInteractive -EncodedCommand " + encoded else: raise ValueError("unknown remote platform") return ssh_command(ssh_bin, entry, script)