susops 3.0.0rc3.dev1__py3-none-any.whl

susops/core/socat.py

@@ -0,0 +1,200 @@
+"""UDP port forwarding via socat over SSH ControlMaster.
+
+Architecture:
+  - Local UDP forward: socat EXEC approach — no SSH port forward slave needed.
+    One process: local socat pipes each UDP conversation through ControlMaster
+    to a remote socat instance (spawned per conversation via EXEC).
+  - Remote UDP forward: SSH -R slave + remote socat + local socat.
+    Three processes: an intermediate TCP port bridges the two socat instances.
+
+Error handling: FileNotFoundError when socat is not installed locally;
+subprocess exit errors when the remote host blocks command execution or
+lacks socat — both surface through the process manager log file.
+"""
+from __future__ import annotations
+
+import shlex
+from pathlib import Path
+
+from susops.core.config import Connection, PortForward
+from susops.core.ports import get_random_free_port
+from susops.core.process import ProcessManager
+from susops.core.ssh import socket_path
+
+__all__ = [
+    "UDP_PROCESS_PREFIX",
+    "start_udp_forward",
+    "stop_udp_forward",
+    "stop_all_udp_forwards_for_connection",
+    "is_udp_forward_running",
+]
+
+UDP_PROCESS_PREFIX = "susops-udp"
+
+
+def _fw_tag(fw: PortForward, direction: str) -> str:
+    """Return the identifying tag for a forward (tag field or direction-port)."""
+    return fw.tag or f"{direction}-{fw.src_port}"
+
+
+def _udp_process_name(conn_tag: str, fw_tag: str, suffix: str) -> str:
+    """Build a process name like susops-udp-<conn>-<fw_tag>-<suffix>."""
+    return f"{UDP_PROCESS_PREFIX}-{conn_tag}-{fw_tag}-{suffix}"
+
+
+def start_udp_forward(
+        conn: Connection,
+        fw: PortForward,
+        direction: str,
+        process_mgr: ProcessManager,
+        workspace: Path,
+) -> None:
+    """Start socat process(es) for a UDP port forward.
+
+    direction="local":  one local socat process using EXEC through ControlMaster.
+    direction="remote": SSH -R slave + remote socat (via SSH) + local socat.
+
+    Raises FileNotFoundError if socat is not installed locally.
+    Remote errors (socat missing, shell access blocked) surface as immediate
+    process exit — check the log file at workspace/logs/<name>.log.
+    """
+    sock = socket_path(conn.tag, workspace)
+    tag = _fw_tag(fw, direction)
+    log_dir = workspace / "logs"
+    log_dir.mkdir(parents=True, exist_ok=True)
+
+    if direction == "local":
+        _start_local_udp(conn, fw, sock, tag, process_mgr, log_dir)
+    else:
+        _start_remote_udp(conn, fw, sock, tag, process_mgr, log_dir)
+
+
+def _start_local_udp(
+        conn: Connection,
+        fw: PortForward,
+        sock: Path,
+        tag: str,
+        process_mgr: ProcessManager,
+        log_dir: Path,
+) -> None:
+    """Local UDP forward: socat EXEC piped through SSH ControlMaster.
+
+    Each UDP conversation forks one SSH session (multiplexed via ControlMaster).
+    -T15 closes idle forked children after 15 seconds.
+    """
+    name = _udp_process_name(conn.tag, tag, "lsocat")
+    ssh_exec = (
+        f"ssh -o ControlPath={shlex.quote(str(sock))} -T {conn.ssh_host} "
+        f"socat - UDP4-SENDTO:{fw.dst_addr}:{fw.dst_port}"
+    )
+    cmd = [
+        "socat",
+        "-T15",
+        f"UDP4-RECVFROM:{fw.src_port},reuseaddr,fork",
+        f"EXEC:'{ssh_exec}'",
+    ]
+    log_file = log_dir / f"{name}.log"
+    with open(log_file, "a") as log:
+        process_mgr.start(name, cmd, stdout=log, stderr=log)
+
+
+def _start_remote_udp(
+        conn: Connection,
+        fw: PortForward,
+        sock: Path,
+        tag: str,
+        process_mgr: ProcessManager,
+        log_dir: Path,
+) -> None:
+    """Remote UDP forward: local socat + SSH -R slave + remote socat (via SSH).
+
+    Startup order matters:
+      1. lsocat  — bind local:intermediate so it is ready when rsocat connects.
+      2. SSH -R  — request remote SSH server to bind remote:intermediate.
+      3. rsocat  — connect to remote:intermediate (with retry for binding lag).
+
+    The rsocat shell command retries up to 5 times (1 s apart) because the
+    SSH server may not finish binding the remote port before rsocat executes.
+    """
+    intermediate = get_random_free_port()
+
+    # 1. Local socat: TCP intermediate → UDP local service (start first)
+    lsocat_name = _udp_process_name(conn.tag, tag, "lsocat")
+    lsocat_cmd = [
+        "socat",
+        f"TCP4-LISTEN:{intermediate},reuseaddr,fork",
+        f"UDP4-SENDTO:{fw.dst_addr}:{fw.dst_port}",
+    ]
+    log_file = log_dir / f"{lsocat_name}.log"
+    with open(log_file, "a") as log:
+        process_mgr.start(lsocat_name, lsocat_cmd, stdout=log, stderr=log)
+
+    # 2. SSH -R slave: bind intermediate port on remote, forward to local
+    ssh_name = _udp_process_name(conn.tag, tag, "ssh")
+    ssh_cmd = [
+        "ssh", "-N", "-T",
+        "-o", f"ControlPath={sock}",
+        "-R", f"{intermediate}:localhost:{intermediate}",
+        conn.ssh_host,
+    ]
+    log_file = log_dir / f"{ssh_name}.log"
+    with open(log_file, "a") as log:
+        process_mgr.start(ssh_name, ssh_cmd, stdout=log, stderr=log)
+
+    # 3. Remote socat: UDP → TCP intermediate (retry loop for port binding lag)
+    rsocat_name = _udp_process_name(conn.tag, tag, "rsocat")
+    rsocat_shell = (
+        f"for _i in 1 2 3 4 5; do "
+        f"socat -T15 UDP4-RECVFROM:{fw.src_port},reuseaddr,fork "
+        f"TCP4:localhost:{intermediate} && break; sleep 1; done"
+    )
+    rsocat_cmd = [
+        "ssh", "-T",
+        "-o", f"ControlPath={sock}",
+        conn.ssh_host,
+        rsocat_shell,
+    ]
+    log_file = log_dir / f"{rsocat_name}.log"
+    with open(log_file, "a") as log:
+        process_mgr.start(rsocat_name, rsocat_cmd, stdout=log, stderr=log)
+
+
+def stop_udp_forward(
+        conn_tag: str,
+        fw_tag: str,
+        process_mgr: ProcessManager,
+) -> bool:
+    """Stop all socat/SSH processes for a single UDP forward.
+
+    Returns True if at least one process was stopped.
+    """
+    prefix = f"{UDP_PROCESS_PREFIX}-{conn_tag}-{fw_tag}-"
+    stopped_any = False
+    for name in list(process_mgr.status_all().keys()):
+        if name.startswith(prefix):
+            if process_mgr.stop(name):
+                stopped_any = True
+    return stopped_any
+
+
+def is_udp_forward_running(
+        conn_tag: str,
+        fw: PortForward,
+        direction: str,
+        process_mgr: ProcessManager,
+) -> bool:
+    """Return True if the main socat process (lsocat) for this UDP forward is alive."""
+    tag = _fw_tag(fw, direction)
+    name = _udp_process_name(conn_tag, tag, "lsocat")
+    return process_mgr.is_running(name)
+
+
+def stop_all_udp_forwards_for_connection(
+        conn_tag: str,
+        process_mgr: ProcessManager,
+) -> None:
+    """Stop all UDP socat processes for every forward on a connection."""
+    prefix = f"{UDP_PROCESS_PREFIX}-{conn_tag}-"
+    for name in list(process_mgr.status_all().keys()):
+        if name.startswith(prefix):
+            process_mgr.stop(name)

susops/core/ssh.py

@@ -0,0 +1,330 @@
+"""SSH tunnel management using ssh + ProcessManager.
+
+Architecture:
+  - One ControlMaster process per connection (manages the multiplexed socket).
+  - Each port forward (local, remote, share) is its own lightweight slave process
+    that attaches to the master via the Unix socket.
+  - Forwards can be added/removed without restarting the master.
+"""
+from __future__ import annotations
+
+import subprocess
+import time
+from pathlib import Path
+
+from susops.core.config import Connection, PortForward
+from susops.core.process import ProcessManager
+
+__all__ = [
+    "build_ssh_cmd",  # legacy alias — kept for CLI compat
+    "build_master_cmd",
+    "start_tunnel",  # starts master (forwards bundled in cmd)
+    "start_master",
+    "start_forward",  # ssh -O forward — registers live forward via socket
+    "cancel_forward",  # ssh -O cancel — releases master-held port
+    "stop_tunnel",  # stops master (and all its forwards)
+    "is_tunnel_running",
+    "is_socket_alive",
+    "find_master_pid",
+    "socket_path",
+    "test_ssh_connectivity",
+    "SSH_PROCESS_PREFIX",
+    "FWD_PROCESS_PREFIX",
+]
+
+SSH_PROCESS_PREFIX = "susops-ssh"
+FWD_PROCESS_PREFIX = "susops-fwd"
+_SOCKET_DIR = "sockets"
+
+
+def _master_name(tag: str) -> str:
+    return f"{SSH_PROCESS_PREFIX}-{tag}"
+
+
+def _forward_name(conn_tag: str, fw_tag: str) -> str:
+    return f"{FWD_PROCESS_PREFIX}-{conn_tag}-{fw_tag}"
+
+
+def socket_path(tag: str, workspace: Path) -> Path:
+    """Return the ControlMaster Unix socket path for a connection."""
+    return workspace / _SOCKET_DIR / f"{tag}.sock"
+
+
+def build_master_cmd(conn: Connection, sock: Path) -> list[str]:
+    """Build the ControlMaster command.
+
+    Intentionally contains no -L/-R forward flags — forwards are registered
+    live via ``start_forward`` (``ssh -O forward``) after the master socket is
+    ready. This keeps process arguments minimal and avoids exposing forward
+    destinations in ``ps aux``.
+
+    ControlPersist is intentionally omitted: with -N the process stays in the
+    foreground with a stable, trackable PID. Reconnection is handled by the
+    Python-side _ReconnectMonitor which restarts the master on dropout.
+    """
+    return [
+        "ssh",
+        "-N", "-T",
+        "-D", str(conn.socks_proxy_port),
+        "-o", "ControlMaster=yes",
+        "-o", f"ControlPath={sock}",
+        "-o", "ServerAliveInterval=5",
+        "-o", "ServerAliveCountMax=3",
+        conn.ssh_host,
+    ]
+
+
+# Legacy alias — keeps existing callers (CLI, old tests) working.
+def build_ssh_cmd(conn: Connection) -> list[str]:
+    """Legacy: build the monolithic SSH command (no ControlMaster).
+
+    Kept for backwards compatibility. New code should use build_master_cmd.
+    """
+    cmd: list[str] = ["ssh"]
+
+    cmd += [
+        "-N", "-T",
+        "-D", str(conn.socks_proxy_port),
+        "-o", "ExitOnForwardFailure=yes",
+        "-o", "ServerAliveInterval=30",
+        "-o", "ServerAliveCountMax=3",
+    ]
+
+    for fw in conn.forwards.local:
+        cmd += ["-L", f"{fw.src_addr}:{fw.src_port}:{fw.dst_addr}:{fw.dst_port}"]
+
+    for fw in conn.forwards.remote:
+        cmd += ["-R", f"{fw.src_addr}:{fw.src_port}:{fw.dst_addr}:{fw.dst_port}"]
+
+    cmd.append(conn.ssh_host)
+    return cmd
+
+
+def start_master(
+        conn: Connection,
+        process_mgr: ProcessManager,
+        workspace: Path,
+) -> int:
+    """Start the ControlMaster SSH process for a connection.
+
+    Returns the PID of the started process.
+    Raises ValueError if socks_proxy_port is 0.
+    Raises RuntimeError if the process fails to start.
+    """
+    if conn.socks_proxy_port == 0:
+        raise ValueError(
+            f"Connection '{conn.tag}' has no SOCKS port assigned. "
+            "Assign one before starting."
+        )
+
+    sock = socket_path(conn.tag, workspace)
+    sock.parent.mkdir(parents=True, exist_ok=True)
+
+    # Remove stale socket so the new master can take ownership.
+    # A live master is never reached here (facade checks is_tunnel_running +
+    # is_socket_alive before calling start_master), so a socket at this
+    # point is always left over from a dead master.
+    if sock.exists():
+        sock.unlink()
+
+    name = _master_name(conn.tag)
+    cmd = build_master_cmd(conn, sock)
+
+    log_file = workspace / "logs" / f"{name}.log"
+    log_file.parent.mkdir(parents=True, exist_ok=True)
+
+    with open(log_file, "a") as log:
+        pid = process_mgr.start(name, cmd, stdout=log, stderr=log)
+
+    return pid
+
+
+def start_forward(
+        conn: Connection,
+        fw: PortForward,
+        direction: str,
+        workspace: Path,
+) -> None:
+    """Register a TCP port forward with the running ControlMaster via ``ssh -O forward``.
+
+    Sends the forward request through the Unix socket — no SSH handshake,
+    no new TCP connection. The master holds the port until ``cancel_forward``
+    is called or the master exits.
+
+    Raises RuntimeError if the socket is not ready or the forward fails.
+    """
+    sock = socket_path(conn.tag, workspace)
+
+    for _ in range(100):
+        if sock.exists():
+            break
+        time.sleep(0.1)
+    else:
+        raise RuntimeError(f"ControlMaster socket {sock} not ready after 10 s")
+
+    flag = "-L" if direction == "local" else "-R"
+    fwd_spec = f"{fw.src_addr}:{fw.src_port}:{fw.dst_addr}:{fw.dst_port}"
+    result = subprocess.run(
+        ["ssh", "-O", "forward", "-o", f"ControlPath={sock}", flag, fwd_spec, conn.ssh_host],
+        capture_output=True,
+        timeout=5,
+    )
+    if result.returncode != 0:
+        stderr = result.stderr.decode(errors="replace").strip()
+        raise RuntimeError(f"ssh -O forward failed: {stderr}")
+
+
+def cancel_forward(
+        conn: Connection,
+        fw: PortForward,
+        direction: str,
+        workspace: Path,
+) -> None:
+    """Release a master-held TCP port forward via ``ssh -O cancel``.
+
+    Best-effort: silently ignores errors (master may already be gone).
+    """
+    sock = socket_path(conn.tag, workspace)
+    if not sock.exists():
+        return
+    flag = "-L" if direction == "local" else "-R"
+    fwd_spec = f"{fw.src_addr}:{fw.src_port}:{fw.dst_addr}:{fw.dst_port}"
+    try:
+        subprocess.run(
+            ["ssh", "-O", "cancel", "-o", f"ControlPath={sock}", flag, fwd_spec, conn.ssh_host],
+            capture_output=True,
+            timeout=5,
+        )
+    except (subprocess.TimeoutExpired, FileNotFoundError, OSError):
+        pass
+
+
+def start_tunnel(
+        conn: Connection,
+        process_mgr: ProcessManager,
+        workspace: Path,
+) -> int:
+    """Start the ControlMaster for a connection.
+
+    All enabled TCP forwards are bundled in the master command and python
+    restarts them automatically on reconnect. Returns the PID of the master.
+    """
+    return start_master(conn, process_mgr, workspace)
+
+
+def stop_tunnel(
+        tag: str,
+        process_mgr: ProcessManager,
+        workspace: Path | None = None,
+        ssh_host: str | None = None,
+) -> bool:
+    """Stop the ControlMaster for a connection.
+
+    Sends SIGTERM to the ssh master process via ProcessManager, then sends
+    ``ssh -O exit`` through the socket for a clean shutdown. workspace and
+    ssh_host are both required for the -O exit step; if either is absent that
+    step is skipped.
+
+    Returns True if the master was stopped, False if it wasn't running.
+    """
+    # Clean up any lingering forward slave PIDs (no-op with Approach B but safe
+    # to keep for processes left over from older susops versions).
+    prefix = f"{FWD_PROCESS_PREFIX}-{tag}-"
+    for name in list(process_mgr.status_all().keys()):
+        if name.startswith(prefix):
+            process_mgr.stop(name)
+
+    stopped = process_mgr.stop(_master_name(tag))
+
+    # Best-effort clean shutdown via the ControlMaster socket.
+    if workspace is not None and ssh_host is not None:
+        sock = socket_path(tag, workspace)
+        if sock.exists():
+            try:
+                subprocess.run(
+                    ["ssh", "-O", "exit", "-o", f"ControlPath={sock}", ssh_host],
+                    capture_output=True,
+                    timeout=3,
+                )
+            except (subprocess.TimeoutExpired, FileNotFoundError, OSError):
+                pass
+
+    return stopped
+
+
+def is_tunnel_running(tag: str, process_mgr: ProcessManager) -> bool:
+    """Return True if the ControlMaster SSH process for tag is currently running."""
+    return process_mgr.is_running(_master_name(tag))
+
+
+def find_master_pid(tag: str, workspace: Path) -> int | None:
+    """Find the PID of the ControlMaster by scanning /proc cmdline (Linux only).
+
+    Used to recover the PID when the PID file is stale but the socket is alive.
+    Returns None if the process cannot be found or /proc is unavailable.
+    """
+    sock_str = str(socket_path(tag, workspace))
+    proc = Path("/proc")
+    if not proc.exists():
+        return None
+    for pid_dir in proc.iterdir():
+        if not pid_dir.name.isdigit():
+            continue
+        try:
+            cmdline = (pid_dir / "cmdline").read_bytes().replace(b"\x00", b" ").decode(errors="replace")
+            if "ControlMaster=yes" in cmdline and sock_str in cmdline:
+                return int(pid_dir.name)
+        except OSError:
+            continue
+    return None
+
+
+def is_socket_alive(tag: str, workspace: Path) -> bool:
+    """Return True if the ControlMaster socket is responsive.
+
+    Uses `ssh -O check` to verify the master is healthy.
+    """
+    sock = socket_path(tag, workspace)
+    if not sock.exists():
+        return False
+    try:
+        result = subprocess.run(
+            ["ssh", "-O", "check", "-o", f"ControlPath={sock}", "placeholder"],
+            capture_output=True,
+            timeout=3,
+        )
+        return result.returncode == 0
+    except (subprocess.TimeoutExpired, FileNotFoundError):
+        return False
+
+
+def test_ssh_connectivity(ssh_host: str, timeout: int = 5) -> bool:
+    """Test SSH connectivity to a host without establishing a full tunnel.
+
+    Uses ssh with BatchMode=yes and a short timeout. Returns True if the
+    host is reachable (even if auth fails — we just need network reachability).
+
+    Note: A return code of 255 means connection failed. Code 1 means auth
+    failed (host is reachable but key not accepted). We treat both 0 and 1
+    as "reachable" since the SSH port is open.
+    """
+    cmd = [
+        "ssh",
+        "-q",
+        "-o", "BatchMode=yes",
+        "-o", f"ConnectTimeout={timeout}",
+        "-o", "StrictHostKeyChecking=no",
+        "-T",
+        ssh_host,
+        "true",
+    ]
+    try:
+        result = subprocess.run(
+            cmd,
+            capture_output=True,
+            timeout=timeout + 2,
+        )
+        # rc=0: success, rc=1: connected but auth failed, rc=255: connection error
+        return result.returncode != 255
+    except (subprocess.TimeoutExpired, FileNotFoundError):
+        return False

susops/core/ssh_config.py

@@ -0,0 +1,40 @@
+"""SSH config parser for hostname autocompletion."""
+from __future__ import annotations
+
+import re
+from pathlib import Path
+
+__all__ = ["get_ssh_hosts"]
+
+_SSH_CONFIG = Path.home() / ".ssh" / "config"
+
+
+def get_ssh_hosts(ssh_config_path: Path = _SSH_CONFIG) -> list[str]:
+    """Parse ~/.ssh/config and return all non-wildcard Host entries.
+
+    Returns a sorted list of hostnames suitable for autocomplete.
+    Skips wildcard entries like '*' or '*.example.com'.
+    """
+    if not ssh_config_path.exists():
+        return []
+
+    hosts: list[str] = []
+    try:
+        content = ssh_config_path.read_text(errors="replace")
+    except OSError:
+        return []
+
+    for line in content.splitlines():
+        line = line.strip()
+        # Match "Host <name>" lines (case-insensitive)
+        match = re.match(r'^[Hh]ost\s+(.+)$', line)
+        if not match:
+            continue
+        # A Host line can have multiple space-separated entries
+        for entry in match.group(1).split():
+            # Skip wildcards
+            if '*' in entry or '?' in entry:
+                continue
+            hosts.append(entry)
+
+    return sorted(set(hosts))