The update popup kept the last 64 KiB of raw bytes and decoded the whole buffer on every refresh. Cutting the buffer, or reading a chunk that ended inside a multi-byte sequence, made the UTF-8 decode fail and sent everything through the GB18030 fallback, garbling the display. Clip the retained tail at a UTF-8 boundary and decode still-growing streams with an incremental decoder that holds back an incomplete trailing sequence instead of treating it as corruption. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
805 lines
28 KiB
Python
805 lines
28 KiB
Python
from __future__ import annotations
|
|
|
|
import codecs
|
|
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 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
|
|
|
|
|
|
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] = (),
|
|
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.
|
|
"""
|
|
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)
|
|
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_bin, "-T", "-o", "BatchMode=yes", "-o", "StrictHostKeyChecking=yes",
|
|
"-o", "ConnectionAttempts=1", "-o", "ConnectTimeout=4",
|
|
*ssh_destination_arguments(entry.target), script]
|