runtime-sdk 0.4.49__py3-none-win_amd64.whl

runtime_sdk/proxy.py

@@ -0,0 +1,600 @@
+from __future__ import annotations
+
+import contextlib
+import json
+import os
+import shlex
+import signal
+import socket
+import subprocess
+import sys
+import threading
+import time
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any, Callable
+
+from websockets.sync.client import connect as websocket_connect
+
+from runtime_sdk import computers, output as cli_output, terminal
+from runtime_sdk.auth import authenticated_client
+from runtime_sdk.client import DEFAULT_WARMUP_TIMEOUT, RuntimeAPIError
+from runtime_sdk.config import RuntimeConfig
+
+
+def _runtime_state_dir() -> Path:
+    root = os.environ.get("XDG_STATE_HOME")
+    if root:
+        return Path(root) / "runtime"
+    if sys.platform == "win32":
+        root = os.environ.get("LOCALAPPDATA")
+        if root:
+            return Path(root) / "Runtime"
+    return Path.home() / ".local" / "state" / "runtime"
+
+
+def _proxy_dir() -> Path:
+    return _runtime_state_dir() / "proxies"
+
+
+def _proxy_state_path(local_port: int) -> Path:
+    return _proxy_dir() / f"proxy-{local_port}.json"
+
+
+def _proxy_log_path(local_port: int) -> Path:
+    return _proxy_dir() / f"proxy-{local_port}.log"
+
+
+def _pid_running(pid: int) -> bool:
+    if pid <= 0:
+        return False
+    if sys.platform == "win32":
+        return _windows_process_started_at(pid) is not None
+    try:
+        os.kill(pid, 0)
+        return True
+    except OSError:
+        return False
+
+
+def _pid_command(pid: int) -> str:
+    if pid <= 0:
+        return ""
+    try:
+        return subprocess.check_output(
+            ["ps", "-o", "command=", "-p", str(pid)],
+            text=True,
+            stderr=subprocess.DEVNULL,
+        ).strip()
+    except (OSError, subprocess.SubprocessError):
+        return ""
+
+
+def _proxy_pid_matches(
+    pid: int, local_port: int, process_started_at: int | None = None
+) -> bool:
+    if sys.platform == "win32":
+        current_started_at = _windows_process_started_at(pid)
+        return current_started_at is not None and current_started_at == process_started_at
+    if not _pid_running(pid):
+        return False
+    command = _pid_command(pid)
+    if not command:
+        return False
+    with contextlib.suppress(ValueError):
+        argv = shlex.split(command)
+        if "_proxy_agent" not in argv:
+            return False
+        if "--local-port" not in argv:
+            return False
+        index = argv.index("--local-port")
+        if index + 1 >= len(argv) or argv[index + 1] != str(local_port):
+            return False
+        return "runtime_sdk.cli" in argv or any(
+            part.endswith("runtime_sdk.cli") for part in argv
+        )
+    return False
+
+
+def _load_proxy_records(*, clean_stale: bool = False) -> list[dict[str, Any]]:
+    records: list[dict[str, Any]] = []
+    directory = _proxy_dir()
+    if not directory.exists():
+        return records
+    for path in sorted(directory.glob("proxy-*.json")):
+        try:
+            record = json.loads(path.read_text())
+        except (OSError, ValueError):
+            continue
+        if not isinstance(record, dict):
+            continue
+        record["state_path"] = str(path)
+        pid = int(record.get("pid") or 0)
+        local_port = int(record.get("local_port") or 0)
+        process_started_at = int(record.get("process_started_at") or 0) or None
+        record["status"] = (
+            "running"
+            if _proxy_pid_matches(pid, local_port, process_started_at)
+            else "stopped"
+        )
+        if clean_stale and record["status"] != "running":
+            with contextlib.suppress(OSError):
+                path.unlink()
+            continue
+        records.append(record)
+    return records
+
+
+def _parse_proxy_port_spec(spec: str) -> tuple[int, int]:
+    text = str(spec or "").strip()
+    if not text:
+        raise RuntimeAPIError("port mapping is required")
+    if ":" in text:
+        local_text, remote_text = text.split(":", 1)
+    else:
+        local_text, remote_text = text, text
+    try:
+        local_port = int(local_text)
+        remote_port = int(remote_text)
+    except ValueError as exc:
+        raise RuntimeAPIError(f"invalid port mapping: {text}") from exc
+    for port in (local_port, remote_port):
+        if port <= 0 or port > 65535:
+            raise RuntimeAPIError("port must be between 1 and 65535")
+    return local_port, remote_port
+
+
+def _spawn_proxy_agent(
+    base_url: str,
+    computer_id: str,
+    display_name: str,
+    local_port: int,
+    remote_port: int,
+) -> int:
+    proxy_dir = _proxy_dir()
+    proxy_dir.mkdir(parents=True, exist_ok=True)
+    log_path = _proxy_log_path(local_port)
+    log_file = log_path.open("a", encoding="utf-8")
+    try:
+        popen_options: dict[str, Any] = {
+            "stdin": subprocess.DEVNULL,
+            "stdout": log_file,
+            "stderr": log_file,
+            "close_fds": True,
+            "env": os.environ.copy(),
+        }
+        if sys.platform == "win32":
+            popen_options["creationflags"] = (
+                subprocess.CREATE_NEW_PROCESS_GROUP | subprocess.DETACHED_PROCESS
+            )
+        else:
+            popen_options["start_new_session"] = True
+        proc = subprocess.Popen(
+            [
+                sys.executable,
+                "-m",
+                "runtime_sdk.cli",
+                "_proxy_agent",
+                "--base-url",
+                base_url,
+                "--computer-id",
+                computer_id,
+                "--display-name",
+                display_name,
+                "--local-port",
+                str(local_port),
+                "--remote-port",
+                str(remote_port),
+            ],
+            **popen_options,
+        )
+    finally:
+        log_file.close()
+
+    deadline = time.time() + max(4.0, DEFAULT_WARMUP_TIMEOUT + 2.0)
+    state_path = _proxy_state_path(local_port)
+    while time.time() < deadline:
+        if state_path.exists():
+            return proc.pid
+        code = proc.poll()
+        if code is not None:
+            break
+        time.sleep(0.1)
+
+    message = f"proxy failed to start on localhost:{local_port}"
+    with contextlib.suppress(OSError):
+        tail = log_path.read_text(encoding="utf-8").strip()
+        if tail:
+            message = tail.splitlines()[-1]
+    raise RuntimeAPIError(message)
+
+
+def _remove_proxy_state(local_port: int) -> None:
+    with contextlib.suppress(OSError):
+        _proxy_state_path(local_port).unlink()
+
+
+def _stop_proxy_ports(
+    local_ports: list[int], *, stop_all: bool = False
+) -> list[dict[str, Any]]:
+    records = _load_proxy_records(clean_stale=False)
+    stopped: list[dict[str, Any]] = []
+    wanted = set(local_ports)
+    for record in records:
+        port = int(record.get("local_port") or 0)
+        if not stop_all and port not in wanted:
+            continue
+        pid = int(record.get("pid") or 0)
+        process_started_at = int(record.get("process_started_at") or 0) or None
+        if _proxy_pid_matches(pid, port, process_started_at):
+            with contextlib.suppress(OSError):
+                _terminate_process(pid)
+        _remove_proxy_state(port)
+        record["status"] = "stopped"
+        stopped.append(record)
+    return stopped
+
+
+def start_temporary_proxy(
+    config: RuntimeConfig, computer_id: str, remote_port: int
+) -> tuple[int, Callable[[], None]]:
+    _wait_for_proxy_preflight(config, computer_id, remote_port)
+
+    listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+    try:
+        listener.bind(("127.0.0.1", 0))
+        listener.listen()
+        listener.settimeout(1.0)
+    except OSError as exc:
+        with contextlib.suppress(OSError):
+            listener.close()
+        raise RuntimeAPIError(f"proxy listen failed: {exc}") from exc
+
+    local_port = int(listener.getsockname()[1])
+    stop = threading.Event()
+
+    def serve() -> None:
+        try:
+            while not stop.is_set():
+                try:
+                    conn, _ = listener.accept()
+                except socket.timeout:
+                    continue
+                except OSError:
+                    if stop.is_set():
+                        break
+                    continue
+                thread = threading.Thread(
+                    target=_handle_proxy_client,
+                    args=(config, computer_id, remote_port, conn),
+                    daemon=True,
+                )
+                thread.start()
+        finally:
+            with contextlib.suppress(OSError):
+                listener.close()
+
+    thread = threading.Thread(target=serve, daemon=True)
+    thread.start()
+
+    def cleanup() -> None:
+        stop.set()
+        with contextlib.suppress(OSError):
+            listener.close()
+        thread.join(timeout=2)
+
+    return local_port, cleanup
+
+
+def _wait_for_proxy_preflight(
+    config: RuntimeConfig, computer_id: str, remote_port: int
+) -> None:
+    deadline = time.monotonic() + 20
+    last_error: Exception | None = None
+    while time.monotonic() < deadline:
+        try:
+            _proxy_preflight(config, computer_id, remote_port)
+            return
+        except Exception as exc:
+            last_error = exc
+            time.sleep(0.5)
+    raise RuntimeAPIError(
+        f"proxy preflight failed for {computer_id}:{remote_port}: {last_error}"
+    )
+
+
+def handle_proxy_start(
+    config: RuntimeConfig, computer_id: str | None, specs: list[str] | None
+) -> int:
+    if (err := cli_output.require_auth(config)) is not None:
+        return err
+    if not computer_id:
+        return cli_output.report_error("computer id is required")
+    if not specs:
+        return cli_output.report_error("at least one port mapping is required")
+
+    client = authenticated_client(config)
+    resolved_id = computers.resolve_computer_ref(client, computer_id)
+    existing_records = _load_proxy_records(clean_stale=True)
+    existing_ports = {int(r.get("local_port") or 0) for r in existing_records}
+    mappings: list[tuple[int, int]] = []
+    seen_ports: set[int] = set()
+    for spec in specs:
+        local_port, remote_port = _parse_proxy_port_spec(spec)
+        if local_port in existing_ports:
+            return cli_output.report_error(
+                f"local port {local_port} is already in use by a runtime proxy"
+            )
+        if local_port in seen_ports:
+            return cli_output.report_error(
+                f"local port {local_port} was specified more than once"
+            )
+        seen_ports.add(local_port)
+        mappings.append((local_port, remote_port))
+
+    started: list[dict[str, Any]] = []
+    started_ports: list[int] = []
+    for local_port, remote_port in mappings:
+        try:
+            _spawn_proxy_agent(
+                config.base_url,
+                resolved_id,
+                computer_id,
+                local_port,
+                remote_port,
+            )
+        except RuntimeAPIError as exc:
+            if started_ports:
+                _stop_proxy_ports(started_ports, stop_all=False)
+            return cli_output.report_error(str(exc), status_code=exc.status_code)
+        except Exception as exc:
+            if started_ports:
+                _stop_proxy_ports(started_ports, stop_all=False)
+            return cli_output.report_error(str(exc))
+        started_ports.append(local_port)
+        started.append(
+            {
+                "computer_id": resolved_id,
+                "computer": computer_id,
+                "local_port": local_port,
+                "remote_port": remote_port,
+            }
+        )
+
+    return cli_output.report_success({"proxies": started})
+
+
+def handle_proxy_ls(config: RuntimeConfig) -> int:
+    proxies = _load_proxy_records(clean_stale=True)
+    return cli_output.report_success({"proxies": proxies})
+
+
+def handle_proxy_stop(
+    config: RuntimeConfig, local_ports: list[int] | None, *, stop_all: bool = False
+) -> int:
+    if not stop_all and not local_ports:
+        return cli_output.report_error("pass one or more local ports, or use --all")
+
+    stopped = _stop_proxy_ports(local_ports or [], stop_all=stop_all)
+
+    return cli_output.report_success({"proxies": stopped})
+
+
+def _proxy_preflight(
+    config: RuntimeConfig, computer_id: str, remote_port: int
+) -> None:
+    client = authenticated_client(config)
+    ticket_payload = client.create_proxy_ticket(computer_id, remote_port)
+    ws_url = terminal.resolve_ws_url(config.base_url, ticket_payload.get("ws_url"))
+    if not ws_url:
+        raise RuntimeAPIError("proxy endpoint did not return a websocket url")
+    with websocket_connect(
+        ws_url, open_timeout=DEFAULT_WARMUP_TIMEOUT, close_timeout=1, max_size=None
+    ):
+        return
+
+
+def run_proxy_agent(
+    config: RuntimeConfig,
+    computer_id: str,
+    display_name: str,
+    local_port: int,
+    remote_port: int,
+) -> int:
+    listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+    try:
+        listener.bind(("127.0.0.1", local_port))
+        listener.listen()
+        listener.settimeout(1.0)
+    except OSError as exc:
+        print(f"proxy listen failed on localhost:{local_port}: {exc}", file=sys.stderr)
+        return 1
+
+    try:
+        _proxy_preflight(config, computer_id, remote_port)
+    except Exception as exc:
+        with contextlib.suppress(OSError):
+            listener.close()
+        print(
+            f"proxy preflight failed for localhost:{local_port} -> {computer_id}:{remote_port}: {exc}",
+            file=sys.stderr,
+        )
+        return 1
+
+    proxy_dir = _proxy_dir()
+    proxy_dir.mkdir(parents=True, exist_ok=True)
+    state = {
+        "pid": os.getpid(),
+        "computer_id": computer_id,
+        "display_name": display_name,
+        "local_port": local_port,
+        "remote_port": remote_port,
+        "started_at": datetime.now(timezone.utc).isoformat(),
+    }
+    if sys.platform == "win32":
+        state["process_started_at"] = _windows_process_started_at(os.getpid())
+    _proxy_state_path(local_port).write_text(
+        json.dumps(state, indent=2) + "\n", encoding="utf-8"
+    )
+
+    stop_event = threading.Event()
+
+    def handle_signal(_signum: int, _frame: Any) -> None:
+        stop_event.set()
+        with contextlib.suppress(OSError):
+            listener.close()
+
+    signal.signal(signal.SIGTERM, handle_signal)
+    signal.signal(signal.SIGINT, handle_signal)
+
+    try:
+        while not stop_event.is_set():
+            try:
+                conn, _ = listener.accept()
+            except socket.timeout:
+                continue
+            except OSError:
+                if stop_event.is_set():
+                    break
+                continue
+            thread = threading.Thread(
+                target=_handle_proxy_client,
+                args=(config, computer_id, remote_port, conn),
+                daemon=True,
+            )
+            thread.start()
+    finally:
+        with contextlib.suppress(OSError):
+            listener.close()
+        _remove_proxy_state(local_port)
+    return 0
+
+
+def _terminate_process(pid: int) -> None:
+    if sys.platform != "win32":
+        os.kill(pid, signal.SIGTERM)
+        return
+
+    import ctypes
+    from ctypes import wintypes
+
+    process_terminate = 0x0001
+    kernel32 = ctypes.WinDLL("kernel32", use_last_error=True)
+    kernel32.OpenProcess.argtypes = [wintypes.DWORD, wintypes.BOOL, wintypes.DWORD]
+    kernel32.OpenProcess.restype = wintypes.HANDLE
+    kernel32.TerminateProcess.argtypes = [wintypes.HANDLE, wintypes.UINT]
+    kernel32.TerminateProcess.restype = wintypes.BOOL
+    kernel32.CloseHandle.argtypes = [wintypes.HANDLE]
+    kernel32.CloseHandle.restype = wintypes.BOOL
+    handle = kernel32.OpenProcess(process_terminate, False, pid)
+    if not handle:
+        raise OSError(ctypes.get_last_error(), "failed to open process")
+    try:
+        if not kernel32.TerminateProcess(handle, 0):
+            raise OSError(ctypes.get_last_error(), "failed to terminate process")
+    finally:
+        kernel32.CloseHandle(handle)
+
+
+def _windows_process_started_at(pid: int) -> int | None:
+    import ctypes
+    from ctypes import wintypes
+
+    process_query_limited_information = 0x1000
+    kernel32 = ctypes.WinDLL("kernel32", use_last_error=True)
+    kernel32.OpenProcess.argtypes = [wintypes.DWORD, wintypes.BOOL, wintypes.DWORD]
+    kernel32.OpenProcess.restype = wintypes.HANDLE
+    kernel32.GetProcessTimes.argtypes = [
+        wintypes.HANDLE,
+        ctypes.POINTER(wintypes.FILETIME),
+        ctypes.POINTER(wintypes.FILETIME),
+        ctypes.POINTER(wintypes.FILETIME),
+        ctypes.POINTER(wintypes.FILETIME),
+    ]
+    kernel32.GetProcessTimes.restype = wintypes.BOOL
+    kernel32.CloseHandle.argtypes = [wintypes.HANDLE]
+    kernel32.CloseHandle.restype = wintypes.BOOL
+    handle = kernel32.OpenProcess(process_query_limited_information, False, pid)
+    if not handle:
+        return None
+    creation = wintypes.FILETIME()
+    exit_time = wintypes.FILETIME()
+    kernel_time = wintypes.FILETIME()
+    user_time = wintypes.FILETIME()
+    try:
+        if not kernel32.GetProcessTimes(
+            handle,
+            ctypes.byref(creation),
+            ctypes.byref(exit_time),
+            ctypes.byref(kernel_time),
+            ctypes.byref(user_time),
+        ):
+            return None
+        return (int(creation.dwHighDateTime) << 32) | int(creation.dwLowDateTime)
+    finally:
+        kernel32.CloseHandle(handle)
+
+
+def _handle_proxy_client(
+    config: RuntimeConfig, computer_id: str, remote_port: int, conn: socket.socket
+) -> None:
+    client = authenticated_client(config)
+    try:
+        ticket_payload = client.create_proxy_ticket(computer_id, remote_port)
+        ws_url = terminal.resolve_ws_url(config.base_url, ticket_payload.get("ws_url"))
+        if not ws_url:
+            raise RuntimeAPIError("proxy endpoint did not return a websocket url")
+        with websocket_connect(
+            ws_url, open_timeout=DEFAULT_WARMUP_TIMEOUT, close_timeout=1, max_size=None
+        ) as ws:
+            stop = threading.Event()
+
+            def local_to_ws() -> None:
+                try:
+                    while not stop.is_set():
+                        data = conn.recv(65536)
+                        if not data:
+                            break
+                        ws.send(data)
+                except Exception:
+                    pass
+                finally:
+                    stop.set()
+                    with contextlib.suppress(Exception):
+                        ws.close()
+
+            def ws_to_local() -> None:
+                try:
+                    while not stop.is_set():
+                        data = ws.recv()
+                        if data is None:
+                            break
+                        if isinstance(data, str):
+                            chunk = data.encode()
+                        else:
+                            chunk = data
+                        if chunk:
+                            conn.sendall(chunk)
+                except Exception:
+                    pass
+                finally:
+                    stop.set()
+                    with contextlib.suppress(OSError):
+                        conn.shutdown(socket.SHUT_RDWR)
+
+            t1 = threading.Thread(target=local_to_ws, daemon=True)
+            t2 = threading.Thread(target=ws_to_local, daemon=True)
+            t1.start()
+            t2.start()
+            t1.join()
+            t2.join()
+    except Exception as exc:
+        print(
+            f"proxy connection failed for localhost:{conn.getsockname()[1]} -> {computer_id}:{remote_port}: {exc}",
+            file=sys.stderr,
+        )
+    finally:
+        with contextlib.suppress(OSError):
+            conn.close()