1558 lines
53 KiB
Python
1558 lines
53 KiB
Python
from __future__ import annotations
|
||
|
||
import argparse
|
||
from contextlib import contextmanager
|
||
import curses
|
||
from dataclasses import dataclass, field, replace
|
||
import fcntl
|
||
from functools import wraps
|
||
import hashlib
|
||
import json
|
||
import os
|
||
from pathlib import Path
|
||
import shlex
|
||
import shutil
|
||
import stat
|
||
import subprocess
|
||
import sys
|
||
import tempfile
|
||
import threading
|
||
from typing import Callable, Iterable, Mapping, Sequence
|
||
from urllib.parse import quote
|
||
|
||
from .aliases import AliasRegistry, AliasRegistryError, default_alias_registry_path
|
||
from .domain import validate_session_name
|
||
from .remotes import (
|
||
RegistryError,
|
||
RemoteEntry,
|
||
RemoteRegistry,
|
||
default_registry_path,
|
||
delete_remote_session,
|
||
ssh_destination_arguments,
|
||
stop_remote_session,
|
||
validate_target,
|
||
)
|
||
from .selector import (
|
||
CatalogWindow,
|
||
DeleteSessionOutcome,
|
||
SelectorError,
|
||
SelectorResult,
|
||
SessionSelector,
|
||
StopSessionOutcome,
|
||
)
|
||
|
||
|
||
DEFAULT_SOCKET_NAME = "herdr-hub"
|
||
DEFAULT_TMUX_SESSION = "hub"
|
||
DEFAULT_SESSION_NAME = "default"
|
||
BOOTSTRAP_WINDOW_NAME = "__herdr_hub_bootstrap__"
|
||
WINDOW_KEY_OPTION = "@herdr_hub_key"
|
||
WINDOW_COMMAND_OPTION = "@herdr_hub_command"
|
||
WINDOW_ENDPOINT_OPTION = "@herdr_hub_endpoint"
|
||
WINDOW_SHORT_LABEL_OPTION = "@herdr_hub_short_label"
|
||
SESSION_DISMISSED_OPTION = "@herdr_hub_dismissed"
|
||
|
||
|
||
class HubError(RuntimeError):
|
||
"""A user-facing failure."""
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class Endpoint:
|
||
kind: str
|
||
session: str = DEFAULT_SESSION_NAME
|
||
target: str | None = None
|
||
remote_keybindings: str = "local"
|
||
remote_transport: str = "native"
|
||
|
||
def __post_init__(self) -> None:
|
||
if self.kind not in {"local", "remote"}:
|
||
raise ValueError(f"unsupported endpoint kind: {self.kind}")
|
||
validate_session_name(self.session)
|
||
if self.kind == "local":
|
||
if self.target is not None:
|
||
raise ValueError("a local endpoint cannot have a remote target")
|
||
if self.remote_transport != "native":
|
||
raise ValueError("a local endpoint cannot use a remote transport")
|
||
return
|
||
validate_remote_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'")
|
||
|
||
@classmethod
|
||
def local(cls, session: str = DEFAULT_SESSION_NAME) -> Endpoint:
|
||
return cls(kind="local", session=session)
|
||
|
||
@classmethod
|
||
def remote(
|
||
cls,
|
||
target: str,
|
||
session: str = DEFAULT_SESSION_NAME,
|
||
remote_keybindings: str = "local",
|
||
remote_transport: str = "native",
|
||
) -> Endpoint:
|
||
return cls(
|
||
kind="remote",
|
||
target=target,
|
||
session=session,
|
||
remote_keybindings=remote_keybindings,
|
||
remote_transport=remote_transport,
|
||
)
|
||
|
||
@property
|
||
def key(self) -> str:
|
||
if self.kind == "local":
|
||
return f"local:{quote(self.session, safe='')}"
|
||
assert self.target is not None
|
||
return f"remote:{quote(self.target, safe='')}:{quote(self.session, safe='')}"
|
||
|
||
@property
|
||
def label(self) -> str:
|
||
if self.kind == "local":
|
||
return truncate_label(f"L {self.session}")
|
||
assert self.target is not None
|
||
return truncate_label(f"R {display_target(self.target)}/{self.session}")
|
||
|
||
@property
|
||
def short_label(self) -> str:
|
||
if self.kind == "local":
|
||
label = "local"
|
||
if self.session != DEFAULT_SESSION_NAME:
|
||
label += f"/{self.session}"
|
||
return truncate_label(label, 32)
|
||
assert self.target is not None
|
||
label = display_target(self.target)
|
||
if self.session != DEFAULT_SESSION_NAME:
|
||
label += f"/{self.session}"
|
||
return truncate_label(label, 32)
|
||
|
||
def command(self, herdr_bin: str, ssh_bin: str = "ssh") -> list[str]:
|
||
command = [herdr_bin]
|
||
if self.kind == "remote":
|
||
assert self.target is not None
|
||
if self.remote_transport == "ssh":
|
||
remote_command = ["herdr"]
|
||
if self.session != DEFAULT_SESSION_NAME:
|
||
remote_command.extend(["--session", self.session])
|
||
return [
|
||
ssh_bin,
|
||
"-tt",
|
||
"-o",
|
||
"ServerAliveInterval=15",
|
||
"-o",
|
||
"ServerAliveCountMax=3",
|
||
*ssh_destination_arguments(self.target),
|
||
shlex.join(remote_command),
|
||
]
|
||
command.extend(["--remote", self.target])
|
||
if self.remote_keybindings != "local":
|
||
command.extend(["--remote-keybindings", self.remote_keybindings])
|
||
if self.session != DEFAULT_SESSION_NAME:
|
||
command.extend(["--session", self.session])
|
||
return command
|
||
|
||
|
||
@dataclass
|
||
class DiscoveryReport:
|
||
endpoints: list[Endpoint] = field(default_factory=list)
|
||
warnings: list[str] = field(default_factory=list)
|
||
|
||
def extend(self, other: DiscoveryReport) -> None:
|
||
self.endpoints.extend(other.endpoints)
|
||
self.warnings.extend(other.warnings)
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class ManagedWindow:
|
||
window_id: str
|
||
key: str
|
||
name: str
|
||
dead: bool
|
||
active: bool
|
||
endpoint: Endpoint | None = None
|
||
|
||
|
||
RunFunction = Callable[..., subprocess.CompletedProcess[str]]
|
||
|
||
|
||
def serialized_hub_mutation(method): # type: ignore[no-untyped-def]
|
||
@wraps(method)
|
||
def wrapper(self, *args, **kwargs): # type: ignore[no-untyped-def]
|
||
with self.mutation_lock():
|
||
return method(self, *args, **kwargs)
|
||
|
||
return wrapper
|
||
|
||
|
||
def validate_remote_target(target: str | None) -> None:
|
||
validate_target(target or "")
|
||
|
||
|
||
def truncate_label(label: str, limit: int = 48) -> str:
|
||
clean = "".join(
|
||
character if ord(character) >= 32 and ord(character) != 127 else "?"
|
||
for character in label
|
||
)
|
||
if len(clean) <= limit:
|
||
return clean
|
||
return clean[: limit - 1] + "…"
|
||
|
||
|
||
def display_target(target: str) -> str:
|
||
if target.startswith("ssh://"):
|
||
return target.removeprefix("ssh://")
|
||
return target
|
||
|
||
|
||
def endpoint_sort_key(endpoint: Endpoint) -> tuple[int, str, str]:
|
||
if endpoint.kind == "local":
|
||
session_rank = "" if endpoint.session == DEFAULT_SESSION_NAME else endpoint.session
|
||
return (0, "", session_rank)
|
||
return (1, endpoint.target or "", endpoint.session)
|
||
|
||
|
||
def dedupe_endpoints(endpoints: Iterable[Endpoint]) -> list[Endpoint]:
|
||
by_key: dict[str, Endpoint] = {}
|
||
for endpoint in endpoints:
|
||
by_key.setdefault(endpoint.key, endpoint)
|
||
return sorted(by_key.values(), key=endpoint_sort_key)
|
||
|
||
|
||
def discover_local_sessions(
|
||
herdr_bin: str,
|
||
*,
|
||
run: RunFunction = subprocess.run,
|
||
) -> DiscoveryReport:
|
||
try:
|
||
result = run(
|
||
[herdr_bin, "session", "list", "--json"],
|
||
capture_output=True,
|
||
text=True,
|
||
timeout=5,
|
||
check=False,
|
||
)
|
||
except (OSError, subprocess.SubprocessError) as error:
|
||
return DiscoveryReport(warnings=[f"local Herdr session discovery failed: {error}"])
|
||
|
||
if result.returncode != 0:
|
||
detail = (result.stderr or result.stdout or "unknown error").strip()
|
||
return DiscoveryReport(
|
||
warnings=[f"local Herdr session discovery failed: {detail}"]
|
||
)
|
||
|
||
try:
|
||
payload = json.loads(result.stdout)
|
||
sessions = payload["sessions"]
|
||
except (json.JSONDecodeError, KeyError, TypeError) as error:
|
||
return DiscoveryReport(
|
||
warnings=[f"local Herdr returned invalid session JSON: {error}"]
|
||
)
|
||
|
||
endpoints: list[Endpoint] = []
|
||
for item in sessions:
|
||
if not isinstance(item, dict) or item.get("running") is not True:
|
||
continue
|
||
name = item.get("name")
|
||
if not isinstance(name, str):
|
||
continue
|
||
try:
|
||
endpoints.append(Endpoint.local(name))
|
||
except ValueError:
|
||
continue
|
||
return DiscoveryReport(endpoints=dedupe_endpoints(endpoints))
|
||
|
||
|
||
def parse_endpoint_argv(argv: Sequence[str]) -> Endpoint | None:
|
||
if not argv or Path(argv[0]).name != "herdr":
|
||
return None
|
||
|
||
target: str | None = None
|
||
session = DEFAULT_SESSION_NAME
|
||
keybindings = "local"
|
||
index = 1
|
||
while index < len(argv):
|
||
argument = argv[index]
|
||
if argument == "--":
|
||
break
|
||
if argument == "--remote":
|
||
if index + 1 >= len(argv):
|
||
return None
|
||
target = argv[index + 1]
|
||
index += 2
|
||
continue
|
||
if argument.startswith("--remote="):
|
||
target = argument.split("=", 1)[1]
|
||
index += 1
|
||
continue
|
||
if argument == "--session":
|
||
if index + 1 >= len(argv):
|
||
return None
|
||
session = argv[index + 1]
|
||
index += 2
|
||
continue
|
||
if argument.startswith("--session="):
|
||
session = argument.split("=", 1)[1]
|
||
index += 1
|
||
continue
|
||
if argument == "--remote-keybindings":
|
||
if index + 1 >= len(argv):
|
||
return None
|
||
keybindings = argv[index + 1]
|
||
index += 2
|
||
continue
|
||
if argument.startswith("--remote-keybindings="):
|
||
keybindings = argument.split("=", 1)[1]
|
||
index += 1
|
||
continue
|
||
index += 1
|
||
|
||
try:
|
||
return (
|
||
Endpoint.remote(target, session, keybindings)
|
||
if target is not None
|
||
else Endpoint.local(session)
|
||
)
|
||
except ValueError:
|
||
return None
|
||
|
||
|
||
def parse_remote_argv(argv: Sequence[str]) -> Endpoint | None:
|
||
endpoint = parse_endpoint_argv(argv)
|
||
return endpoint if endpoint is not None and endpoint.kind == "remote" else None
|
||
|
||
|
||
def serialize_endpoint(endpoint: Endpoint) -> str:
|
||
return json.dumps(
|
||
{
|
||
"version": 1,
|
||
"kind": endpoint.kind,
|
||
"session": endpoint.session,
|
||
"target": endpoint.target,
|
||
"remote_keybindings": endpoint.remote_keybindings,
|
||
"remote_transport": endpoint.remote_transport,
|
||
},
|
||
ensure_ascii=False,
|
||
separators=(",", ":"),
|
||
)
|
||
|
||
|
||
def deserialize_endpoint(metadata: str, command: str = "") -> Endpoint | None:
|
||
if metadata:
|
||
try:
|
||
payload = json.loads(metadata)
|
||
if not isinstance(payload, dict) or payload.get("version") != 1:
|
||
raise ValueError
|
||
kind = payload.get("kind")
|
||
session = payload.get("session")
|
||
target = payload.get("target")
|
||
keybindings = payload.get("remote_keybindings", "local")
|
||
transport = payload.get("remote_transport", "native")
|
||
if not isinstance(kind, str) or not isinstance(session, str):
|
||
raise ValueError
|
||
if not isinstance(keybindings, str):
|
||
raise ValueError
|
||
if not isinstance(transport, str):
|
||
raise ValueError
|
||
if kind == "local":
|
||
return Endpoint.local(session)
|
||
if kind == "remote" and isinstance(target, str):
|
||
return Endpoint.remote(target, session, keybindings, transport)
|
||
raise ValueError
|
||
except (json.JSONDecodeError, TypeError, ValueError):
|
||
pass
|
||
if not command:
|
||
return None
|
||
try:
|
||
arguments = shlex.split(command)
|
||
except ValueError:
|
||
return None
|
||
if arguments and arguments[0] == "exec":
|
||
arguments = arguments[1:]
|
||
return parse_endpoint_argv(arguments)
|
||
|
||
|
||
def decode_cmdline(raw: bytes) -> list[str]:
|
||
return [
|
||
os.fsdecode(argument)
|
||
for argument in raw.rstrip(b"\0").split(b"\0")
|
||
if argument
|
||
]
|
||
|
||
|
||
def discover_remote_processes(
|
||
proc_root: Path = Path("/proc"),
|
||
*,
|
||
user_id: int | None = None,
|
||
) -> DiscoveryReport:
|
||
if user_id is None:
|
||
user_id = os.getuid()
|
||
if not proc_root.is_dir():
|
||
return DiscoveryReport(
|
||
warnings=[f"remote discovery unavailable: {proc_root} is not mounted"]
|
||
)
|
||
|
||
endpoints: list[Endpoint] = []
|
||
unreadable = 0
|
||
try:
|
||
entries = sorted(
|
||
(entry for entry in proc_root.iterdir() if entry.name.isdigit()),
|
||
key=lambda entry: int(entry.name),
|
||
)
|
||
except OSError as error:
|
||
return DiscoveryReport(warnings=[f"remote process discovery failed: {error}"])
|
||
|
||
for entry in entries:
|
||
cmdline_path = entry / "cmdline"
|
||
try:
|
||
if cmdline_path.stat().st_uid != user_id:
|
||
continue
|
||
raw = cmdline_path.read_bytes()
|
||
except (FileNotFoundError, ProcessLookupError):
|
||
continue
|
||
except (PermissionError, OSError):
|
||
unreadable += 1
|
||
continue
|
||
endpoint = parse_remote_argv(decode_cmdline(raw))
|
||
if endpoint is not None:
|
||
endpoints.append(endpoint)
|
||
|
||
warnings: list[str] = []
|
||
if unreadable:
|
||
warnings.append(f"could not inspect {unreadable} same-user process command line(s)")
|
||
return DiscoveryReport(endpoints=dedupe_endpoints(endpoints), warnings=warnings)
|
||
|
||
|
||
def discover_endpoints(herdr_bin: str, proc_root: Path = Path("/proc")) -> DiscoveryReport:
|
||
report = discover_local_sessions(herdr_bin)
|
||
report.extend(discover_remote_processes(proc_root))
|
||
report.endpoints = dedupe_endpoints(report.endpoints)
|
||
return report
|
||
|
||
|
||
class TmuxHub:
|
||
def __init__(
|
||
self,
|
||
*,
|
||
tmux_bin: str,
|
||
ssh_bin: str = "ssh",
|
||
socket_name: str = DEFAULT_SOCKET_NAME,
|
||
session_name: str = DEFAULT_TMUX_SESSION,
|
||
config_file: str | None = None,
|
||
script_path: str | None = None,
|
||
python_bin: str | None = None,
|
||
aliases: Mapping[str, str] | None = None,
|
||
run: RunFunction = subprocess.run,
|
||
) -> None:
|
||
self.tmux_bin = tmux_bin
|
||
self.ssh_bin = ssh_bin
|
||
self.socket_name = socket_name
|
||
self.session_name = session_name
|
||
self.config_file = config_file
|
||
self.script_path = script_path or str(Path(sys.argv[0]).absolute())
|
||
self.python_bin = python_bin or sys.executable
|
||
self.aliases = dict(aliases or {})
|
||
self._run = run
|
||
self._local_mutation_lock = threading.RLock()
|
||
self._mutation_depth = 0
|
||
|
||
@property
|
||
def mutation_lock_path(self) -> Path:
|
||
identity = f"{os.getuid()}\0{self.socket_name}\0{self.session_name}"
|
||
digest = hashlib.sha256(identity.encode("utf-8")).hexdigest()[:20]
|
||
roots = []
|
||
if configured_runtime := os.environ.get("XDG_RUNTIME_DIR"):
|
||
roots.append(Path(configured_runtime).expanduser())
|
||
temporary_root = Path(tempfile.gettempdir())
|
||
if temporary_root not in roots:
|
||
roots.append(temporary_root)
|
||
|
||
failures: list[str] = []
|
||
for runtime_root in roots:
|
||
lock_directory = runtime_root / f"herdr-hub-{os.getuid()}"
|
||
try:
|
||
lock_directory.mkdir(mode=0o700, parents=False, exist_ok=True)
|
||
metadata = lock_directory.lstat()
|
||
if (
|
||
not stat.S_ISDIR(metadata.st_mode)
|
||
or metadata.st_uid != os.getuid()
|
||
):
|
||
failures.append(f"unsafe directory {lock_directory}")
|
||
continue
|
||
os.chmod(lock_directory, 0o700)
|
||
return lock_directory / f"{digest}.lock"
|
||
except OSError as error:
|
||
failures.append(f"{lock_directory}: {error}")
|
||
raise HubError("could not prepare Hub lock directory: " + "; ".join(failures))
|
||
|
||
@contextmanager
|
||
def mutation_lock(self): # type: ignore[no-untyped-def]
|
||
"""Serialize state-changing Hub commands across CLI processes."""
|
||
with self._local_mutation_lock:
|
||
if self._mutation_depth:
|
||
self._mutation_depth += 1
|
||
try:
|
||
yield
|
||
finally:
|
||
self._mutation_depth -= 1
|
||
return
|
||
|
||
descriptor = os.open(
|
||
self.mutation_lock_path,
|
||
os.O_CREAT | os.O_RDWR,
|
||
0o600,
|
||
)
|
||
try:
|
||
os.fchmod(descriptor, 0o600)
|
||
fcntl.flock(descriptor, fcntl.LOCK_EX)
|
||
self._mutation_depth = 1
|
||
try:
|
||
yield
|
||
finally:
|
||
self._mutation_depth = 0
|
||
fcntl.flock(descriptor, fcntl.LOCK_UN)
|
||
finally:
|
||
os.close(descriptor)
|
||
|
||
def base_command(self) -> list[str]:
|
||
command = [self.tmux_bin, "-L", self.socket_name]
|
||
if self.config_file is not None:
|
||
command.extend(["-f", self.config_file])
|
||
return command
|
||
|
||
def run(
|
||
self,
|
||
arguments: Sequence[str],
|
||
*,
|
||
check: bool = True,
|
||
capture_output: bool = True,
|
||
text: bool = True,
|
||
) -> subprocess.CompletedProcess[str]:
|
||
try:
|
||
result = self._run(
|
||
[*self.base_command(), *arguments],
|
||
capture_output=capture_output,
|
||
text=text,
|
||
check=False,
|
||
)
|
||
except OSError as error:
|
||
raise HubError(f"could not run tmux: {error}") from error
|
||
if check and result.returncode != 0:
|
||
detail = (result.stderr or result.stdout or "tmux command failed").strip()
|
||
raise HubError(detail)
|
||
return result
|
||
|
||
def exists(self) -> bool:
|
||
result = self.run(
|
||
["has-session", "-t", self.session_name],
|
||
check=False,
|
||
)
|
||
return result.returncode == 0
|
||
|
||
@serialized_hub_mutation
|
||
def ensure_server(self) -> bool:
|
||
if self.exists():
|
||
self.configure()
|
||
return False
|
||
bootstrap = shlex.join(["sleep", "infinity"])
|
||
self.run(
|
||
[
|
||
"new-session",
|
||
"-d",
|
||
"-s",
|
||
self.session_name,
|
||
"-n",
|
||
BOOTSTRAP_WINDOW_NAME,
|
||
bootstrap,
|
||
]
|
||
)
|
||
self.configure()
|
||
return True
|
||
|
||
def configure(self) -> None:
|
||
option_commands = [
|
||
["set-option", "-s", "set-clipboard", "on"],
|
||
["set-option", "-g", "prefix", "C-a"],
|
||
["set-option", "-g", "status", "on"],
|
||
["set-option", "-g", "status-position", "top"],
|
||
["set-option", "-g", "status-interval", "2"],
|
||
["set-option", "-g", "mouse", "on"],
|
||
["set-option", "-g", "focus-events", "on"],
|
||
["set-option", "-g", "status-left-length", "18"],
|
||
["set-option", "-g", "status-right-length", "42"],
|
||
[
|
||
"set-option",
|
||
"-g",
|
||
"status-left",
|
||
"#{?#{<:#{client_width},50},"
|
||
"#[bold#,fg=colour81] S #[default],"
|
||
"#[bold#,fg=colour81] Sessions #[default]}",
|
||
],
|
||
[
|
||
"set-option",
|
||
"-g",
|
||
"status-right",
|
||
"#{?#{<:#{client_width},80},,"
|
||
"#[dim]#{?pane_dead,连接已断开 · C-a R 重连,"
|
||
"C-a g choose · q close · r refresh}#[default]}",
|
||
],
|
||
["set-option", "-g", "status-style", "bg=colour235,fg=colour250"],
|
||
["set-window-option", "-g", "window-status-separator", ""],
|
||
[
|
||
"set-window-option",
|
||
"-g",
|
||
"window-status-format",
|
||
"#{?#{<:#{client_width},80},"
|
||
" #{?pane_dead,×,○} #{=12:@herdr_hub_short_label} ,"
|
||
" #{?pane_dead,× #W [断开],○ #W} }",
|
||
],
|
||
[
|
||
"set-window-option",
|
||
"-g",
|
||
"window-status-current-format",
|
||
"#[bold,bg=colour31,fg=colour231]"
|
||
"#{?#{<:#{client_width},80},"
|
||
" #{?pane_dead,×,●} #{=12:@herdr_hub_short_label} ,"
|
||
" #{?pane_dead,× #W [断开],● #W} }"
|
||
"#[default]",
|
||
],
|
||
["set-window-option", "-g", "automatic-rename", "off"],
|
||
["set-window-option", "-g", "allow-rename", "off"],
|
||
["set-window-option", "-g", "remain-on-exit", "on"],
|
||
[
|
||
"set-window-option",
|
||
"-g",
|
||
"remain-on-exit-format",
|
||
"#[bold,fg=colour196]连接已断开#[default]"
|
||
"（退出状态 #{pane_dead_status}） · 按 Ctrl+A R 重连",
|
||
],
|
||
]
|
||
for command in option_commands:
|
||
self.run(command)
|
||
|
||
refresh = shlex.join([self.python_bin, self.script_path, "refresh"])
|
||
choose = shlex.join([self.python_bin, self.script_path, "choose"])
|
||
reconnect = shlex.join([self.python_bin, self.script_path, "reconnect"])
|
||
close = shlex.join(
|
||
[self.python_bin, self.script_path, "close", "--window", "#{window_id}"]
|
||
)
|
||
confirm_close = shlex.join(["run-shell", "-b", close])
|
||
# Remove tmux's default/legacy '&' binding when reconfiguring both new
|
||
# and already-running Hub servers.
|
||
self.run(["unbind-key", "-T", "prefix", "&"], check=False)
|
||
self.run(["unbind-key", "-T", "prefix", "Q"], check=False)
|
||
binding_commands = [
|
||
["bind-key", "-T", "prefix", "C-a", "send-prefix"],
|
||
["bind-key", "-T", "prefix", "r", "run-shell", "-b", refresh],
|
||
[
|
||
"bind-key",
|
||
"-T",
|
||
"prefix",
|
||
"g",
|
||
"display-popup",
|
||
"-EE",
|
||
"-w",
|
||
"90%",
|
||
"-h",
|
||
"90%",
|
||
"-T",
|
||
"Sessions",
|
||
choose,
|
||
],
|
||
[
|
||
"bind-key",
|
||
"-T",
|
||
"root",
|
||
"MouseDown1StatusLeft",
|
||
"display-popup",
|
||
"-EE",
|
||
"-w",
|
||
"90%",
|
||
"-h",
|
||
"90%",
|
||
"-T",
|
||
"Sessions",
|
||
choose,
|
||
],
|
||
["bind-key", "-T", "prefix", "R", "run-shell", "-b", reconnect],
|
||
[
|
||
"bind-key",
|
||
"-T",
|
||
"prefix",
|
||
"q",
|
||
"confirm-before",
|
||
"-p",
|
||
"关闭当前 Session Tab？ (y/n)",
|
||
confirm_close,
|
||
],
|
||
]
|
||
for command in binding_commands:
|
||
self.run(command)
|
||
|
||
self.refresh_labels()
|
||
|
||
def managed_windows(self) -> list[ManagedWindow]:
|
||
if not self.exists():
|
||
return []
|
||
line_format = (
|
||
f"#{{window_id}}\t#{{{WINDOW_KEY_OPTION}}}\t#{{window_name}}\t"
|
||
f"#{{pane_dead}}\t#{{window_active}}\t#{{{WINDOW_ENDPOINT_OPTION}}}\t"
|
||
f"#{{{WINDOW_COMMAND_OPTION}}}"
|
||
)
|
||
result = self.run(
|
||
["list-windows", "-t", self.session_name, "-F", line_format]
|
||
)
|
||
windows: list[ManagedWindow] = []
|
||
for line in result.stdout.splitlines():
|
||
fields = line.split("\t", 6)
|
||
if len(fields) != 7 or not fields[1]:
|
||
continue
|
||
windows.append(
|
||
ManagedWindow(
|
||
window_id=fields[0],
|
||
key=fields[1],
|
||
name=fields[2],
|
||
dead=fields[3] == "1",
|
||
active=fields[4] == "1",
|
||
endpoint=deserialize_endpoint(fields[5], fields[6]),
|
||
)
|
||
)
|
||
return windows
|
||
|
||
def window_for_key(self, key: str) -> ManagedWindow | None:
|
||
return next((window for window in self.managed_windows() if window.key == key), None)
|
||
|
||
def refresh_labels(self) -> None:
|
||
for window in self.managed_windows():
|
||
self._apply_window_label(window)
|
||
|
||
def _apply_window_label(self, window: ManagedWindow) -> None:
|
||
alias = self.aliases.get(window.key)
|
||
display_label = truncate_label(alias) if alias else window.name
|
||
short_label = truncate_label(alias, 32) if alias else truncate_label(window.name, 32)
|
||
if window.endpoint is not None:
|
||
display_label = truncate_label(alias) if alias else window.endpoint.label
|
||
short_label = (
|
||
truncate_label(alias, 32) if alias else window.endpoint.short_label
|
||
)
|
||
if window.name != display_label:
|
||
self.run(["rename-window", "-t", window.window_id, display_label])
|
||
self.run(
|
||
[
|
||
"set-option",
|
||
"-w",
|
||
"-t",
|
||
window.window_id,
|
||
WINDOW_SHORT_LABEL_OPTION,
|
||
short_label,
|
||
]
|
||
)
|
||
|
||
@serialized_hub_mutation
|
||
def apply_alias(self, key: str, label: str) -> ManagedWindow | None:
|
||
if label:
|
||
self.aliases[key] = label
|
||
else:
|
||
self.aliases.pop(key, None)
|
||
window = self.window_for_key(key)
|
||
if window is not None:
|
||
self._apply_window_label(window)
|
||
return self.window_for_key(key)
|
||
|
||
def dismissed_keys(self) -> set[str]:
|
||
result = self.run(
|
||
[
|
||
"show-options",
|
||
"-v",
|
||
"-t",
|
||
self.session_name,
|
||
SESSION_DISMISSED_OPTION,
|
||
],
|
||
check=False,
|
||
)
|
||
if result.returncode != 0 or not result.stdout.strip():
|
||
return set()
|
||
try:
|
||
payload = json.loads(result.stdout)
|
||
except json.JSONDecodeError:
|
||
return set()
|
||
if not isinstance(payload, list):
|
||
return set()
|
||
return {key for key in payload if isinstance(key, str) and key}
|
||
|
||
def set_dismissed_keys(self, keys: Iterable[str]) -> None:
|
||
value = json.dumps(sorted(set(keys)), ensure_ascii=False, separators=(",", ":"))
|
||
self.run(
|
||
[
|
||
"set-option",
|
||
"-t",
|
||
self.session_name,
|
||
SESSION_DISMISSED_OPTION,
|
||
value,
|
||
]
|
||
)
|
||
|
||
def restore_endpoint(self, key: str) -> None:
|
||
dismissed = self.dismissed_keys()
|
||
if key not in dismissed:
|
||
return
|
||
dismissed.remove(key)
|
||
self.set_dismissed_keys(dismissed)
|
||
|
||
@serialized_hub_mutation
|
||
def create_window(self, endpoint: Endpoint, herdr_bin: str) -> ManagedWindow:
|
||
existing = self.window_for_key(endpoint.key)
|
||
if existing is not None:
|
||
return existing
|
||
was_dismissed = endpoint.key in self.dismissed_keys()
|
||
self.restore_endpoint(endpoint.key)
|
||
command = "exec " + shlex.join(endpoint.command(herdr_bin, self.ssh_bin))
|
||
alias = self.aliases.get(endpoint.key)
|
||
display_label = truncate_label(alias) if alias else endpoint.label
|
||
short_label = truncate_label(alias, 32) if alias else endpoint.short_label
|
||
result = self.run(
|
||
[
|
||
"new-window",
|
||
"-d",
|
||
"-P",
|
||
"-F",
|
||
"#{window_id}",
|
||
"-t",
|
||
f"{self.session_name}:",
|
||
"-n",
|
||
display_label,
|
||
command,
|
||
]
|
||
)
|
||
window_id = result.stdout.strip()
|
||
if not window_id:
|
||
if was_dismissed:
|
||
self.set_dismissed_keys({*self.dismissed_keys(), endpoint.key})
|
||
raise HubError("tmux did not return the new window id")
|
||
try:
|
||
for option, value in [
|
||
(WINDOW_KEY_OPTION, endpoint.key),
|
||
(WINDOW_COMMAND_OPTION, command),
|
||
(WINDOW_ENDPOINT_OPTION, serialize_endpoint(endpoint)),
|
||
(WINDOW_SHORT_LABEL_OPTION, short_label),
|
||
]:
|
||
self.run(["set-option", "-w", "-t", window_id, option, value])
|
||
self.run(["set-window-option", "-t", window_id, "remain-on-exit", "on"])
|
||
except HubError:
|
||
self.run(["kill-window", "-t", window_id], check=False)
|
||
if was_dismissed:
|
||
self.set_dismissed_keys({*self.dismissed_keys(), endpoint.key})
|
||
raise
|
||
return ManagedWindow(
|
||
window_id=window_id,
|
||
key=endpoint.key,
|
||
name=display_label,
|
||
dead=False,
|
||
active=False,
|
||
endpoint=endpoint,
|
||
)
|
||
|
||
@serialized_hub_mutation
|
||
def add_missing(self, endpoints: Sequence[Endpoint], herdr_bin: str) -> list[ManagedWindow]:
|
||
existing_keys = {window.key for window in self.managed_windows()}
|
||
dismissed_keys = self.dismissed_keys()
|
||
added: list[ManagedWindow] = []
|
||
for endpoint in endpoints:
|
||
if endpoint.key in existing_keys or endpoint.key in dismissed_keys:
|
||
continue
|
||
window = self.create_window(endpoint, herdr_bin)
|
||
existing_keys.add(endpoint.key)
|
||
added.append(window)
|
||
if self.managed_windows():
|
||
self.remove_bootstrap()
|
||
return added
|
||
|
||
@serialized_hub_mutation
|
||
def remove_bootstrap(self) -> None:
|
||
result = self.run(
|
||
[
|
||
"list-windows",
|
||
"-t",
|
||
self.session_name,
|
||
"-F",
|
||
"#{window_id}\t#{window_name}",
|
||
],
|
||
check=False,
|
||
)
|
||
if result.returncode != 0:
|
||
return
|
||
for line in result.stdout.splitlines():
|
||
fields = line.split("\t", 1)
|
||
if len(fields) == 2 and fields[1] == BOOTSTRAP_WINDOW_NAME:
|
||
self.run(["kill-window", "-t", fields[0]], check=False)
|
||
|
||
def select_window(self, window_id: str) -> None:
|
||
self.run(["select-window", "-t", window_id])
|
||
|
||
@serialized_hub_mutation
|
||
def dismiss_window(self, window_id: str | None = None) -> ManagedWindow:
|
||
window_id = window_id or self.current_window_id()
|
||
windows = self.managed_windows()
|
||
window = next((item for item in windows if item.window_id == window_id), None)
|
||
if window is None or window.endpoint is None:
|
||
raise HubError("the selected window is not managed by herdr-hub")
|
||
if len(windows) <= 1:
|
||
raise HubError("至少保留一个 Session Tab；可用 herdr-hub stop 退出整个 Hub")
|
||
|
||
previous = self.dismissed_keys()
|
||
self.set_dismissed_keys({*previous, window.key})
|
||
result = self.run(["kill-window", "-t", window.window_id], check=False)
|
||
if result.returncode != 0:
|
||
self.set_dismissed_keys(previous)
|
||
detail = (result.stderr or result.stdout or "tmux command failed").strip()
|
||
raise HubError(detail)
|
||
return window
|
||
|
||
def select_primary(self, endpoints: Sequence[Endpoint]) -> None:
|
||
preferred = next(
|
||
(endpoint for endpoint in endpoints if endpoint.key == Endpoint.local().key),
|
||
endpoints[0] if endpoints else None,
|
||
)
|
||
if preferred is None:
|
||
return
|
||
window = self.window_for_key(preferred.key)
|
||
if window is not None:
|
||
self.select_window(window.window_id)
|
||
|
||
def current_window_id(self) -> str:
|
||
arguments = ["display-message", "-p"]
|
||
if pane_id := os.environ.get("TMUX_PANE"):
|
||
arguments.extend(["-t", pane_id])
|
||
arguments.append("#{window_id}")
|
||
result = self.run(arguments)
|
||
window_id = result.stdout.strip()
|
||
if not window_id:
|
||
raise HubError("could not identify the current hub window")
|
||
return window_id
|
||
|
||
@serialized_hub_mutation
|
||
def reconnect(self, window_id: str | None = None) -> None:
|
||
window_id = window_id or self.current_window_id()
|
||
command_result = self.run(
|
||
["show-options", "-w", "-v", "-t", window_id, WINDOW_COMMAND_OPTION],
|
||
check=False,
|
||
)
|
||
command = command_result.stdout.strip()
|
||
if command_result.returncode != 0 or not command:
|
||
raise HubError("the selected window is not managed by herdr-hub")
|
||
self.run(["respawn-window", "-k", "-t", window_id, command])
|
||
|
||
@serialized_hub_mutation
|
||
def replace_endpoint(
|
||
self,
|
||
window: ManagedWindow,
|
||
endpoint: Endpoint,
|
||
herdr_bin: str,
|
||
) -> ManagedWindow:
|
||
command = "exec " + shlex.join(endpoint.command(herdr_bin, self.ssh_bin))
|
||
alias = self.aliases.get(endpoint.key)
|
||
short_label = truncate_label(alias, 32) if alias else endpoint.short_label
|
||
replacements = [
|
||
(WINDOW_COMMAND_OPTION, command),
|
||
(WINDOW_ENDPOINT_OPTION, serialize_endpoint(endpoint)),
|
||
(WINDOW_SHORT_LABEL_OPTION, short_label),
|
||
]
|
||
previous: dict[str, str] = {}
|
||
for option, _ in replacements:
|
||
result = self.run(
|
||
["show-options", "-w", "-v", "-t", window.window_id, option],
|
||
check=False,
|
||
)
|
||
if result.returncode == 0:
|
||
previous[option] = result.stdout.rstrip("\n")
|
||
try:
|
||
for option, value in replacements:
|
||
self.run(
|
||
["set-option", "-w", "-t", window.window_id, option, value]
|
||
)
|
||
self.run(["respawn-window", "-k", "-t", window.window_id, command])
|
||
except HubError:
|
||
for option, value in previous.items():
|
||
self.run(
|
||
["set-option", "-w", "-t", window.window_id, option, value],
|
||
check=False,
|
||
)
|
||
raise
|
||
replaced = self.window_for_key(endpoint.key)
|
||
if replaced is None:
|
||
raise HubError("the selected Session window disappeared during reconnect")
|
||
self._apply_window_label(replaced)
|
||
return self.window_for_key(endpoint.key) or replaced
|
||
|
||
@serialized_hub_mutation
|
||
def activate_endpoint(self, endpoint: Endpoint, herdr_bin: str) -> ManagedWindow:
|
||
window = self.window_for_key(endpoint.key)
|
||
if window is None:
|
||
window = self.create_window(endpoint, herdr_bin)
|
||
self.remove_bootstrap()
|
||
elif window.endpoint is not None and (
|
||
window.endpoint.remote_transport != endpoint.remote_transport
|
||
or window.endpoint.remote_keybindings != endpoint.remote_keybindings
|
||
):
|
||
window = self.replace_endpoint(window, endpoint, herdr_bin)
|
||
elif window.dead:
|
||
self.reconnect(window.window_id)
|
||
self.select_window(window.window_id)
|
||
return window
|
||
|
||
def attach(self) -> None:
|
||
if inside_any_tmux():
|
||
if inside_socket(self.socket_name):
|
||
return
|
||
raise HubError(
|
||
"already inside another tmux server; detach from it before attaching to herdr-hub"
|
||
)
|
||
result = self.run(
|
||
["attach-session", "-t", self.session_name],
|
||
check=False,
|
||
capture_output=False,
|
||
text=True,
|
||
)
|
||
if result.returncode != 0:
|
||
raise HubError("could not attach to the herdr-hub tmux session")
|
||
|
||
def display_message(self, message: str) -> None:
|
||
self.run(["display-message", message], check=False)
|
||
|
||
@serialized_hub_mutation
|
||
def stop(self) -> None:
|
||
if not self.exists():
|
||
return
|
||
self.run(["kill-server"])
|
||
|
||
|
||
def inside_any_tmux() -> bool:
|
||
return bool(os.environ.get("TMUX"))
|
||
|
||
|
||
def inside_socket(socket_name: str) -> bool:
|
||
value = os.environ.get("TMUX", "")
|
||
socket_path = value.split(",", 1)[0]
|
||
return bool(socket_path) and Path(socket_path).name == socket_name
|
||
|
||
|
||
def resolve_binary(name: str, override: str | None = None) -> str:
|
||
if override:
|
||
return override
|
||
resolved = shutil.which(name)
|
||
if resolved is None:
|
||
raise HubError(f"required executable not found on PATH: {name}")
|
||
return resolved
|
||
|
||
|
||
def invocation_script_path() -> str:
|
||
invoked = Path(sys.argv[0]).absolute()
|
||
if invoked.name == "__main__.py":
|
||
repository_launcher = Path(__file__).resolve().parents[1] / "bin" / "herdr-hub"
|
||
return str(repository_launcher)
|
||
return str(invoked)
|
||
|
||
|
||
def make_hub(
|
||
arguments: argparse.Namespace,
|
||
*,
|
||
load_aliases: bool = True,
|
||
) -> TmuxHub:
|
||
tmux_bin = resolve_binary("tmux", os.environ.get("HERDR_HUB_TMUX_BIN"))
|
||
ssh_bin = os.environ.get("HERDR_HUB_SSH_BIN") or shutil.which("ssh") or "ssh"
|
||
aliases = make_alias_registry().load() if load_aliases else {}
|
||
return TmuxHub(
|
||
tmux_bin=tmux_bin,
|
||
ssh_bin=ssh_bin,
|
||
socket_name=os.environ.get("HERDR_HUB_SOCKET", DEFAULT_SOCKET_NAME),
|
||
session_name=os.environ.get("HERDR_HUB_SESSION", DEFAULT_TMUX_SESSION),
|
||
config_file=os.environ.get("HERDR_HUB_TMUX_CONFIG"),
|
||
script_path=invocation_script_path(),
|
||
aliases=aliases,
|
||
)
|
||
|
||
|
||
def resolve_herdr() -> str:
|
||
return resolve_binary("herdr", os.environ.get("HERDR_HUB_HERDR_BIN"))
|
||
|
||
|
||
def resolve_ssh() -> str:
|
||
return resolve_binary("ssh", os.environ.get("HERDR_HUB_SSH_BIN"))
|
||
|
||
|
||
def configured_proc_root() -> Path:
|
||
return Path(os.environ.get("HERDR_HUB_PROC_ROOT", "/proc"))
|
||
|
||
|
||
def make_registry() -> RemoteRegistry:
|
||
return RemoteRegistry(default_registry_path())
|
||
|
||
|
||
def make_alias_registry() -> AliasRegistry:
|
||
return AliasRegistry(default_alias_registry_path())
|
||
|
||
|
||
def remember_remote_endpoints(endpoints: Iterable[Endpoint]) -> str | None:
|
||
entries = [
|
||
RemoteEntry(
|
||
endpoint.target,
|
||
endpoint.remote_keybindings,
|
||
endpoint.remote_transport,
|
||
)
|
||
for endpoint in endpoints
|
||
if endpoint.kind == "remote" and endpoint.target is not None
|
||
]
|
||
if not entries:
|
||
return None
|
||
try:
|
||
make_registry().remember_many(entries)
|
||
except RegistryError as error:
|
||
return f"remote registry: {error}"
|
||
return None
|
||
|
||
|
||
def catalog_windows(windows: Iterable[ManagedWindow]) -> list[CatalogWindow]:
|
||
return [
|
||
CatalogWindow(
|
||
window_id=window.window_id,
|
||
key=window.key,
|
||
label=window.name,
|
||
dead=window.dead,
|
||
active=window.active,
|
||
kind=window.endpoint.kind if window.endpoint is not None else None,
|
||
target=window.endpoint.target if window.endpoint is not None else None,
|
||
session=window.endpoint.session if window.endpoint is not None else None,
|
||
default_label=(
|
||
window.endpoint.label if window.endpoint is not None else window.name
|
||
),
|
||
remote_transport=(
|
||
window.endpoint.remote_transport
|
||
if window.endpoint is not None
|
||
else "native"
|
||
),
|
||
)
|
||
for window in windows
|
||
]
|
||
|
||
|
||
def command_start(arguments: argparse.Namespace) -> int:
|
||
herdr_bin = resolve_herdr()
|
||
report = discover_endpoints(herdr_bin, configured_proc_root())
|
||
if warning := remember_remote_endpoints(report.endpoints):
|
||
report.warnings.append(warning)
|
||
hub = make_hub(arguments)
|
||
created = hub.ensure_server()
|
||
added = hub.add_missing(report.endpoints, herdr_bin)
|
||
windows = hub.managed_windows()
|
||
if not windows:
|
||
hub.stop()
|
||
detail = "; ".join(report.warnings)
|
||
suffix = f" ({detail})" if detail else ""
|
||
raise HubError(
|
||
"no running Herdr sessions or live --remote clients were found; "
|
||
f"use 'herdr-hub open' to create one{suffix}"
|
||
)
|
||
if created:
|
||
hub.select_primary(report.endpoints)
|
||
if added and not created:
|
||
hub.display_message(f"herdr-hub: added {len(added)} session(s)")
|
||
hub.attach()
|
||
return 0
|
||
|
||
|
||
def command_refresh(arguments: argparse.Namespace) -> int:
|
||
herdr_bin = resolve_herdr()
|
||
report = discover_endpoints(herdr_bin, configured_proc_root())
|
||
if warning := remember_remote_endpoints(report.endpoints):
|
||
report.warnings.append(warning)
|
||
hub = make_hub(arguments)
|
||
hub.ensure_server()
|
||
added = hub.add_missing(report.endpoints, herdr_bin)
|
||
message = f"herdr-hub: added {len(added)} session(s)"
|
||
if report.warnings:
|
||
message += f"; {report.warnings[0]}"
|
||
has_windows = bool(hub.managed_windows())
|
||
if not has_windows:
|
||
hub.stop()
|
||
if inside_socket(hub.socket_name) and has_windows:
|
||
hub.display_message(message)
|
||
else:
|
||
print(message)
|
||
return 0
|
||
|
||
|
||
def command_open(arguments: argparse.Namespace) -> int:
|
||
herdr_bin = resolve_herdr()
|
||
try:
|
||
endpoint = (
|
||
Endpoint.remote(
|
||
arguments.remote,
|
||
arguments.session,
|
||
arguments.remote_keybindings,
|
||
arguments.transport,
|
||
)
|
||
if arguments.remote is not None
|
||
else Endpoint.local(arguments.session)
|
||
)
|
||
except ValueError as error:
|
||
raise HubError(str(error)) from error
|
||
if endpoint.kind == "remote":
|
||
assert endpoint.target is not None
|
||
try:
|
||
make_registry().upsert(
|
||
RemoteEntry(
|
||
endpoint.target,
|
||
endpoint.remote_keybindings,
|
||
endpoint.remote_transport,
|
||
)
|
||
)
|
||
except RegistryError as error:
|
||
print(f"herdr-hub: warning: remote registry: {error}", file=sys.stderr)
|
||
hub = make_hub(arguments)
|
||
hub.ensure_server()
|
||
hub.activate_endpoint(endpoint, herdr_bin)
|
||
hub.attach()
|
||
return 0
|
||
|
||
|
||
def command_choose(arguments: argparse.Namespace) -> int:
|
||
hub = make_hub(arguments)
|
||
if not hub.exists():
|
||
raise HubError("herdr-hub is not running")
|
||
windows = hub.managed_windows()
|
||
local_report = discover_local_sessions(resolve_herdr())
|
||
registry_warning = remember_remote_endpoints(
|
||
window.endpoint for window in windows if window.endpoint is not None
|
||
)
|
||
|
||
alias_registry = make_alias_registry()
|
||
remote_registry = make_registry()
|
||
ssh_bin = resolve_ssh()
|
||
|
||
def rename_session(endpoint_key: str, alias: str) -> None:
|
||
alias_registry.set(endpoint_key, alias)
|
||
if hub.apply_alias(endpoint_key, alias) is None:
|
||
raise HubError("the selected Session is no longer open")
|
||
|
||
def delete_session(
|
||
target: str,
|
||
session: str,
|
||
window_id: str | None,
|
||
) -> DeleteSessionOutcome:
|
||
delete_remote_session(ssh_bin, RemoteEntry(target), session)
|
||
endpoint_key = Endpoint.remote(target, session).key
|
||
warnings: list[str] = []
|
||
tab_closed = False
|
||
if window_id is not None:
|
||
try:
|
||
hub.dismiss_window(window_id)
|
||
tab_closed = True
|
||
except HubError as error:
|
||
warnings.append(f"Tab 未关闭：{error}")
|
||
try:
|
||
alias_registry.set(endpoint_key, "")
|
||
except AliasRegistryError as error:
|
||
warnings.append(f"显示名称清理失败：{error}")
|
||
hub.aliases.pop(endpoint_key, None)
|
||
try:
|
||
hub.restore_endpoint(endpoint_key)
|
||
except HubError as error:
|
||
warnings.append(f"隐藏记录清理失败：{error}")
|
||
return DeleteSessionOutcome(
|
||
warning="；".join(warnings) or None,
|
||
windows=tuple(catalog_windows(hub.managed_windows())),
|
||
tab_closed=tab_closed,
|
||
)
|
||
|
||
def stop_session(
|
||
window_id: str,
|
||
target: str,
|
||
session: str,
|
||
) -> StopSessionOutcome:
|
||
stop_remote_session(ssh_bin, RemoteEntry(target), session)
|
||
warning: str | None = None
|
||
tab_closed = False
|
||
try:
|
||
hub.dismiss_window(window_id)
|
||
tab_closed = True
|
||
except HubError as error:
|
||
warning = f"Session 已停止，但 Tab 未关闭：{error}"
|
||
updated = catalog_windows(hub.managed_windows())
|
||
if not tab_closed:
|
||
updated = [
|
||
replace(window, dead=True)
|
||
if window.window_id == window_id
|
||
else window
|
||
for window in updated
|
||
]
|
||
return StopSessionOutcome(tuple(updated), tab_closed, warning)
|
||
|
||
try:
|
||
selector = SessionSelector(
|
||
catalog_windows(windows),
|
||
remote_registry,
|
||
ssh_bin,
|
||
local_sessions=(endpoint.session for endpoint in local_report.endpoints),
|
||
aliases=hub.aliases,
|
||
rename_handler=rename_session,
|
||
delete_session_handler=delete_session,
|
||
stop_session_handler=stop_session,
|
||
)
|
||
startup_warnings = [*local_report.warnings]
|
||
if registry_warning:
|
||
startup_warnings.append(registry_warning)
|
||
if startup_warnings:
|
||
selector.message = "; ".join(startup_warnings)
|
||
result = selector.run()
|
||
except (RegistryError, SelectorError) as error:
|
||
raise HubError(str(error)) from error
|
||
if result is None:
|
||
return 0
|
||
if result.kind == "rename":
|
||
assert result.endpoint_key is not None and result.alias is not None
|
||
rename_session(result.endpoint_key, result.alias)
|
||
window = hub.window_for_key(result.endpoint_key)
|
||
assert window is not None
|
||
action = "renamed" if result.alias else "restored the default name for"
|
||
hub.display_message(f"herdr-hub: {action} {window.name}")
|
||
return 0
|
||
if result.kind == "close":
|
||
assert result.window_id is not None
|
||
hub.dismiss_window(result.window_id)
|
||
return 0
|
||
if result.kind == "window":
|
||
assert result.window_id is not None
|
||
if result.dead:
|
||
hub.reconnect(result.window_id)
|
||
hub.select_window(result.window_id)
|
||
return 0
|
||
if result.kind == "local":
|
||
assert result.session is not None
|
||
hub.activate_endpoint(Endpoint.local(result.session), resolve_herdr())
|
||
return 0
|
||
assert result.target is not None and result.session is not None
|
||
endpoint = Endpoint.remote(
|
||
result.target,
|
||
result.session,
|
||
result.remote_keybindings,
|
||
result.remote_transport,
|
||
)
|
||
try:
|
||
remote_registry.upsert(
|
||
RemoteEntry(
|
||
endpoint.target,
|
||
endpoint.remote_keybindings,
|
||
endpoint.remote_transport,
|
||
)
|
||
)
|
||
except RegistryError as error:
|
||
hub.display_message(f"herdr-hub: remote registry: {error}")
|
||
hub.activate_endpoint(endpoint, resolve_herdr())
|
||
return 0
|
||
|
||
|
||
def command_remote(arguments: argparse.Namespace) -> int:
|
||
registry = make_registry()
|
||
action = arguments.remote_command
|
||
if action == "list":
|
||
entries = registry.load()
|
||
if not entries:
|
||
print("no registered remote targets")
|
||
return 0
|
||
for entry in entries:
|
||
print(
|
||
f"{entry.target}\t{entry.remote_keybindings}\t{entry.remote_transport}"
|
||
)
|
||
return 0
|
||
if action == "add":
|
||
entry = RemoteEntry(
|
||
arguments.target,
|
||
arguments.remote_keybindings,
|
||
arguments.transport,
|
||
)
|
||
changed = registry.upsert(entry)
|
||
print(f"{'registered' if changed else 'already registered'} {entry.target}")
|
||
return 0
|
||
if action == "remove":
|
||
removed = registry.remove(arguments.target)
|
||
print(f"{'removed' if removed else 'not registered'} {arguments.target}")
|
||
return 0
|
||
raise HubError("use 'herdr-hub remote list|add|remove'")
|
||
|
||
|
||
def command_reconnect(arguments: argparse.Namespace) -> int:
|
||
hub = make_hub(arguments, load_aliases=False)
|
||
if not hub.exists():
|
||
raise HubError("herdr-hub is not running")
|
||
hub.reconnect(arguments.window)
|
||
if inside_socket(hub.socket_name):
|
||
hub.display_message("herdr-hub: reconnecting current session")
|
||
return 0
|
||
|
||
|
||
def command_close(arguments: argparse.Namespace) -> int:
|
||
hub = make_hub(arguments, load_aliases=False)
|
||
if not hub.exists():
|
||
raise HubError("herdr-hub is not running")
|
||
try:
|
||
window = hub.dismiss_window(arguments.window)
|
||
except HubError as error:
|
||
if inside_socket(hub.socket_name):
|
||
hub.display_message(f"herdr-hub: {error}")
|
||
return 1
|
||
raise
|
||
message = f"herdr-hub: closed {window.name}; Herdr Session is still running"
|
||
if inside_socket(hub.socket_name):
|
||
hub.display_message(message)
|
||
else:
|
||
print(message)
|
||
return 0
|
||
|
||
|
||
def command_stop(arguments: argparse.Namespace) -> int:
|
||
hub = make_hub(arguments, load_aliases=False)
|
||
if hub.exists():
|
||
hub.stop()
|
||
print("stopped herdr-hub; Herdr servers were left running")
|
||
else:
|
||
print("herdr-hub is not running")
|
||
return 0
|
||
|
||
|
||
def command_doctor(arguments: argparse.Namespace) -> int:
|
||
failures = 0
|
||
print(f"python: {sys.version.split()[0]} ({sys.executable})")
|
||
print(f"curses: {getattr(curses, 'version', 'available')}")
|
||
for binary in ["tmux", "herdr", "ssh"]:
|
||
override = os.environ.get(f"HERDR_HUB_{binary.upper()}_BIN")
|
||
path = shutil.which(binary) if override is None else override
|
||
if path is None:
|
||
print(f"{binary}: missing")
|
||
failures += 1
|
||
continue
|
||
version_args = [path, "-V"] if binary in {"tmux", "ssh"} else [path, "--version"]
|
||
result = subprocess.run(
|
||
version_args,
|
||
capture_output=True,
|
||
text=True,
|
||
check=False,
|
||
)
|
||
version = (result.stdout or result.stderr).strip()
|
||
print(f"{binary}: {version or path}")
|
||
registry = make_registry()
|
||
try:
|
||
registered = registry.load()
|
||
print(f"remotes: {len(registered)} registered ({registry.path})")
|
||
except RegistryError as error:
|
||
print(f"remotes: invalid ({error})")
|
||
failures += 1
|
||
alias_registry = make_alias_registry()
|
||
try:
|
||
aliases = alias_registry.load()
|
||
print(f"aliases: {len(aliases)} configured ({alias_registry.path})")
|
||
except AliasRegistryError as error:
|
||
print(f"aliases: invalid ({error})")
|
||
failures += 1
|
||
proc_root = configured_proc_root()
|
||
print(f"proc: {proc_root} ({'available' if proc_root.is_dir() else 'missing'})")
|
||
|
||
try:
|
||
herdr_bin = resolve_herdr()
|
||
except HubError as error:
|
||
print(f"discovery: unavailable ({error})")
|
||
return 1
|
||
report = discover_endpoints(herdr_bin, proc_root)
|
||
print(f"discovered: {len(report.endpoints)} endpoint(s)")
|
||
for endpoint in report.endpoints:
|
||
print(f" {endpoint.label} [{endpoint.key}]")
|
||
for warning in report.warnings:
|
||
print(f"warning: {warning}")
|
||
return failures
|
||
|
||
|
||
def build_parser() -> argparse.ArgumentParser:
|
||
parser = argparse.ArgumentParser(
|
||
prog="herdr-hub",
|
||
description="Switch local and live --remote Herdr sessions through a tmux top bar.",
|
||
)
|
||
subparsers = parser.add_subparsers(dest="command")
|
||
subparsers.add_parser("start", help="scan, create and attach to the hub")
|
||
subparsers.add_parser("refresh", help="add newly discovered sessions")
|
||
|
||
open_parser = subparsers.add_parser("open", help="open a session directly in the hub")
|
||
open_parser.add_argument("--remote", metavar="TARGET")
|
||
open_parser.add_argument("--session", default=DEFAULT_SESSION_NAME)
|
||
open_parser.add_argument(
|
||
"--remote-keybindings",
|
||
choices=["local", "server"],
|
||
default="local",
|
||
)
|
||
open_parser.add_argument(
|
||
"--transport",
|
||
choices=["native", "ssh"],
|
||
default="native",
|
||
help="remote attach transport; Windows targets require ssh",
|
||
)
|
||
|
||
subparsers.add_parser("choose", help="browse and open local or remote sessions")
|
||
remote_parser = subparsers.add_parser("remote", help="manage remembered remote targets")
|
||
remote_subparsers = remote_parser.add_subparsers(dest="remote_command")
|
||
remote_subparsers.add_parser("list", help="list remembered remote targets")
|
||
remote_add_parser = remote_subparsers.add_parser("add", help="remember a remote target")
|
||
remote_add_parser.add_argument("target", metavar="TARGET")
|
||
remote_add_parser.add_argument(
|
||
"--remote-keybindings",
|
||
choices=["local", "server"],
|
||
default="local",
|
||
)
|
||
remote_add_parser.add_argument(
|
||
"--transport",
|
||
choices=["native", "ssh"],
|
||
default="native",
|
||
help="preferred attach transport; use ssh for Windows targets",
|
||
)
|
||
remote_remove_parser = remote_subparsers.add_parser(
|
||
"remove", help="forget a remote target"
|
||
)
|
||
remote_remove_parser.add_argument("target", metavar="TARGET")
|
||
reconnect_parser = subparsers.add_parser(
|
||
"reconnect", help="respawn a dead managed window"
|
||
)
|
||
reconnect_parser.add_argument("--window")
|
||
close_parser = subparsers.add_parser(
|
||
"close", help="hide a managed window without stopping its Herdr session"
|
||
)
|
||
close_parser.add_argument("--window")
|
||
subparsers.add_parser("doctor", help="check dependencies and discovery")
|
||
subparsers.add_parser("stop", help="stop the hub clients, leaving Herdr servers running")
|
||
return parser
|
||
|
||
|
||
def main(argv: Sequence[str] | None = None) -> int:
|
||
parser = build_parser()
|
||
arguments = parser.parse_args(argv)
|
||
command = arguments.command or "start"
|
||
handlers = {
|
||
"start": command_start,
|
||
"refresh": command_refresh,
|
||
"open": command_open,
|
||
"choose": command_choose,
|
||
"remote": command_remote,
|
||
"reconnect": command_reconnect,
|
||
"close": command_close,
|
||
"doctor": command_doctor,
|
||
"stop": command_stop,
|
||
}
|
||
try:
|
||
return handlers[command](arguments)
|
||
except (AliasRegistryError, HubError, RegistryError, ValueError) as error:
|
||
print(f"herdr-hub: {error}", file=sys.stderr)
|
||
return 1
|
||
except KeyboardInterrupt:
|
||
return 130
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main())
|