435 lines
15 KiB
Python
435 lines
15 KiB
Python
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
|
|
|
|
|
|
REGISTRY_VERSION = 1
|
|
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."""
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RemoteEntry:
|
|
target: str
|
|
remote_keybindings: str = "local"
|
|
|
|
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'")
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RemoteSession:
|
|
name: str
|
|
running: bool
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RemoteQueryResult:
|
|
entry: RemoteEntry
|
|
sessions: tuple[RemoteSession, ...] = ()
|
|
error: str | None = None
|
|
|
|
|
|
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 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,
|
|
}
|
|
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")
|
|
if not isinstance(target, str) or not isinstance(keybindings, str):
|
|
raise RegistryError(f"{path}: remotes[{index}] has invalid fields")
|
|
try:
|
|
entry = RemoteEntry(target, keybindings)
|
|
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,
|
|
) -> 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,
|
|
]
|
|
|
|
|
|
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))
|
|
|
|
|
|
@dataclass
|
|
class _RemoteJob:
|
|
entry: RemoteEntry
|
|
process: subprocess.Popen[str]
|
|
started_at: float
|
|
|
|
|
|
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()
|
|
stdout, stderr = job.process.communicate()
|
|
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",
|
|
)
|
|
)
|
|
else:
|
|
finished.append(
|
|
parse_remote_session_result(
|
|
job.entry,
|
|
job.process.returncode,
|
|
stdout,
|
|
stderr,
|
|
)
|
|
)
|
|
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()
|
|
try:
|
|
process = subprocess.Popen(
|
|
remote_session_command(
|
|
self.ssh_bin,
|
|
entry,
|
|
connect_timeout=self.connect_timeout,
|
|
),
|
|
stdin=subprocess.DEVNULL,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
text=True,
|
|
)
|
|
except (OSError, ValueError) as error:
|
|
self._completed.append(
|
|
RemoteQueryResult(entry, error=f"cannot start remote query: {error}")
|
|
)
|
|
continue
|
|
self._active[entry.target] = _RemoteJob(entry, process, time.monotonic())
|
|
|
|
@staticmethod
|
|
def _stop_process(process: subprocess.Popen[str]) -> 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()
|