sub2node.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
sub2node — CLI entry point
==========================
Replaces start.sh / stop.sh with a clean Python interface.

Usage:
    python3 sub2node.py start [URL]     fetch sub, generate config, start mihomo
    python3 sub2node.py stop            stop mihomo
    python3 sub2node.py probe           probe ports, write proxies.json
    python3 sub2node.py run [URL]       start + probe in one step (recommended)
    python3 sub2node.py status          show mihomo process and pool stats
"""

import argparse
import json
import os
import signal
import subprocess
import sys
import time

_DIR = os.path.dirname(os.path.abspath(__file__))
_PID_FILE = os.path.join(_DIR, "mihomo.pid")
_LOG_FILE = os.path.join(_DIR, "mihomo.log")
_CONFIG   = os.path.join(_DIR, "config.yaml")
_POOL     = os.path.join(_DIR, "proxies.json")


# ── Helpers ───────────────────────────────────────────────────────────────────

def _mihomo_bin() -> str:
    """Locate mihomo binary; exit with a helpful message if missing."""
    for candidate in ("mihomo", "clash-meta"):
        if subprocess.run(
            ["which", candidate], capture_output=True
        ).returncode == 0:
            return candidate
    sys.exit(
        "Error: mihomo not found.\n"
        "  macOS:  brew install mihomo\n"
        "  Linux:  https://github.com/MetaCubeX/mihomo/releases"
    )


def _running_pid() -> int | None:
    if os.path.exists(_PID_FILE):
        try:
            pid = int(open(_PID_FILE).read().strip())
            os.kill(pid, 0)      # raises if process is gone
            return pid
        except (ValueError, ProcessLookupError, PermissionError):
            os.unlink(_PID_FILE)
    return None


def _stop_mihomo() -> bool:
    pid = _running_pid()
    if pid is None:
        return False
    os.kill(pid, signal.SIGTERM)
    for _ in range(20):
        time.sleep(0.2)
        try:
            os.kill(pid, 0)
        except ProcessLookupError:
            break
    if os.path.exists(_PID_FILE):
        os.unlink(_PID_FILE)
    return True


def _resolve_sub(url_arg: str | None) -> str:
    url = (url_arg or os.environ.get("PROXY_SUB", "")).strip()
    if not url:
        sys.exit(
            "Error: subscription URL required.\n"
            "  python3 sub2node.py start <url>\n"
            "  or  export PROXY_SUB=<url>"
        )
    return url


# ── Commands ──────────────────────────────────────────────────────────────────

def cmd_start(url_arg: str | None) -> None:
    sub = _resolve_sub(url_arg)
    mihomo = _mihomo_bin()

    print("→ Generating config …")
    result = subprocess.run(
        [sys.executable, os.path.join(_DIR, "gen_config.py"), sub],
        cwd=_DIR,
    )
    if result.returncode != 0:
        sys.exit("gen_config.py failed.")

    _stop_mihomo()
    print("→ Starting mihomo …")
    log = open(_LOG_FILE, "w")
    proc = subprocess.Popen(
        [mihomo, "-d", _DIR, "-f", _CONFIG],
        stdout=log, stderr=log,
        cwd=_DIR,
    )
    time.sleep(2)
    if proc.poll() is not None:
        log.close()
        print("Error: mihomo exited immediately. Last log lines:")
        with open(_LOG_FILE) as f:
            print("".join(f.readlines()[-20:]))
        sys.exit(1)

    with open(_PID_FILE, "w") as f:
        f.write(str(proc.pid))

    print(f"✓ mihomo running  pid={proc.pid}  log={_LOG_FILE}")
    print("  Run  python3 sub2node.py probe  to discover alive nodes.")


def cmd_stop() -> None:
    if _stop_mihomo():
        print("✓ mihomo stopped.")
    else:
        print("mihomo is not running.")


def cmd_probe(args: argparse.Namespace) -> None:
    probe_args = [sys.executable, os.path.join(_DIR, "probe.py")]
    if hasattr(args, "workers") and args.workers:
        probe_args += ["--workers", str(args.workers)]
    result = subprocess.run(probe_args, cwd=_DIR)
    sys.exit(result.returncode)


def cmd_run(url_arg: str | None, workers: int | None) -> None:
    """start + probe combined."""
    sub = _resolve_sub(url_arg)
    mihomo = _mihomo_bin()

    print("→ Generating config …")
    result = subprocess.run(
        [sys.executable, os.path.join(_DIR, "gen_config.py"), sub],
        cwd=_DIR,
    )
    if result.returncode != 0:
        sys.exit("gen_config.py failed.")

    _stop_mihomo()
    print("→ Starting mihomo …")
    log = open(_LOG_FILE, "w")
    proc = subprocess.Popen(
        [mihomo, "-d", _DIR, "-f", _CONFIG],
        stdout=log, stderr=log, cwd=_DIR,
    )
    time.sleep(2)
    if proc.poll() is not None:
        log.close()
        with open(_LOG_FILE) as f:
            print("".join(f.readlines()[-20:]))
        sys.exit("mihomo failed to start.")

    with open(_PID_FILE, "w") as f:
        f.write(str(proc.pid))
    print(f"✓ mihomo running  pid={proc.pid}")

    print("→ Probing nodes …")
    probe_cmd = [sys.executable, os.path.join(_DIR, "probe.py")]
    if workers:
        probe_cmd += ["--workers", str(workers)]
    subprocess.run(probe_cmd, cwd=_DIR)


def cmd_status() -> None:
    pid = _running_pid()
    if pid:
        print(f"mihomo  running  pid={pid}")
    else:
        print("mihomo  stopped")

    if os.path.exists(_POOL):
        with open(_POOL, encoding="utf-8") as f:
            pool = json.load(f)
        mtime = time.strftime(
            "%Y-%m-%d %H:%M:%S",
            time.localtime(os.path.getmtime(_POOL))
        )
        uniq = len({p["egress_ip"] for p in pool})
        print(f"proxies.json  {len(pool)} alive nodes  {uniq} unique IPs  (updated {mtime})")
        if pool:
            print(f"  fastest: {pool[0]['egress_ip']}  {pool[0]['latency_s']}s  {pool[0]['name']}")
    else:
        print("proxies.json  not found — run: python3 sub2node.py probe")


# ── CLI ───────────────────────────────────────────────────────────────────────

def main() -> None:
    ap = argparse.ArgumentParser(
        prog="sub2node",
        description="sub2node — turn a proxy subscription into a local port pool",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
examples:
  python3 sub2node.py run https://your-airport/token
  PROXY_SUB=https://... python3 sub2node.py run
  python3 sub2node.py status
  python3 sub2node.py stop
""",
    )
    sub = ap.add_subparsers(dest="cmd", metavar="command")

    p_start = sub.add_parser("start",  help="fetch subscription, generate config, start mihomo")
    p_start.add_argument("url", nargs="?", help="subscription URL (or set PROXY_SUB env var)")

    sub.add_parser("stop", help="stop mihomo")

    p_probe = sub.add_parser("probe", help="probe ports and write proxies.json")
    p_probe.add_argument("--workers", type=int, default=None, help="concurrent threads (default: 16)")

    p_run = sub.add_parser("run", help="start + probe in one step (recommended)")
    p_run.add_argument("url", nargs="?", help="subscription URL (or set PROXY_SUB env var)")
    p_run.add_argument("--workers", type=int, default=None, help="probe concurrency (default: 16)")

    sub.add_parser("status", help="show mihomo status and pool summary")

    args = ap.parse_args()

    if args.cmd == "start":
        cmd_start(args.url)
    elif args.cmd == "stop":
        cmd_stop()
    elif args.cmd == "probe":
        cmd_probe(args)
    elif args.cmd == "run":
        cmd_run(args.url, args.workers)
    elif args.cmd == "status":
        cmd_status()
    else:
        ap.print_help()


if __name__ == "__main__":
    main()