142 lines
5.2 KiB
Python
142 lines
5.2 KiB
Python
"""
|
|
Mosh integration — build mosh/ssh commands for attaching shells to nodes.
|
|
|
|
This manager doesn't *exec* the command (that's the operator's terminal
|
|
job). It builds the correct command string with all the right flags for:
|
|
- mosh (default — survives network roams + suspend)
|
|
- ssh (fallback — works everywhere mosh isn't installed)
|
|
- mosh via SSH bootstrap (mosh needs SSH for initial connection anyway)
|
|
|
|
The /api/nodes/{name}/shell endpoints use this to produce copy-to-clipboard
|
|
command strings for the UI "Shell" button.
|
|
"""
|
|
|
|
import shutil
|
|
import subprocess
|
|
from typing import Optional
|
|
|
|
|
|
class MoshManager:
|
|
|
|
def __init__(self):
|
|
self._mosh_path = shutil.which("mosh")
|
|
self._ssh_path = shutil.which("ssh")
|
|
|
|
@property
|
|
def available(self) -> bool:
|
|
return self._mosh_path is not None
|
|
|
|
@property
|
|
def ssh_available(self) -> bool:
|
|
return self._ssh_path is not None
|
|
|
|
def build_command(self, target: str,
|
|
port: Optional[int] = None,
|
|
ssh_key: Optional[str] = None,
|
|
method: str = "mosh",
|
|
mosh_port_range: Optional[str] = None,
|
|
extra_ssh_args: Optional[list] = None) -> str:
|
|
"""Build a shell command string for attaching to `target`.
|
|
|
|
Args:
|
|
target: "user@host" or just "host"
|
|
port: SSH port (default 22)
|
|
ssh_key: path to private key file
|
|
method: "mosh" or "ssh"
|
|
mosh_port_range: e.g. "60000:61000" (for firewalled environments)
|
|
extra_ssh_args: list of extra args to pass to ssh
|
|
|
|
Returns:
|
|
A shell command string ready to paste into a terminal.
|
|
"""
|
|
if method == "ssh":
|
|
return self._build_ssh(target, port, ssh_key, extra_ssh_args)
|
|
elif method == "mosh":
|
|
return self._build_mosh(target, port, ssh_key, mosh_port_range, extra_ssh_args)
|
|
else:
|
|
raise ValueError(f"unknown method: {method}")
|
|
|
|
def _build_ssh(self, target, port, ssh_key, extra_args) -> str:
|
|
parts = ["ssh"]
|
|
if port and port != 22:
|
|
parts.extend(["-p", str(port)])
|
|
if ssh_key:
|
|
parts.extend(["-i", ssh_key])
|
|
if extra_args:
|
|
parts.extend(extra_args)
|
|
parts.append(target)
|
|
return " ".join(parts)
|
|
|
|
def _build_mosh(self, target, port, ssh_key,
|
|
mosh_port_range, extra_ssh_args) -> str:
|
|
# mosh uses SSH for the bootstrap connection, so we pass SSH args
|
|
# via --ssh=...
|
|
ssh_cmd = "ssh"
|
|
ssh_flags = []
|
|
if port and port != 22:
|
|
ssh_flags.extend(["-p", str(port)])
|
|
if ssh_key:
|
|
ssh_flags.extend(["-i", ssh_key])
|
|
if extra_ssh_args:
|
|
ssh_flags.extend(extra_ssh_args)
|
|
|
|
if ssh_flags:
|
|
ssh_cmd = f"ssh {' '.join(ssh_flags)}"
|
|
|
|
parts = ["mosh"]
|
|
parts.extend(["--ssh", ssh_cmd])
|
|
if mosh_port_range:
|
|
parts.extend(["-p", mosh_port_range])
|
|
parts.append(target)
|
|
return " ".join(parts)
|
|
|
|
def connect(self, node) -> bool:
|
|
"""Spawn a mosh connection to `node` (blocks until mosh exits).
|
|
|
|
Returns True if mosh started successfully, False otherwise.
|
|
NOTE: this is interactive — only call from a real TTY."""
|
|
if not self.available:
|
|
# Fall back to ssh if mosh isn't installed
|
|
if self.ssh_available:
|
|
host = node.get("host") if isinstance(node, dict) else node
|
|
try:
|
|
subprocess.run(["ssh", host])
|
|
return True
|
|
except (subprocess.CalledProcessError, FileNotFoundError):
|
|
return False
|
|
return False
|
|
|
|
host = node.get("host") if isinstance(node, dict) else node
|
|
try:
|
|
subprocess.run([self._mosh_path, host])
|
|
return True
|
|
except (subprocess.CalledProcessError, FileNotFoundError):
|
|
return False
|
|
|
|
def test_connection(self, node) -> dict:
|
|
"""Test whether we can reach the node via SSH (mosh bootstrap).
|
|
|
|
Doesn't actually start mosh — just runs `ssh -o BatchMode=yes host true`
|
|
to verify connectivity. Returns a dict with `reachable` + `latency_ms`."""
|
|
import time
|
|
host = node.get("host") if isinstance(node, dict) else node
|
|
if not self.ssh_available:
|
|
return {"reachable": False, "reason": "ssh not installed"}
|
|
start = time.time()
|
|
try:
|
|
r = subprocess.run(
|
|
["ssh", "-o", "BatchMode=yes", "-o", "ConnectTimeout=5",
|
|
host, "true"],
|
|
capture_output=True, timeout=10,
|
|
)
|
|
latency = (time.time() - start) * 1000
|
|
if r.returncode == 0:
|
|
return {"reachable": True, "latency_ms": round(latency, 1)}
|
|
else:
|
|
err = r.stderr.decode().strip()[:200]
|
|
return {"reachable": False, "reason": err, "latency_ms": round(latency, 1)}
|
|
except subprocess.TimeoutExpired:
|
|
return {"reachable": False, "reason": "timeout", "latency_ms": 10000}
|
|
except Exception as e:
|
|
return {"reachable": False, "reason": str(e)}
|