onesvm-browser-server/bench/proxy/probe.py

#!/usr/bin/env python3
"""Probe mihomo proxies. Credentials never written to output."""

from __future__ import annotations

import argparse
import concurrent.futures
import datetime as dt
import json
import os
import subprocess
import sys
import time
import urllib.error
import urllib.parse
import urllib.request
from typing import Any

from lib import api_quote, infer_region, is_info_node

DEFAULT_DELAY_URL = "https://www.google.com/generate_204"
TARGETS = [
    "https://www.google.com/generate_204",
    "https://www.google.com/",
    "https://www.bing.com/",
    "https://www.reddit.com/",
    "https://www.amazon.com/",
    "https://r.jina.ai/",
    "https://duckduckgo.com/",
]
SELECTOR = "🚀节点选择"


def now_iso() -> str:
    return dt.datetime.now().astimezone().isoformat(timespec="seconds")


def api(base: str, path: str, method: str = "GET", body: Any | None = None, timeout: float = 15) -> Any:
    data = None
    headers = {}
    if body is not None:
        data = json.dumps(body).encode("utf-8")
        headers["Content-Type"] = "application/json"
    req = urllib.request.Request(base + path, data=data, method=method, headers=headers)
    with urllib.request.urlopen(req, timeout=timeout) as resp:
        raw = resp.read()
        if not raw:
            return None
        return json.loads(raw.decode("utf-8"))


def list_leaf_proxies(base: str) -> list[dict[str, str]]:
    data = api(base, "/proxies")
    out: list[dict[str, str]] = []
    for name, meta in (data or {}).get("proxies", {}).items():
        ptype = meta.get("type") or ""
        if ptype in {"Selector", "URLTest", "Fallback", "LoadBalance", "Relay", "Compatible", "Direct", "Reject", "Pass", "RejectDrop"}:
            continue
        if name in {"DIRECT", "REJECT", "GLOBAL", "COMPATIBLE", "PASS"}:
            continue
        if is_info_node(name):
            continue
        out.append({"name": name, "type": ptype.lower(), "region": infer_region(name)})
    out.sort(key=lambda x: (x["region"], x["name"]))
    return out


def delay_one(base: str, name: str, url: str, timeout_ms: int) -> dict[str, Any]:
    q = api_quote(name)
    path = f"/proxies/{q}/delay?timeout={timeout_ms}&url={urllib.parse.quote(url, safe='')}"
    started = time.perf_counter()
    try:
        data = api(base, path, timeout=timeout_ms / 1000 + 3)
        delay = data.get("delay") if isinstance(data, dict) else None
        return {
            "name": name,
            "ok": isinstance(delay, int) and delay > 0,
            "delay_ms": delay,
            "error": None,
            "elapsed_s": round(time.perf_counter() - started, 3),
        }
    except urllib.error.HTTPError as exc:
        body = exc.read().decode("utf-8", errors="replace")[:200]
        return {
            "name": name,
            "ok": False,
            "delay_ms": None,
            "error": f"HTTP {exc.code} {body}",
            "elapsed_s": round(time.perf_counter() - started, 3),
        }
    except Exception as exc:  # noqa: BLE001
        return {
            "name": name,
            "ok": False,
            "delay_ms": None,
            "error": f"{type(exc).__name__}: {exc}",
            "elapsed_s": round(time.perf_counter() - started, 3),
        }


def switch_selector(base: str, node: str, group: str = SELECTOR) -> None:
    api(base, f"/proxies/{api_quote(group)}", method="PUT", body={"name": node})


def curl_via_proxy(proxy: str, url: str, method: str = "GET", max_time: int = 20) -> dict[str, Any]:
    cmd = [
        "curl", "-sS", "-o", "/dev/null",
        "-X", method,
        "-w", "%{http_code} %{time_namelookup} %{time_connect} %{time_starttransfer} %{time_total} %{errormsg}",
        "--connect-timeout", "10",
        "--max-time", str(max_time),
        "-x", proxy,
        "-A", "bench-proxy-probe/1.0",
        url,
    ]
    started = time.perf_counter()
    try:
        proc = subprocess.run(cmd, capture_output=True, text=True, check=False)
        parts = (proc.stdout or "").strip().split(" ", 5)
        http_code = parts[0] if parts else ""
        ttfb = float(parts[3]) if len(parts) > 3 and parts[3] else None
        total = float(parts[4]) if len(parts) > 4 and parts[4] else None
        err = (parts[5] if len(parts) > 5 else "") or (proc.stderr or "").strip()
        ok = http_code.isdigit() and http_code[0] in {"2", "3"}
        return {
            "url": url,
            "method": method,
            "ok": ok,
            "http_code": http_code,
            "ttfb_s": ttfb,
            "total_s": total,
            "error": err or None,
            "elapsed_s": round(time.perf_counter() - started, 3),
            "cmd": " ".join(cmd[:-1] + ["<url>"]),
        }
    except Exception as exc:  # noqa: BLE001
        return {
            "url": url,
            "method": method,
            "ok": False,
            "http_code": None,
            "ttfb_s": None,
            "total_s": None,
            "error": f"{type(exc).__name__}: {exc}",
            "elapsed_s": round(time.perf_counter() - started, 3),
            "cmd": "curl -x <proxy> <url>",
        }


def sample_nodes(nodes: list[dict[str, str]], min_n: int = 10) -> list[dict[str, str]]:
    if len(nodes) <= min_n:
        return nodes
    by_region: dict[str, list[dict[str, str]]] = {}
    for n in nodes:
        by_region.setdefault(n["region"], []).append(n)
    picked: list[dict[str, str]] = []
    # Prefer at least one per region, then fill by type diversity.
    for region, items in by_region.items():
        hy2 = [x for x in items if x["type"] in {"hysteria2", "hysteria"}]
        rest = [x for x in items if x not in hy2]
        if hy2:
            picked.append(hy2[0])
        if rest:
            picked.append(rest[0])
        if len(rest) > 3:
            picked.append(rest[len(rest) // 2])
    # unique preserve
    seen = set()
    uniq = []
    for n in picked:
        if n["name"] in seen:
            continue
        seen.add(n["name"])
        uniq.append(n)
    i = 0
    while len(uniq) < min_n and i < len(nodes):
        if nodes[i]["name"] not in seen:
            uniq.append(nodes[i])
            seen.add(nodes[i]["name"])
        i += 1
    return uniq


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--api", default=os.environ.get("MIHOMO_API", "http://127.0.0.1:19090"))
    ap.add_argument("--proxy", default=os.environ.get("MIXED_PROXY", "http://127.0.0.1:17890"))
    ap.add_argument("--out", default=os.path.join(os.path.dirname(__file__), "probe-results.json"))
    ap.add_argument("--delay-url", default=DEFAULT_DELAY_URL)
    ap.add_argument("--timeout-ms", type=int, default=8000)
    ap.add_argument("--all", action="store_true", help="probe every real node (default: sample >=10)")
    ap.add_argument("--min-sample", type=int, default=12)
    ap.add_argument("--workers", type=int, default=6)
    args = ap.parse_args()

    version = api(args.api, "/version")
    leaves = list_leaf_proxies(args.api)
    targets = leaves if args.all else sample_nodes(leaves, args.min_sample)

    print(f"[probe] mihomo={version} leaves={len(leaves)} testing={len(targets)}", flush=True)
    delay_rows: list[dict[str, Any]] = []
    with concurrent.futures.ThreadPoolExecutor(max_workers=args.workers) as pool:
        futs = [
            pool.submit(delay_one, args.api, n["name"], args.delay_url, args.timeout_ms)
            for n in targets
        ]
        for fut in concurrent.futures.as_completed(futs):
            row = fut.result()
            meta = next(x for x in targets if x["name"] == row["name"])
            row.update({"type": meta["type"], "region": meta["region"]})
            delay_rows.append(row)
            status = f"{row['delay_ms']}ms" if row["ok"] else row["error"]
            print(f"  delay {row['region']:6} {row['type']:10} {row['name']} -> {status}", flush=True)

    delay_rows.sort(key=lambda r: (not r["ok"], r["delay_ms"] is None, r["delay_ms"] or 10**9))
    alive = [r for r in delay_rows if r["ok"]]
    top3 = alive[:3]

    matrix: list[dict[str, Any]] = []
    for node in top3:
        print(f"[probe] matrix via {node['name']}", flush=True)
        switch_selector(args.api, node["name"])
        time.sleep(0.4)
        site_rows = []
        for url in TARGETS:
            method = "GET"
            row = curl_via_proxy(args.proxy, url, method=method)
            # some sites dislike HEAD; GET homepage / 204 only
            site_rows.append(row)
            print(
                f"    {row['http_code']} ttfb={row['ttfb_s']} {url} err={row['error']}",
                flush=True,
            )
        matrix.append(
            {
                "node": node["name"],
                "type": node["type"],
                "region": node["region"],
                "delay_ms": node["delay_ms"],
                "sites": site_rows,
            }
        )

    result = {
        "probed_at": now_iso(),
        "mihomo": version,
        "api": args.api,
        "mixed_proxy": args.proxy,
        "delay_url": args.delay_url,
        "timeout_ms": args.timeout_ms,
        "leaf_count": len(leaves),
        "tested_count": len(targets),
        "alive_count": len(alive),
        "alive_ge_3": len(alive) >= 3,
        "delay": delay_rows,
        "top3": [
            {"name": r["name"], "type": r["type"], "region": r["region"], "delay_ms": r["delay_ms"]}
            for r in top3
        ],
        "matrix": matrix,
        "commands": {
            "up": "PROXY_SUB_URL=... bash bench/proxy/up.sh",
            "probe": "python3 bench/proxy/probe.py --all",
            "down": "bash bench/proxy/down.sh",
        },
    }
    os.makedirs(os.path.dirname(args.out) or ".", exist_ok=True)
    with open(args.out, "w", encoding="utf-8") as fh:
        json.dump(result, fh, ensure_ascii=False, indent=2)
        fh.write("\n")
    print(f"[probe] wrote {args.out} alive={len(alive)}/{len(targets)} ge3={len(alive) >= 3}")
    return 0 if alive else 1


if __name__ == "__main__":
    sys.exit(main())