from __future__ import annotations from collections import deque from dataclasses import dataclass import fcntl import json import os from pathlib import Path import subprocess import tempfile import time from typing import Iterable, Sequence from urllib.parse import SplitResult, unquote, urlsplit from .domain import is_valid_session_name REGISTRY_VERSION = 1 REMOTE_DIRECT_COMMAND = "herdr session list --json" REMOTE_LIST_SCRIPT = ( 'if [ -x "$HOME/.local/bin/herdr" ]; then ' 'exec "$HOME/.local/bin/herdr" session list --json; ' "elif command -v herdr >/dev/null 2>&1; then " "exec herdr session list --json; " "else echo 'herdr is not installed on the remote host' >&2; exit 127; fi" ) 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 remote_session_command( ssh_bin: str, entry: RemoteEntry, *, connect_timeout: int = 4, posix_fallback: bool = False, ) -> list[str]: return [ ssh_bin, "-T", "-o", "BatchMode=yes", "-o", "StrictHostKeyChecking=yes", "-o", "ConnectionAttempts=1", "-o", f"ConnectTimeout={connect_timeout}", *ssh_destination_arguments(entry.target), REMOTE_LIST_SCRIPT if posix_fallback else REMOTE_DIRECT_COMMAND, ] def remote_delete_session_command( ssh_bin: str, entry: RemoteEntry, session: str, *, connect_timeout: int = 4, posix_fallback: bool = False, ) -> list[str]: if not is_valid_session_name(session): raise ValueError("invalid remote Session name") direct_command = f"herdr session delete --json -- {session}" fallback_command = ( 'if [ -x "$HOME/.local/bin/herdr" ]; then ' f'exec "$HOME/.local/bin/herdr" session delete --json -- {session}; ' "elif command -v herdr >/dev/null 2>&1; then " f"exec herdr session delete --json -- {session}; " "else echo 'herdr is not installed on the remote host' >&2; exit 127; fi" ) return [ ssh_bin, "-T", "-o", "BatchMode=yes", "-o", "StrictHostKeyChecking=yes", "-o", "ConnectionAttempts=1", "-o", f"ConnectTimeout={connect_timeout}", *ssh_destination_arguments(entry.target), fallback_command if posix_fallback else direct_command, ] def remote_stop_session_command( ssh_bin: str, entry: RemoteEntry, session: str, *, connect_timeout: int = 4, posix_fallback: bool = False, ) -> list[str]: if not is_valid_session_name(session): raise ValueError("invalid remote Session name") direct_command = f"herdr session stop -- {session}" fallback_command = ( 'if [ -x "$HOME/.local/bin/herdr" ]; then ' f'exec "$HOME/.local/bin/herdr" session stop -- {session}; ' "elif command -v herdr >/dev/null 2>&1; then " f"exec herdr session stop -- {session}; " "else echo 'herdr is not installed on the remote host' >&2; exit 127; fi" ) return [ ssh_bin, "-T", "-o", "BatchMode=yes", "-o", "StrictHostKeyChecking=yes", "-o", "ConnectionAttempts=1", "-o", f"ConnectTimeout={connect_timeout}", *ssh_destination_arguments(entry.target), fallback_command if posix_fallback else direct_command, ] 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.""" result: subprocess.CompletedProcess[bytes] | None = None for posix_fallback in (False, True): command = remote_delete_session_command( 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 RemoteSessionDeleteError( f"remote Session delete timed out after {timeout:g}s" ) from error except OSError as error: raise RemoteSessionDeleteError( f"could not start remote Session delete: {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 RemoteSessionDeleteError(" ".join(detail.splitlines())) 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.""" result: subprocess.CompletedProcess[bytes] | None = None for posix_fallback in (False, True): command = remote_stop_session_command( 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 RemoteSessionStopError( f"remote Session stop timed out after {timeout:g}s" ) from error except OSError as error: raise RemoteSessionStopError( f"could not start remote Session stop: {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 RemoteSessionStopError(" ".join(detail.splitlines())) def decode_remote_output(raw: bytes) -> str: if not raw: return "" 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 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() 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]: 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() 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) @staticmethod def _stop_process(process: subprocess.Popen[bytes]) -> None: if process.poll() is not None: process.communicate() return process.terminate() try: process.communicate(timeout=0.25) except subprocess.TimeoutExpired: process.kill() process.communicate() def remote_control_command( ssh_bin: str, entry: RemoteEntry, arguments: Sequence[str], *, platform: str, unset_environment: Sequence[str] = (), ) -> 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. """ import base64 import re import shlex 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) argv = shlex.join(arguments) script = ( unset + 'if [ -x "$HOME/.local/bin/herdr" ]; then ' f'exec "$HOME/.local/bin/herdr" {argv}; ' 'elif command -v herdr >/dev/null 2>&1; then ' f"exec herdr {argv}; " "else echo 'herdr is not installed on the remote host' >&2; exit 127; fi" ) 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) script_text = ( "$ErrorActionPreference='Stop'; " + unset + "$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_bin, "-T", "-o", "BatchMode=yes", "-o", "StrictHostKeyChecking=yes", "-o", "ConnectionAttempts=1", "-o", "ConnectTimeout=4", *ssh_destination_arguments(entry.target), script]