gitea-adminandClaude Fable 5.1
8cac4370bf
AliasRegistry and RemoteRegistry duplicated the lock-file, read-modify- write and atomic-replace logic almost line for line. Move it into a generic JsonRegistry base; each registry now only declares its error type, empty value, parser and serializer. Behaviour, file format, permissions and the corrupt-file guarantee are unchanged. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
791 lines
27 KiB
Python
791 lines
27 KiB
Python
from __future__ import annotations
|
|
|
|
import base64
|
|
import codecs
|
|
from collections import deque
|
|
from dataclasses import dataclass
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import re
|
|
import shlex
|
|
import subprocess
|
|
import time
|
|
from typing import Callable, Iterable, Sequence
|
|
from urllib.parse import SplitResult, unquote, urlsplit
|
|
|
|
from .domain import is_valid_session_name
|
|
from .jsonstore import JsonRegistry
|
|
|
|
|
|
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(JsonRegistry[list[RemoteEntry]]):
|
|
error_type = RegistryError
|
|
|
|
def __init__(self, path: Path | None = None) -> None:
|
|
super().__init__(path or default_registry_path())
|
|
|
|
def _empty(self) -> list[RemoteEntry]:
|
|
return []
|
|
|
|
def _parse(self, payload: object) -> list[RemoteEntry]:
|
|
return _parse_registry(payload, self.path)
|
|
|
|
def _serialize(self, entries: list[RemoteEntry]) -> object:
|
|
return {
|
|
"version": REGISTRY_VERSION,
|
|
"remotes": [
|
|
{
|
|
"target": entry.target,
|
|
"remote_keybindings": entry.remote_keybindings,
|
|
"remote_transport": entry.remote_transport,
|
|
}
|
|
for entry in entries
|
|
],
|
|
}
|
|
|
|
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 _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)
|