nullgate 1.1.0__py3-none-any.whl

nullgate/session.py

@@ -0,0 +1,509 @@
+"""Session lifecycle orchestration for Nullgate.
+
+Provides handlers for open, shut, inspect, enter, trace, cycle, and upgrade.
+"""
+
+from __future__ import annotations
+
+import argparse
+import importlib.metadata
+import os
+import shutil
+import subprocess
+import sys
+import time
+from pathlib import Path
+from typing import Any
+
+from nullgate import __version__, transports, wsroute
+from nullgate.account import current_username
+import nullgate.runtime as st
+
+DEFAULT_PORT = 4822
+DEFAULT_SLOT = 1
+DISTRIBUTION_URL = "https://shadowdocks.github.io/nullgate-dist"
+INSTALLER_URL = f"{DISTRIBUTION_URL}/install.sh"
+
+
+def _version() -> str:
+    try:
+        return importlib.metadata.version("nullgate")
+    except Exception:
+        return __version__
+
+
+def die(message: str) -> None:
+    raise SystemExit(message)
+
+
+def _resolve_launch_settings(
+    args: Any, saved: dict[str, Any] | None = None
+) -> tuple[Path, int, int]:
+    saved = saved or {}
+    root_arg = getattr(args, "workspace", None) or saved.get("root") or "."
+    port_arg = getattr(args, "port", None) or saved.get("port") or DEFAULT_PORT
+    slot_arg = getattr(args, "slot", None) or saved.get("slot") or DEFAULT_SLOT
+
+    port_text = str(port_arg)
+    if not port_text.isdigit() or not 1 <= int(port_text) <= 65535:
+        die("Port must be an integer between 1 and 65535")
+    port = int(port_text)
+
+    if not str(slot_arg).isdigit() or int(slot_arg) < 1:
+        die("Slot must be a positive integer")
+    slot = int(slot_arg)
+
+    try:
+        root = Path(root_arg).expanduser().resolve(strict=True)
+    except OSError:
+        die(f"Workspace root directory does not exist: {root_arg}")
+
+    return root, port, slot
+
+
+class _Color:
+    def __init__(self, enabled: bool) -> None:
+        self._on = enabled
+
+    def _wrap(self, code: str, text: str) -> str:
+        return f"\x1b[{code}m{text}\x1b[0m" if self._on else text
+
+    def green(self, text: str) -> str:
+        return self._wrap("32", text)
+
+    def red(self, text: str) -> str:
+        return self._wrap("31", text)
+
+    def dim(self, text: str) -> str:
+        return self._wrap("2", text)
+
+    def bold(self, text: str) -> str:
+        return self._wrap("1", text)
+
+
+def _colorizer() -> _Color:
+    return _Color(sys.stdout.isatty())
+
+
+def _proc_row(color: _Color, running: bool, pid: int | None) -> str:
+    if running:
+        return color.green(f"running   pid {pid}")
+    return color.red("stopped")
+
+
+def _print_url_row(color: _Color, known: bool, url: str) -> None:
+    del color
+    if known:
+        print(f"  url       {url}")
+    else:
+        print("  url       pending   " + _Color(sys.stdout.isatty()).dim("(nullgate trace transport -f)"))
+
+
+def _print_auth_rows(color: _Color, state: Path) -> None:
+    accept = st.manifest_get("accept")
+    allow = st.manifest_get("allow_tcp_forwarding")
+    if accept == "1":
+        print("  auth      none " + color.dim("(--accept); anyone can connect"))
+    else:
+        password_file = state / "password"
+        if password_file.is_file():
+            print(f"  auth      {current_username()} / {password_file.read_text()}")
+        keys = Path(os.path.expanduser("~/.ssh/authorized_keys"))
+        if keys.is_file() and keys.stat().st_size > 0:
+            print("  keys      enabled")
+        else:
+            print("  keys      disabled " + color.dim("(add keys to ~/.ssh/authorized_keys)"))
+    if allow == "1":
+        print("  forward   enabled")
+    else:
+        print("  forward   disabled " + color.dim("(--allow-tcp-forwarding)"))
+    confine = st.manifest_get("confine_sftp")
+    if not confine:
+        confine = st.read_settings().get("confine_sftp")
+    if confine == "1":
+        print("  sftp      confined " + color.dim("(--confine-sftp)"))
+    else:
+        print("  sftp      host")
+
+
+def _setup_line(host: str, proxy_command: str) -> str:
+    lines = [
+        "Host " + host,
+        "  ProxyCommand " + proxy_command,
+        "  StrictHostKeyChecking no",
+        "  UserKnownHostsFile /dev/null",
+        "  LogLevel ERROR",
+    ]
+    printed = " ".join(f"'{line}'" for line in lines)
+    if "*" in host:
+        grep_pattern = "^Host " + host.replace("*", "\\*")
+    else:
+        grep_pattern = f"^Host {host}$"
+    return (
+        f"grep -qs '{grep_pattern}' ~/.ssh/config || "
+        f"printf '%s\\n' {printed} >> ~/.ssh/config"
+    )
+
+
+def _fallback_line(ssh_user: str, host: str, proxy_command: str) -> str:
+    return (
+        f"ssh {ssh_user}@{host} -o 'ProxyCommand={proxy_command}' "
+        "-o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null "
+        "-o LogLevel=ERROR"
+    )
+
+
+def _tail_stderr(path: Path, count: int) -> None:
+    try:
+        lines = path.read_text(errors="replace").splitlines()
+    except OSError:
+        return
+    for line in lines[-count:]:
+        print(line, file=sys.stderr)
+
+
+def _tail_file(path: Path, follow: bool) -> None:
+    if not follow:
+        try:
+            lines = path.read_text(errors="replace").splitlines()[-100:]
+        except OSError:
+            lines = []
+        for line in lines:
+            print(line)
+        return
+    with open(path, "rb") as handle:
+        handle.seek(0, os.SEEK_END)
+        try:
+            while True:
+                line = handle.readline()
+                if line:
+                    sys.stdout.buffer.write(line)
+                    sys.stdout.buffer.flush()
+                else:
+                    time.sleep(0.2)
+        except KeyboardInterrupt:
+            return
+
+
+def _show_status() -> int:
+    color = _colorizer()
+    state = st.state_dir()
+    gateway_pid_file = state / "gateway.pid"
+    transport_pid_file = state / "transport.pid"
+    gateway_running = st.is_running(gateway_pid_file)
+    transport_running = st.is_running(transport_pid_file)
+    transport = st.manifest_get("transport") or transports.DEFAULT_TRANSPORT
+    target = transports.resolve_target(transport, state)
+
+    print(f"nullgate {_version()} \u00b7 {transport}")
+    print()
+    print(f"  gateway   {_proc_row(color, gateway_running, st.read_pid(gateway_pid_file))}")
+    print(f"  transport {_proc_row(color, transport_running, st.read_pid(transport_pid_file))}")
+    _print_url_row(color, target is not None, target.url if target else "")
+    _print_auth_rows(color, state)
+    if target is not None:
+        print()
+        print("  enter     " + color.bold(f"ssh {target.user or current_username()}@{target.host}"))
+        print("            " + color.dim("first time here: nullgate enter"))
+    return 0 if (gateway_running and transport_running) else 1
+
+
+def cmd_open(args: Any) -> int:
+    saved = st.read_settings()
+    transport = getattr(args, "transport", None)
+    if transport is None:
+        transport = saved.get("transport") or transports.DEFAULT_TRANSPORT
+    if transport not in transports.VALID_TRANSPORTS:
+        die(
+            f"Unrecognized transport: {transport} "
+            f"(supported: {', '.join(transports.VALID_TRANSPORTS)})"
+        )
+    hostname = getattr(args, "hostname", None) or saved.get("hostname") or ""
+    endpoint = getattr(args, "endpoint", None) or saved.get("endpoint") or ""
+    token = (
+        getattr(args, "token", None)
+        or os.environ.get("NULLGATE_CLOUDFLARED_TOKEN")
+        or saved.get("token")
+        or ""
+    )
+
+    root, port, slot = _resolve_launch_settings(args, saved)
+    state = st.setup_state()
+    gateway_pid_file = state / "gateway.pid"
+    transport_pid_file = state / "transport.pid"
+    gateway_log = state / "gateway.log"
+    transport_log = state / "transport.log"
+
+    if st.is_running(gateway_pid_file):
+        die("Gateway service is already running")
+    if st.is_running(transport_pid_file):
+        die("Transport service is already running")
+    for pid_file in (gateway_pid_file, transport_pid_file):
+        if not pid_file.is_dir():
+            pid_file.unlink(missing_ok=True)
+    if st.port_open(port):
+        die(f"Port {port} is already in use")
+
+    accept_arg = getattr(args, "accept", None)
+    accept = saved.get("accept") == "1" if accept_arg is None else bool(accept_arg)
+    forwarding_arg = getattr(args, "allow_tcp_forwarding", None)
+    allow_tcp_forwarding = (
+        saved.get("allow_tcp_forwarding") == "1"
+        if forwarding_arg is None
+        else bool(forwarding_arg)
+    )
+    confine_arg = getattr(args, "confine_sftp", None)
+    if confine_arg is not None:
+        confine_sftp = bool(confine_arg)
+    else:
+        raw_env = os.environ.get("NULLGATE_CONFINE_SFTP")
+        if raw_env is not None:
+            confine_sftp = raw_env.strip().lower() in ("1", "true", "yes")
+        else:
+            confine_sftp = saved.get("confine_sftp") == "1"
+
+    if accept:
+        (state / "password").unlink(missing_ok=True)
+        password = ""
+    elif transport == "upterm":
+        transports.prompt_authorized_key()
+        keys = Path(os.path.expanduser("~/.ssh/authorized_keys"))
+        if not keys.is_file() or not keys.stat().st_size:
+            die("The upterm transport requires an SSH authorized key or the --accept flag")
+        (state / "password").unlink(missing_ok=True)
+        password = ""
+    else:
+        transports.prompt_authorized_key()
+        password = st.generate_password()
+        (state / "password").write_text(password)
+
+    session_id = st.get_or_create_session()
+    gateway_log.write_text("")
+    transport_log.write_text("")
+
+    gateway_proc = transports.launch_gateway(
+        transport=transport,
+        root=root,
+        port=port,
+        password=password,
+        accept=accept,
+        allow_tcp_forwarding=allow_tcp_forwarding,
+        confine_sftp=confine_sftp,
+        state=state,
+    )
+    if not st.write_pid(gateway_pid_file, gateway_proc.pid):
+        st.kill_untracked(gateway_proc.pid)
+        die("Unable to track gateway process")
+
+    try:
+        ready = False
+        for _ in range(120):
+            if not st.is_running(gateway_pid_file):
+                break
+            if st.port_open(port):
+                ready = True
+                break
+            time.sleep(0.5)
+        if not ready:
+            st.stop_one(gateway_pid_file)
+            print("Gateway service failed to launch. Recent logs:", file=sys.stderr)
+            _tail_stderr(gateway_log, 100)
+            return 1
+
+        transport_proc = transports.launch_transport(
+            transport=transport,
+            port=port,
+            slot=slot,
+            session=session_id,
+            state=state,
+            endpoint=endpoint,
+            hostname=hostname,
+            token=token,
+            accept=accept,
+        )
+        if not st.write_pid(transport_pid_file, transport_proc.pid):
+            st.kill_untracked(transport_proc.pid)
+            st.stop_one(gateway_pid_file)
+            die("Unable to track transport process")
+
+        time.sleep(2)
+        if not st.is_running(transport_pid_file):
+            st.stop_one(gateway_pid_file)
+            transport_pid_file.unlink(missing_ok=True)
+            print("Transport service failed to launch. Recent logs:", file=sys.stderr)
+            _tail_stderr(transport_log, 100)
+            return 1
+
+        if transport == "srvus":
+            transports.wait_for_srvus_host(transport_pid_file, transport_log)
+        elif transport == "upterm":
+            if not transports.wait_for_upterm_session(transport_pid_file, state / "upterm.json"):
+                st.stop_one(transport_pid_file)
+                st.stop_one(gateway_pid_file)
+                print("Upterm session initialization timed out. Recent logs:", file=sys.stderr)
+                _tail_stderr(transport_log, 100)
+                return 1
+
+        manifest: dict[str, Any] = {
+            "transport": transport,
+            "port": str(port),
+            "accept": "1" if accept else "0",
+            "allow_tcp_forwarding": "1" if allow_tcp_forwarding else "0",
+            "confine_sftp": "1" if confine_sftp else "0",
+            "sftp_mode": "confined" if confine_sftp else "host",
+            "sftp_root": str(root),
+        }
+        if transport == "cloudflare":
+            manifest["endpoint"] = endpoint
+            manifest["session"] = session_id
+        elif transport == "cloudflared":
+            manifest["hostname"] = hostname
+        elif transport == "upterm":
+            manifest["endpoint"] = endpoint or wsroute.DEFAULT_ENDPOINT
+        st.write_manifest(manifest)
+
+        st.write_settings(
+            {
+                "transport": transport,
+                "root": str(root),
+                "port": str(port),
+                "slot": str(slot),
+                "hostname": hostname,
+                "endpoint": endpoint,
+                "token": token,
+                "accept": "1" if accept else "0",
+                "allow_tcp_forwarding": "1" if allow_tcp_forwarding else "0",
+                "confine_sftp": "1" if confine_sftp else "0",
+            }
+        )
+    except BaseException:
+        st.stop_one(transport_pid_file)
+        st.stop_one(gateway_pid_file)
+        raise
+
+    _show_status()
+    return 0
+
+
+def cmd_shut(args: Any = None) -> int:
+    del args
+    st.setup_state()
+    state = st.state_dir()
+    st.stop_one(state / "transport.pid")
+    st.stop_one(state / "gateway.pid")
+    (state / "manifest.json").unlink(missing_ok=True)
+    print("Nullgate stopped.")
+    return 0
+
+
+def cmd_inspect(args: Any = None) -> int:
+    del args
+    st.setup_state()
+    return _show_status()
+
+
+def cmd_enter(args: Any = None) -> int:
+    del args
+    st.setup_state()
+    state = st.state_dir()
+    transport = st.manifest_get("transport") or transports.DEFAULT_TRANSPORT
+    target = transports.resolve_target(transport, state)
+    if target is None:
+        die("No published address yet; check nullgate inspect.")
+    color = _colorizer()
+    ssh_user = target.user or current_username()
+    print("Run once on the machine you are connecting from:")
+    print()
+    print("  " + _setup_line(target.setup_host, target.proxy_command))
+    print()
+    print("Then, from now on:")
+    print()
+    print("  " + color.bold(f"ssh {ssh_user}@{target.host}"))
+    print()
+    print("Or, without changing any config:")
+    print()
+    print("  " + _fallback_line(ssh_user, target.host, target.proxy_command))
+    if target.note:
+        print()
+        print(color.dim(target.note))
+    return 0
+
+
+def cmd_trace(args: Any) -> int:
+    st.setup_state()
+    state = st.state_dir()
+    gateway_log = state / "gateway.log"
+    transport_log = state / "transport.log"
+    for log in (gateway_log, transport_log):
+        if not log.is_file():
+            log.write_text("")
+    if args.target == "all":
+        _tail_file(gateway_log, follow=args.f)
+        _tail_file(transport_log, follow=args.f)
+    elif args.target == "gateway":
+        _tail_file(gateway_log, follow=args.f)
+    else:
+        _tail_file(transport_log, follow=args.f)
+    return 0
+
+
+def cmd_cycle(args: Any = None) -> int:
+    del args
+    st.setup_state()
+    if not st.read_settings():
+        die("No saved configuration found to cycle; run open with desired options first.")
+    cmd_shut()
+    return cmd_open(
+        argparse.Namespace(
+            workspace=None,
+            port=None,
+            slot=None,
+            transport=None,
+            endpoint=None,
+            hostname=None,
+            token=None,
+            accept=None,
+            allow_tcp_forwarding=None,
+            confine_sftp=None,
+        )
+    )
+
+
+def cmd_upgrade(args: Any) -> int:
+    if not shutil.which("uv"):
+        die("uv is required; install it and re-run `nullgate upgrade`")
+    if shutil.which("curl") is None:
+        die("curl is required to upgrade")
+
+    current_version = _version()
+    requested_version = args.version or os.environ.get("NULLGATE_VERSION")
+    distribution_url = (
+        os.environ.get("NULLGATE_DIST_URL") or DISTRIBUTION_URL
+    ).rstrip("/")
+    install_url = os.environ.get("NULLGATE_INSTALL_URL") or (
+        f"{distribution_url}/install.sh"
+    )
+    separator = "&" if "?" in install_url else "?"
+    install_request = f"{install_url}{separator}ts={int(time.time())}"
+    try:
+        installer = subprocess.run(
+            ["curl", "-fsSL", install_request],
+            check=True,
+            capture_output=True,
+        )
+    except subprocess.CalledProcessError:
+        die(f"Failed to fetch installer from {install_url}")
+    destination = requested_version or "latest stable release"
+    print(f"nullgate {current_version} upgrading to {destination}")
+    environment = dict(os.environ)
+    environment.pop("NULLGATE_PACKAGE", None)
+    if requested_version:
+        environment["NULLGATE_VERSION"] = requested_version
+    else:
+        environment.pop("NULLGATE_VERSION", None)
+    completed = subprocess.run(
+        ["sh"],
+        input=installer.stdout,
+        env=environment,
+    )
+    return completed.returncode