multiproxy.py

#!/usr/bin/env python3
"""Resilient SOCKS5 proxy over the fastest usable SSH route."""

from __future__ import annotations

import argparse
import concurrent.futures
import contextlib
import itertools
import json
import os
import signal
import socket
import subprocess
import sys
import threading
import time
from pathlib import Path
from typing import Any

APP = "multiproxy"
CONFIG_DIR = Path(os.environ.get("XDG_CONFIG_HOME", Path.home() / ".config")) / APP
STATE_DIR = Path(os.environ.get("XDG_STATE_HOME", Path.home() / ".local/state")) / APP
CONFIG_FILE = CONFIG_DIR / "config.json"
PID_FILE = STATE_DIR / "multiproxy.pid"
STATUS_FILE = STATE_DIR / "status.json"
LOG_FILE = STATE_DIR / "multiproxy.log"
ROUTES_FILE = STATE_DIR / "routes.json"

DEFAULT_CONFIG: dict[str, Any] = {
    "port": 1080,
    "listen": "127.0.0.1",
    "test_url": "https://www.google.com/generate_204",
    "connect_timeout": 5,
    "test_timeout": 12,
    "health_interval": 20,
    "failure_threshold": 2,
    "ssh_config": "~/.ssh/config",
    "exits": ["cam", "srcf"],
    "max_jump_hops": 2,
    "discovery_workers": 8,
    "per_host_workers": 1,
    "probe_launch_interval": 0.35,
    "min_working_routes": 2,
}


def ensure_dirs() -> None:
    CONFIG_DIR.mkdir(parents=True, exist_ok=True)
    STATE_DIR.mkdir(parents=True, exist_ok=True)


def load_config(path: Path = CONFIG_FILE, create: bool = True) -> dict[str, Any]:
    ensure_dirs()
    if not path.exists():
        if not create:
            raise SystemExit(f"Config not found: {path}")
        path.parent.mkdir(parents=True, exist_ok=True)
        initial = dict(DEFAULT_CONFIG)
        initial["ssh_config"] = str((Path.home() / ".ssh/config").resolve())
        path.write_text(json.dumps(initial, indent=2) + "\n")
        try:
            path.chmod(0o600)
        except OSError:
            pass
    try:
        cfg = json.loads(path.read_text())
    except (OSError, json.JSONDecodeError) as exc:
        raise SystemExit(f"Cannot read {path}: {exc}") from exc
    if not cfg.get("exits") and not cfg.get("routes"):
        raise SystemExit(f"No exits configured in {path}")
    ssh_config = Path(os.path.expanduser(str(cfg.get("ssh_config", "~/.ssh/config"))))
    if not ssh_config.is_absolute():
        beside = path.parent / ssh_config
        ssh_config = beside if beside.exists() else Path.cwd() / ssh_config
    cfg["ssh_config"] = str(ssh_config.resolve())
    return cfg


def ssh_base(cfg: dict[str, Any], route: dict[str, Any]) -> list[str]:
    command = [
        "ssh", "-F", str(cfg["ssh_config"]), "-o", "BatchMode=yes", "-o", f"ConnectTimeout={cfg['connect_timeout']}",
        "-o", "ConnectionAttempts=1", "-o", "StrictHostKeyChecking=accept-new", "-o", "Compression=no",
    ]
    jumps = route.get("jumps", [])
    if jumps:
        command += ["-J", ",".join(jumps)]
    command += route.get("ssh_options", [])
    return command


def probe_route(cfg: dict[str, Any], route: dict[str, Any]) -> dict[str, Any]:
    started = time.monotonic()
    cmd = ssh_base(cfg, route) + [route["destination"], "true"]
    timeout = float(cfg["connect_timeout"]) * (len(route.get("jumps", [])) + 1) + 2
    try:
        result = subprocess.run(cmd, stdin=subprocess.DEVNULL, stdout=subprocess.DEVNULL,
                                stderr=subprocess.PIPE, text=True, timeout=timeout)
        elapsed = round((time.monotonic() - started) * 1000)
        return {"name": route["name"], "ok": result.returncode == 0, "latency_ms": elapsed,
                "error": result.stderr.strip().splitlines()[-1] if result.stderr.strip() else ""}
    except subprocess.TimeoutExpired:
        return {"name": route["name"], "ok": False, "latency_ms": None, "error": "timed out"}
    except OSError as exc:
        return {"name": route["name"], "ok": False, "latency_ms": None, "error": str(exc)}


class ProbeLimiter:
    """Pace SSH handshakes and protect each first-hop host from bursts."""

    def __init__(self, cfg: dict[str, Any], routes: list[dict[str, Any]]) -> None:
        per_host = max(1, int(cfg.get("per_host_workers", 1)))
        self._route_endpoints = {route["name"]: self.endpoint(cfg, route) for route in routes}
        endpoints = set(self._route_endpoints.values())
        self._host_slots = {endpoint: threading.BoundedSemaphore(per_host) for endpoint in endpoints}
        self._interval = max(0.0, float(cfg.get("probe_launch_interval", 0.35)))
        self._launch_lock = threading.Lock()
        self._next_launch = 0.0

    @staticmethod
    def endpoint(cfg: dict[str, Any], route: dict[str, Any]) -> str:
        jumps = route.get("jumps", [])
        endpoint = jumps[0] if jumps else route["destination"]
        seen: set[str] = set()
        # Follow configured ProxyJump dependencies to the actual first network
        # endpoint. This groups e.g. both `scm1` and `lxt-node3` under lxt-node3
        # when scm1 itself uses `ProxyJump lxt-node3`.
        while endpoint not in seen:
            seen.add(endpoint)
            try:
                resolved = subprocess.run(
                    ["ssh", "-F", str(cfg["ssh_config"]), "-G", endpoint],
                    stdin=subprocess.DEVNULL, stdout=subprocess.PIPE,
                    stderr=subprocess.DEVNULL, text=True, timeout=2,
                )
            except (OSError, subprocess.TimeoutExpired):
                break
            proxyjump = next((line.split(None, 1)[1] for line in resolved.stdout.splitlines()
                              if line.startswith("proxyjump ")), "none")
            if proxyjump == "none":
                break
            endpoint = proxyjump.split(",", 1)[0]
        return endpoint

    @contextlib.contextmanager
    def slot(self, route: dict[str, Any]):
        semaphore = self._host_slots[self._route_endpoints[route["name"]]]
        with semaphore:
            with self._launch_lock:
                now = time.monotonic()
                delay = max(0.0, self._next_launch - now)
                self._next_launch = max(now, self._next_launch) + self._interval
            if delay:
                time.sleep(delay)
            yield


def probe_routes(cfg: dict[str, Any], routes: list[dict[str, Any]]) -> list[dict[str, Any]]:
    workers = min(int(cfg.get("discovery_workers", 8)), len(routes)) or 1
    limiter = ProbeLimiter(cfg, routes)

    def limited_probe(route: dict[str, Any]) -> dict[str, Any]:
        with limiter.slot(route):
            return probe_route(cfg, route)

    with concurrent.futures.ThreadPoolExecutor(max_workers=workers) as pool:
        return list(pool.map(limited_probe, routes))


def ssh_hosts(path: Path, seen: set[Path] | None = None) -> list[str]:
    """Return concrete Host aliases, recursively following OpenSSH Includes."""
    hosts: list[str] = []
    seen = seen or set()
    path = path.expanduser().resolve()
    if path in seen:
        return hosts
    seen.add(path)
    try:
        lines = path.read_text().splitlines()
    except OSError as exc:
        raise SystemExit(f"Cannot read SSH config {path}: {exc}") from exc
    for raw in lines:
        fields = raw.split("#", 1)[0].strip().split()
        if fields and fields[0].lower() == "include":
            for pattern in fields[1:]:
                include = Path(os.path.expanduser(pattern))
                if not include.is_absolute():
                    include = path.parent / include
                for included_path in sorted(include.parent.glob(include.name)):
                    for host in ssh_hosts(included_path, seen):
                        if host not in hosts:
                            hosts.append(host)
        elif fields and fields[0].lower() == "host":
            for host in fields[1:]:
                if not any(mark in host for mark in "*!?") and host not in hosts:
                    hosts.append(host)
    return hosts


def route_name(destination: str, jumps: tuple[str, ...] | list[str]) -> str:
    return destination if not jumps else f"{destination} via {' -> '.join(jumps)}"


def discover_candidates(cfg: dict[str, Any]) -> list[dict[str, Any]]:
    if cfg.get("routes") and not cfg.get("exits"):
        return cfg["routes"]  # backwards compatibility
    exits = list(cfg["exits"])
    ignored = set(exits) | {"github.com"}
    jumps = [host for host in ssh_hosts(Path(cfg["ssh_config"])) if host not in ignored]
    candidates: list[dict[str, Any]] = []
    for destination in exits:
        for count in range(int(cfg.get("max_jump_hops", 2)) + 1):
            for chain in itertools.permutations(jumps, count):
                candidates.append({"name": route_name(destination, chain),
                                   "destination": destination, "jumps": list(chain)})
    return candidates


def load_cached_routes() -> list[dict[str, Any]]:
    try:
        return json.loads(ROUTES_FILE.read_text()).get("routes", [])
    except (OSError, json.JSONDecodeError, AttributeError):
        return []


def save_cached_routes(routes: list[dict[str, Any]], results: list[dict[str, Any]]) -> None:
    ensure_dirs()
    latency = {result["name"]: result["latency_ms"] for result in results if result["ok"]}
    remembered = [{**route, "latency_ms": latency[route["name"]]}
                  for route in routes if route["name"] in latency]
    temp = ROUTES_FILE.with_suffix(".tmp")
    temp.write_text(json.dumps({
        "discovered_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "routes": remembered,
    }, indent=2) + "\n")
    temp.replace(ROUTES_FILE)


def rank_routes(cfg: dict[str, Any], discover: bool = False) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
    routes = discover_candidates(cfg) if discover else load_cached_routes()
    if not routes:
        routes = discover_candidates(cfg)
        discover = True
    results: list[dict[str, Any]] = []
    if discover:
        # Lower-hop routes are normally faster. Probe each hop-count tier in
        # parallel and only expand the search when it has not found enough
        # failover choices.
        minimum = int(cfg.get("min_working_routes", 2))
        hop_counts = sorted({len(route.get("jumps", [])) for route in routes})
        for hop_count in hop_counts:
            tier = [route for route in routes if len(route.get("jumps", [])) == hop_count]
            results.extend(probe_routes(cfg, tier))
            if sum(result["ok"] for result in results) >= minimum:
                break
    else:
        results = probe_routes(cfg, routes)
    by_name = {route["name"]: route for route in routes}
    usable = sorted((result for result in results if result["ok"]),
                    key=lambda result: result["latency_ms"])
    ranked = [by_name[result["name"]] for result in usable]
    if discover:
        save_cached_routes(ranked, results)
    return ranked, results


def curl_test(cfg: dict[str, Any], quiet: bool = False) -> bool:
    proxy = f"socks5h://{cfg['listen']}:{cfg['port']}"
    cmd = ["curl", "--silent", "--show-error", "--output", "/dev/null", "--fail",
           "--max-time", str(cfg["test_timeout"]), "--proxy", proxy, cfg["test_url"]]
    result = subprocess.run(cmd, stdout=subprocess.DEVNULL,
                            stderr=subprocess.DEVNULL if quiet else subprocess.PIPE, text=True)
    if not quiet and result.returncode:
        print((result.stderr or "proxy test failed").strip(), file=sys.stderr)
    return result.returncode == 0


def tunnel_command(cfg: dict[str, Any], route: dict[str, Any]) -> list[str]:
    return ssh_base(cfg, route) + [
        "-N", "-T", "-D", f"{cfg['listen']}:{cfg['port']}",
        "-o", "ExitOnForwardFailure=yes", "-o", "ServerAliveInterval=15",
        "-o", "ServerAliveCountMax=2", "-o", "TCPKeepAlive=yes", route["destination"],
    ]


def write_status(**values: Any) -> None:
    ensure_dirs()
    values["updated_at"] = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
    temp = STATUS_FILE.with_suffix(".tmp")
    temp.write_text(json.dumps(values, indent=2) + "\n")
    temp.replace(STATUS_FILE)


def stop_child(child: subprocess.Popen[Any] | None) -> None:
    if child is None or child.poll() is not None:
        return
    child.terminate()
    try:
        child.wait(timeout=4)
    except subprocess.TimeoutExpired:
        child.kill()
        child.wait()


def daemon_loop(cfg: dict[str, Any]) -> int:
    ensure_dirs()
    PID_FILE.write_text(str(os.getpid()) + "\n")
    child: subprocess.Popen[Any] | None = None
    stopping = False

    def request_stop(_signum: int, _frame: Any) -> None:
        nonlocal stopping
        stopping = True

    signal.signal(signal.SIGTERM, request_stop)
    signal.signal(signal.SIGINT, request_stop)
    backoff = 2
    try:
        while not stopping:
            ranked, probes = rank_routes(cfg)
            if not ranked:
                ranked, probes = rank_routes(cfg, discover=True)
            if not ranked:
                write_status(state="waiting", pid=os.getpid(), error="no SSH route is reachable", probes=probes)
                time.sleep(min(backoff, 30)); backoff = min(backoff * 2, 30)
                continue
            for route in ranked:
                if stopping:
                    break
                write_status(state="connecting", pid=os.getpid(), route=route["name"], probes=probes)
                child = subprocess.Popen(tunnel_command(cfg, route), stdin=subprocess.DEVNULL)
                deadline = time.monotonic() + float(cfg["connect_timeout"]) + 2
                while time.monotonic() < deadline and child.poll() is None and not stopping:
                    if port_open(str(cfg["listen"]), int(cfg["port"])):
                        break
                    time.sleep(.25)
                if child.poll() is not None or not port_open(str(cfg["listen"]), int(cfg["port"])):
                    stop_child(child); child = None
                    continue
                failures = 0
                backoff = 2
                while not stopping and child.poll() is None:
                    healthy = curl_test(cfg, quiet=True)
                    failures = 0 if healthy else failures + 1
                    write_status(state="healthy" if healthy else "degraded", pid=os.getpid(),
                                 ssh_pid=child.pid, route=route["name"], consecutive_failures=failures)
                    if failures >= int(cfg["failure_threshold"]):
                        break
                    for _ in range(int(cfg["health_interval"] * 2)):
                        if stopping or child.poll() is not None:
                            break
                        time.sleep(.5)
                stop_child(child); child = None
                if not stopping:
                    write_status(state="reconnecting", pid=os.getpid(), route=route["name"])
            if not stopping:
                time.sleep(backoff); backoff = min(backoff * 2, 30)
    finally:
        stop_child(child)
        write_status(state="stopped", pid=os.getpid())
        try:
            PID_FILE.unlink()
        except FileNotFoundError:
            pass
    return 0


def port_open(host: str, port: int) -> bool:
    try:
        with socket.create_connection((host, port), timeout=.5):
            return True
    except OSError:
        return False


def read_pid() -> int | None:
    try:
        pid = int(PID_FILE.read_text().strip())
        os.kill(pid, 0)
        return pid
    except (FileNotFoundError, ValueError, ProcessLookupError, PermissionError):
        return None


def cmd_start(args: argparse.Namespace, cfg: dict[str, Any]) -> int:
    pid = read_pid()
    if pid:
        print(f"multiproxy is already running (PID {pid})")
        return 0
    if port_open(str(cfg["listen"]), int(cfg["port"])):
        print(f"Port {cfg['port']} is already in use; stop the existing proxy first.", file=sys.stderr)
        return 1
    ensure_dirs()
    with LOG_FILE.open("ab", buffering=0) as log:
        proc = subprocess.Popen([sys.executable, str(Path(__file__).resolve()), "--config", str(args.config),
                                 "run"], stdin=subprocess.DEVNULL, stdout=log, stderr=log,
                                start_new_session=True)
    for _ in range(30):
        if read_pid() == proc.pid:
            print(f"Started multiproxy (PID {proc.pid}); use `multiproxy status` to inspect it.")
            return 0
        if proc.poll() is not None:
            break
        time.sleep(.1)
    print(f"Failed to start; inspect {LOG_FILE}", file=sys.stderr)
    return 1


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Self-healing SOCKS5 proxy over the best SSH route")
    parser.add_argument("--config", type=Path, default=CONFIG_FILE)
    commands = parser.add_subparsers(dest="command", required=True)
    commands.add_parser("init", help="create and print the config path")
    routes_parser = commands.add_parser("routes", help="discover, probe, rank, and remember SSH routes")
    routes_parser.add_argument("-j", "--workers", type=int, default=8,
                               help="global parallel SSH probes (default: 8; per-host limits still apply)")
    commands.add_parser("start", help="start the watchdog in the background")
    commands.add_parser("run", help="run the watchdog in the foreground")
    commands.add_parser("stop", help="stop the watchdog and tunnel")
    commands.add_parser("status", help="show current route and health")
    commands.add_parser("test", help="test Google through SOCKS5 on port 1080")
    return parser


def main() -> int:
    parser = build_parser(); args = parser.parse_args()
    cfg = load_config(args.config)
    if args.command == "init":
        print(args.config); return 0
    if args.command == "routes":
        if args.workers < 1:
            parser.error("routes --workers must be at least 1")
        cfg["discovery_workers"] = args.workers
        _, results = rank_routes(cfg, discover=True)
        for result in sorted(results, key=lambda r: (not r["ok"], r["latency_ms"] or 10**9)):
            detail = f"{result['latency_ms']} ms" if result["ok"] else result["error"]
            print(f"{'OK  ' if result['ok'] else 'FAIL'} {result['name']}: {detail}")
        return 0 if any(r["ok"] for r in results) else 1
    if args.command == "run":
        return daemon_loop(cfg)
    if args.command == "start":
        return cmd_start(args, cfg)
    if args.command == "stop":
        pid = read_pid()
        if not pid:
            print("multiproxy is not running"); return 0
        os.kill(pid, signal.SIGTERM)
        for _ in range(50):
            if read_pid() is None:
                print("Stopped multiproxy"); return 0
            time.sleep(.1)
        print(f"Stop requested for PID {pid}"); return 0
    if args.command == "status":
        pid = read_pid()
        try:
            status = json.loads(STATUS_FILE.read_text())
        except (FileNotFoundError, json.JSONDecodeError):
            status = {"state": "unknown"}
        status["running"] = bool(pid)
        print(json.dumps(status, indent=2)); return 0 if pid else 1
    if args.command == "test":
        started = time.monotonic(); ok = curl_test(cfg)
        print(f"{'PASS' if ok else 'FAIL'}: {cfg['test_url']} via socks5h://{cfg['listen']}:{cfg['port']} "
              f"({time.monotonic() - started:.2f}s)")
        return 0 if ok else 1
    return 2


if __name__ == "__main__":
    raise SystemExit(main())