sshmd 0.1.0__py3-none-any.whl

sshm/daemon.py

@@ -0,0 +1,289 @@
+"""sshmd — SSH session manager daemon."""
+
+from __future__ import annotations
+
+import logging
+import os
+import signal
+import sys
+import threading
+from collections.abc import Callable
+from typing import Any
+
+from . import protocol
+from .config import (
+    PortForward,
+    add_port_forward,
+    find_entry,
+    load_entries,
+    remove_host,
+    remove_port_forward,
+    rename_host,
+    set_enabled,
+)
+from .ipc import IPC_HOST, IPC_PORT, IpcServer, StreamingResponse
+from .process import ProcessManager
+from .state import log_file, new_token, pid_file, write_token
+
+log = logging.getLogger("sshm")
+
+WATCHDOG_INTERVAL = 5.0
+
+
+def _required(req: dict[str, Any], *keys: str) -> list[Any]:
+    values = [req.get(k) for k in keys]
+    if not all(values):
+        raise ValueError(f"Missing {' or '.join(keys)}")
+    return values
+
+
+class Daemon:
+    def __init__(self) -> None:
+        self.pm = ProcessManager()
+        self.token = new_token()
+        self.server = IpcServer(handler=self.handle_request, token=self.token)
+        self._shutdown_event = threading.Event()
+        self._handlers: dict[str, Callable[[dict[str, Any]], Any]] = {
+            protocol.CMD_STATUS: self._cmd_status,
+            protocol.CMD_LIST: self._cmd_list,
+            protocol.CMD_CONNECT: self._cmd_connect,
+            protocol.CMD_ATTACH: self._cmd_attach,
+            protocol.CMD_DETACH: self._cmd_detach,
+            protocol.CMD_RESIZE: self._cmd_resize,
+            protocol.CMD_DISCONNECT: self._cmd_disconnect,
+            protocol.CMD_PORT_ADD: self._cmd_port_add,
+            protocol.CMD_PORT_REMOVE: self._cmd_port_remove,
+            protocol.CMD_ENABLE: self._cmd_enable,
+            protocol.CMD_DISABLE: self._cmd_disable,
+            protocol.CMD_REMOVE: self._cmd_remove,
+            protocol.CMD_RENAME: self._cmd_rename,
+            protocol.CMD_SHUTDOWN: self._cmd_shutdown,
+        }
+
+    def handle_request(self, req: dict[str, Any]) -> dict[str, Any] | StreamingResponse:
+        cmd = req.get("cmd", "")
+        handler = self._handlers.get(cmd)
+        if handler is None:
+            return protocol.err(f"Unknown command: {cmd}")
+        try:
+            return handler(req)
+        except ValueError as e:
+            # Request validation errors — report to the client without traceback noise
+            return protocol.err(str(e))
+        except Exception as e:
+            log.exception("Error handling %s", cmd)
+            return protocol.err(str(e))
+
+    # --- Command handlers ---
+
+    def _cmd_status(self, req: dict[str, Any]):
+        return protocol.ok({"status": "running", "sessions": len(self.pm.get_sessions())})
+
+    def _cmd_list(self, req: dict[str, Any]):
+        alias = req.get("alias")
+        if alias:
+            return protocol.ok([s.to_dict() for s in self.pm.get_sessions(alias)])
+
+        result = []
+        for e in load_entries():
+            if e.alias in ("*", ""):
+                continue
+            active = self.pm.get_sessions(e.alias)
+            result.append({
+                "alias": e.alias,
+                "hostname": e.hostname,
+                "user": e.user,
+                "port": e.port,
+                "enabled": e.enabled,
+                "connections": len(active),
+                "attached": sum(1 for s in active if s.attached),
+                "port_forwards": [pf.to_str() for pf in e.port_forwards],
+            })
+        return protocol.ok(result)
+
+    def _cmd_connect(self, req: dict[str, Any]):
+        [alias] = _required(req, "alias")
+        if not find_entry(alias):
+            return protocol.err(f"Unknown alias: {alias}")
+        session = self.pm.connect(alias, req.get("name"))
+        return protocol.ok(session.to_dict())
+
+    def _cmd_attach(self, req: dict[str, Any]):
+        [alias] = _required(req, "alias")
+        if not find_entry(alias):
+            return protocol.err(f"Unknown alias: {alias}")
+        session = self.pm.attach(alias, req.get("name"), cli_pid=req.get("cli_pid"))
+        if not session:
+            return protocol.err(f"No available session for '{alias}'")
+
+        # attach() has now RESERVED the session (attached=True). If anything below
+        # raises (e.g. a bogus cols/rows), release the reservation so the session
+        # isn't leaked out of the attachable pool.
+        try:
+            # Size the remote terminal to the attaching client before streaming starts
+            cols, rows = req.get("cols"), req.get("rows")
+            if cols and rows:
+                session.set_winsize(int(cols), int(rows))
+
+            # The IPC server bridges this connection after sending the response
+            return StreamingResponse(
+                protocol.ok(session.to_dict()), session, cli_pid=req.get("cli_pid")
+            )
+        except Exception:
+            session.detach()
+            raise
+
+    def _cmd_resize(self, req: dict[str, Any]):
+        alias, name, cols, rows = _required(req, "alias", "name", "cols", "rows")
+        for s in self.pm.get_sessions(alias):
+            if s.name == name:
+                s.set_winsize(int(cols), int(rows))
+                return protocol.ok({"resized": [int(cols), int(rows)]})
+        return protocol.err(f"No session {alias}/{name}")
+
+    def _cmd_detach(self, req: dict[str, Any]):
+        alias, name = _required(req, "alias", "name")
+        # Detach is mostly handled by the bridge ending; explicit detach works too
+        for s in self.pm.get_sessions(alias):
+            if s.name == name and s.attached:
+                s.detach()
+        # For enabled aliases, keep an unattached session ready
+        entry = find_entry(alias)
+        if entry and entry.enabled:
+            self.pm.ensure_unattached(alias)
+        return protocol.ok({"detached": True})
+
+    def _cmd_disconnect(self, req: dict[str, Any]):
+        alias, name = _required(req, "alias", "name")
+        return protocol.ok({"disconnected": self.pm.disconnect(alias, name)})
+
+    def _cmd_port_add(self, req: dict[str, Any]):
+        alias, direction, rule = _required(req, "alias", "direction", "rule")
+        if direction == "D":
+            pf = PortForward.socks(int(rule))  # SOCKS proxy: rule is just the port
+        else:
+            pf = PortForward.parse_rule(rule, direction)
+        add_port_forward(alias, pf)
+        self._rebuild_alias_sessions(alias)
+        return protocol.ok({"added": pf.to_str()})
+
+    def _cmd_port_remove(self, req: dict[str, Any]):
+        alias, rule = _required(req, "alias", "rule")  # rule is already serialized
+        remove_port_forward(alias, rule)
+        self._rebuild_alias_sessions(alias)
+        return protocol.ok({"removed": rule})
+
+    def _cmd_enable(self, req: dict[str, Any]):
+        [alias] = _required(req, "alias")
+        set_enabled(alias, True)
+        self.pm.ensure_unattached(alias)
+        return protocol.ok({"enabled": alias})
+
+    def _cmd_disable(self, req: dict[str, Any]):
+        [alias] = _required(req, "alias")
+        set_enabled(alias, False)
+        return protocol.ok({"disabled": alias})
+
+    def _cmd_remove(self, req: dict[str, Any]):
+        [alias] = _required(req, "alias")
+        self.pm.disconnect_alias(alias)
+        remove_host(alias)
+        return protocol.ok({"removed": alias})
+
+    def _cmd_rename(self, req: dict[str, Any]):
+        old_alias, new_alias = _required(req, "alias", "new_alias")
+        if not find_entry(old_alias):
+            return protocol.err(f"Unknown alias: {old_alias}")
+        if find_entry(new_alias):
+            return protocol.err(f"Alias '{new_alias}' already exists")
+        # Sessions are keyed by alias; drop the old ones (an enabled host gets a
+        # fresh session under the new alias on the next watchdog tick).
+        self.pm.disconnect_alias(old_alias)
+        rename_host(old_alias, new_alias)
+        return protocol.ok({"renamed": new_alias})
+
+    def _cmd_shutdown(self, req: dict[str, Any]):
+        self._shutdown_event.set()
+        return protocol.ok({"shutting_down": True})
+
+    # --- Background maintenance ---
+
+    def _rebuild_alias_sessions(self, alias: str) -> None:
+        """Restart active sessions so they pick up config changes (e.g. forwards)."""
+        for s in self.pm.get_sessions(alias):
+            self.pm.rebuild_session(alias, s.name)
+
+    def _watchdog(self) -> None:
+        while not self._shutdown_event.is_set():
+            try:
+                self.pm.check_orphaned_attaches()
+                self.pm.check_health()
+                self._ensure_enabled_sessions()
+            except Exception:
+                log.exception("Watchdog error")
+            self._shutdown_event.wait(WATCHDOG_INTERVAL)
+
+    def _ensure_enabled_sessions(self) -> None:
+        for entry in load_entries():
+            if entry.enabled and entry.alias not in ("*", ""):
+                try:
+                    self.pm.ensure_unattached(entry.alias)
+                except Exception:
+                    # One failing alias must not block the others
+                    log.exception("Failed to ensure session for %s", entry.alias)
+
+    def run(self) -> None:
+        try:
+            self.server.start()
+        except OSError as e:
+            # Don't touch the live daemon's pid/token files if the port is taken
+            log.error(
+                "Cannot bind %s:%s (%s) — is another sshmd already running?",
+                IPC_HOST, IPC_PORT, e,
+            )
+            return
+
+        # Publish credentials only after the port is ours
+        write_token(self.token)
+        pf = pid_file()
+        pf.write_text(str(os.getpid()), encoding="utf-8")
+
+        log.info("sshmd started (PID %s)", os.getpid())
+        log.info("IPC server listening on %s:%s", IPC_HOST, IPC_PORT)
+
+        # First watchdog tick auto-connects enabled aliases
+        watchdog_thread = threading.Thread(target=self._watchdog, daemon=True)
+        watchdog_thread.start()
+
+        def handle_signal(sig, frame):
+            log.info("Received signal %s, shutting down...", sig)
+            self._shutdown_event.set()
+
+        signal.signal(signal.SIGTERM, handle_signal)
+        signal.signal(signal.SIGINT, handle_signal)
+
+        self._shutdown_event.wait()
+
+        log.info("Shutting down...")
+        self.server.stop()
+        self.pm.disconnect_all()
+        pf.unlink(missing_ok=True)
+        log.info("sshmd stopped")
+
+
+def main() -> None:
+    handlers: list[logging.Handler] = [logging.FileHandler(log_file(), encoding="utf-8")]
+    if sys.stderr is not None:  # absent under pythonw / detached start
+        handlers.append(logging.StreamHandler())
+
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s [%(name)s] %(levelname)s: %(message)s",
+        handlers=handlers,
+    )
+
+    Daemon().run()
+
+
+if __name__ == "__main__":
+    main()

sshm/ipc.py

@@ -0,0 +1,249 @@
+"""IPC client/server for sshm CLI <-> daemon communication over localhost TCP."""
+
+from __future__ import annotations
+
+import hmac
+import socket
+import subprocess
+import sys
+import threading
+import time
+from collections.abc import Callable
+from typing import Any
+
+from . import protocol
+from .procutil import daemon_interpreter, detached_popen_flags, pid_alive
+from .state import pid_file, read_token, resolve_port, token_file
+
+IPC_HOST = "127.0.0.1"
+IPC_PORT = resolve_port()
+BUFFER_SIZE = 65536
+REQUEST_TIMEOUT = 10.0
+
+
+def _recv_line(sock: socket.socket) -> tuple[bytes, bytes] | None:
+    """Read until a newline.
+
+    Returns ``(line, leftover)`` where ``line`` is the bytes up to (excluding)
+    the first newline and ``leftover`` is anything already received past it.
+    Returns None if the peer closed before sending a newline. The leftover
+    matters for the streaming handshake: the daemon writes the JSON response and
+    then immediately starts streaming session output, and on localhost both can
+    land in a single recv() — without preserving the tail those first bytes of
+    terminal output would be lost (or corrupt the JSON decode).
+    """
+    data = b""
+    while b"\n" not in data:
+        chunk = sock.recv(BUFFER_SIZE)
+        if not chunk:
+            return None
+        data += chunk
+    line, _, leftover = data.partition(b"\n")
+    return line, leftover
+
+
+# --- Server ---
+
+
+class StreamingResponse:
+    """Marks a response whose connection should become a raw I/O bridge to a session."""
+
+    def __init__(self, response: dict[str, Any], session, cli_pid: int | None = None):
+        self.response = response
+        self.session = session  # SshSession to bridge to
+        self.cli_pid = cli_pid
+
+
+Handler = Callable[[dict[str, Any]], "dict[str, Any] | StreamingResponse"]
+
+
+class IpcServer:
+    def __init__(self, handler: Handler, token: str):
+        self.handler = handler
+        self.token = token
+        self._sock: socket.socket | None = None
+        self._running = False
+        self._thread: threading.Thread | None = None
+
+    def start(self) -> None:
+        self._sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        if sys.platform == "win32":
+            # On Windows SO_REUSEADDR lets another process bind an actively
+            # listening port and hijack traffic; exclusive use makes it fail.
+            self._sock.setsockopt(socket.SOL_SOCKET, socket.SO_EXCLUSIVEADDRUSE, 1)
+        else:
+            self._sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+        self._sock.bind((IPC_HOST, IPC_PORT))
+        self._sock.listen(5)
+        self._sock.settimeout(1.0)
+        self._running = True
+        self._thread = threading.Thread(target=self._accept_loop, daemon=True)
+        self._thread.start()
+
+    def _accept_loop(self) -> None:
+        while self._running:
+            try:
+                conn, _ = self._sock.accept()
+            except TimeoutError:
+                continue
+            except OSError:
+                break
+            threading.Thread(target=self._handle_client, args=(conn,), daemon=True).start()
+
+    def _handle_client(self, conn: socket.socket) -> None:
+        try:
+            conn.settimeout(REQUEST_TIMEOUT)
+            received = _recv_line(conn)
+            if received is None:
+                return
+            data, _ = received  # a request is exactly one line; nothing trails it
+
+            request = protocol.decode(data)
+
+            token = request.get("token")
+            if not isinstance(token, str) or not hmac.compare_digest(token, self.token):
+                conn.sendall(protocol.encode(protocol.err("Invalid token")))
+                return
+
+            response = self.handler(request)
+
+            if isinstance(response, StreamingResponse):
+                # Send the JSON response, then turn this connection into a raw bridge.
+                # attach() already reserved the session (attached=True); if anything
+                # fails before bridge() starts its own finally-detach, release that
+                # reservation so the session isn't leaked out of the attachable pool.
+                try:
+                    conn.sendall(protocol.encode(response.response))
+                    conn.settimeout(None)
+                    response.session.bridge(conn, cli_pid=response.cli_pid)  # blocks until detach
+                except Exception:
+                    response.session.detach()
+                    raise
+            else:
+                conn.sendall(protocol.encode(response))
+
+        except Exception as e:
+            try:
+                conn.sendall(protocol.encode(protocol.err(str(e))))
+            except Exception:
+                pass
+        finally:
+            try:
+                conn.close()
+            except Exception:
+                pass
+
+    def stop(self) -> None:
+        self._running = False
+        if self._sock:
+            self._sock.close()
+        if self._thread:
+            self._thread.join(timeout=5)
+
+
+# --- Client ---
+
+
+class DaemonNotRunning(Exception):
+    pass
+
+
+def _open_request(cmd: str, **kwargs) -> socket.socket:
+    """Connect to the daemon and send a request; the response is not read yet."""
+    token = read_token()
+    if not token:
+        raise DaemonNotRunning("Daemon not running (no token file)")
+
+    request = protocol.make_request(cmd, token, **kwargs)
+
+    try:
+        sock = socket.create_connection((IPC_HOST, IPC_PORT), timeout=REQUEST_TIMEOUT)
+    except ConnectionRefusedError as e:
+        raise DaemonNotRunning("Cannot connect to sshmd (connection refused)") from e
+
+    sock.sendall(protocol.encode(request))
+    return sock
+
+
+def send_request(cmd: str, **kwargs) -> dict[str, Any]:
+    sock = _open_request(cmd, **kwargs)
+    try:
+        received = _recv_line(sock)
+    finally:
+        sock.close()
+    if received is None:
+        raise DaemonNotRunning("Connection closed")
+    data, _ = received
+    return protocol.decode(data)
+
+
+def connect_streaming(cmd: str, **kwargs) -> tuple[socket.socket, dict[str, Any], bytes]:
+    """Send a request; return (socket, response, leftover).
+
+    ``leftover`` is any streamed output that arrived bundled with the response
+    line and must be replayed to the terminal before reading further from the
+    socket. The socket stays open for streaming.
+    """
+    sock = _open_request(cmd, **kwargs)
+
+    received = _recv_line(sock)
+    if received is None:
+        sock.close()
+        raise DaemonNotRunning("Connection closed")
+
+    data, leftover = received
+    resp = protocol.decode(data)
+    if not resp.get("ok"):
+        sock.close()
+        return sock, resp, b""
+
+    sock.settimeout(None)  # no timeout while streaming
+    return sock, resp, leftover
+
+
+# --- Daemon lifecycle ---
+
+
+def is_daemon_running() -> bool:
+    pf = pid_file()
+    if not pf.exists():
+        return False
+    try:
+        pid = int(pf.read_text().strip())
+    except (ValueError, OSError):
+        # The daemon unlinks the pid file on shutdown; a read race means "not running"
+        return False
+    return pid_alive(pid)
+
+
+def ensure_daemon() -> None:
+    if is_daemon_running():
+        # The pid is alive, but confirm it's actually our daemon answering on the
+        # port — guards against PID reuse (a stale pid file pointing at an
+        # unrelated process) and a wedged daemon. If it doesn't respond, fall
+        # through and (re)spawn.
+        try:
+            if send_request(protocol.CMD_STATUS).get("ok"):
+                return
+        except Exception:
+            pass
+
+    subprocess.Popen(
+        [daemon_interpreter(), "-m", "sshm.daemon"],
+        stdin=subprocess.DEVNULL,
+        stdout=subprocess.DEVNULL,
+        stderr=subprocess.DEVNULL,
+        **detached_popen_flags(),
+    )
+
+    for _ in range(30):
+        time.sleep(0.2)
+        if token_file().exists():
+            try:
+                if send_request(protocol.CMD_STATUS).get("ok"):
+                    return
+            except Exception:
+                pass
+    raise RuntimeError(
+        f"Failed to start sshmd daemon (is port {IPC_PORT} taken by another process?)"
+    )