multinetspeeder/cli.py

from __future__ import annotations

import argparse
import asyncio
import json
import re
from collections import defaultdict
from pathlib import Path

from . import __version__
from .config_tcp import TcpConfig
from .config_udp import UdpConfig
from .relay_server_tcp import TcpRelayServer
from .relay_server_udp import UdpRelayServer
from .relay_client_tcp import TcpRelayManager
from .edge_udp import UdpEdge
from .edge_tcp import TcpEdge


TCP_WIN_RE = re.compile(
    r"tcp win session=(?P<session>\d+) target=(?P<host>[^:]+):(?P<port>\d+) winner=(?P<winner>\S+) .*?direct=(?P<direct>\d+) .*?relay=(?P<relay>\d+)"
)
UDP_WIN_RE = re.compile(
    r"udp flow=(?P<flow>\d+) winner=(?P<winner>\S+) target=(?P<host>[^:]+):(?P<port>\d+)"
)
UDP_SOCKS_SUMMARY_RE = re.compile(
    r"udp summary bind=(?P<bind>\S+) flows=(?P<flows>\d+) winners=(?P<winners>\d+) "
    r"winner_breakdown=direct=(?P<direct>\d+),relay=(?P<relay>\d+) "
    r"sample=(?P<sample>.*?) candidates=(?P<candidates>\[.*?\]) "
    r"sent=(?P<sent>\d+) recv=(?P<recv>\d+) dup=(?P<dup>\d+) "
    r"direct_paths=(?P<direct_paths>\d+) relay_paths=(?P<relay_paths>\d+) relay_errors=(?P<relay_errors>.*)"
)


def normalize_winner(name: str) -> str:
    return "direct" if name.startswith("direct") else name


def parse_summary_log(log_path: Path) -> dict[str, object]:
    tcp_total = 0
    tcp_direct = 0
    tcp_relay = 0
    tcp_winners: dict[str, int] = defaultdict(int)
    tcp_targets: dict[str, dict[str, int]] = defaultdict(lambda: defaultdict(int))

    udp_total = 0
    udp_direct = 0
    udp_relay = 0
    udp_winners: dict[str, int] = defaultdict(int)
    udp_targets: dict[str, dict[str, int]] = defaultdict(lambda: defaultdict(int))
    udp_snapshot: dict[str, object] | None = None

    for line in log_path.read_text(errors="replace").splitlines():
        tcp_match = TCP_WIN_RE.search(line)
        if tcp_match:
            tcp_total += 1
            winner = normalize_winner(tcp_match.group("winner"))
            host = tcp_match.group("host")
            port = tcp_match.group("port")
            key = f"{host}:{port}"
            tcp_winners[winner] += 1
            tcp_targets[key][winner] += 1
            if winner == "direct":
                tcp_direct += 1
            else:
                tcp_relay += 1
            continue

        udp_match = UDP_WIN_RE.search(line)
        if udp_match:
            udp_total += 1
            winner = normalize_winner(udp_match.group("winner"))
            host = udp_match.group("host")
            port = udp_match.group("port")
            key = f"{host}:{port}"
            udp_winners[winner] += 1
            udp_targets[key][winner] += 1
            if winner == "direct":
                udp_direct += 1
            else:
                udp_relay += 1
            continue

        udp_snapshot_match = UDP_SOCKS_SUMMARY_RE.search(line)
        if udp_snapshot_match:
            udp_snapshot = {
                "bind": udp_snapshot_match.group("bind"),
                "flows": int(udp_snapshot_match.group("flows")),
                "winners": int(udp_snapshot_match.group("winners")),
                "direct": int(udp_snapshot_match.group("direct")),
                "relay": int(udp_snapshot_match.group("relay")),
                "sample": udp_snapshot_match.group("sample"),
                "candidates": udp_snapshot_match.group("candidates"),
                "sent": int(udp_snapshot_match.group("sent")),
                "recv": int(udp_snapshot_match.group("recv")),
                "dup": int(udp_snapshot_match.group("dup")),
                "direct_paths": int(udp_snapshot_match.group("direct_paths")),
                "relay_paths": int(udp_snapshot_match.group("relay_paths")),
                "relay_errors": udp_snapshot_match.group("relay_errors"),
            }

    return {
        "path": log_path,
        "tcp_total": tcp_total,
        "tcp_direct": tcp_direct,
        "tcp_relay": tcp_relay,
        "tcp_winners": tcp_winners,
        "tcp_targets": tcp_targets,
        "udp_total": udp_total,
        "udp_direct": udp_direct,
        "udp_relay": udp_relay,
        "udp_winners": udp_winners,
        "udp_targets": udp_targets,
        "udp_snapshot": udp_snapshot,
    }


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(prog="mynetspeeder")
    parser.add_argument("--version", action="version", version=f"%(prog)s {__version__}")
    sub = parser.add_subparsers(dest="command", required=True)

    relay_tcp = sub.add_parser("relay-tcp", help="在子节点 VPS 上启动 TCP relay")
    relay_tcp.add_argument("--listen-host", default="0.0.0.0")
    relay_tcp.add_argument("--listen-port", type=int, default=9009)
    relay_tcp.add_argument("--token", required=True)
    relay_tcp.set_defaults(handler=handle_relay_tcp)

    relay_udp = sub.add_parser("relay-udp", help="在子节点 VPS 上启动 UDP relay")
    relay_udp.add_argument("--listen-host", default="0.0.0.0")
    relay_udp.add_argument("--listen-port", type=int, default=9010)
    relay_udp.add_argument("--token", required=True)
    relay_udp.set_defaults(handler=handle_relay_udp)

    edge_tcp = sub.add_parser("edge-tcp", help="在当前主 VPS 上启动 TCP 透明加速")
    edge_tcp.add_argument("--listen-host", default="127.0.0.1")
    edge_tcp.add_argument("--listen-port", type=int, default=19080)
    edge_tcp.add_argument("--config", required=True)
    edge_tcp.add_argument("--kernel", choices=("auto", "20", "24"), default="auto")
    edge_tcp.set_defaults(handler=handle_edge_tcp)

    edge_udp = sub.add_parser("edge-udp", help="在当前主 VPS 上启动 UDP SOCKS5 加速")
    edge_udp.add_argument("--listen-host", default="127.0.0.1")
    edge_udp.add_argument("--listen-port", type=int, default=19180)
    edge_udp.add_argument("--config", required=True)
    edge_udp.set_defaults(handler=handle_edge_udp)

    probe = sub.add_parser("probe", help="查看子节点探测与在线状态")
    probe.add_argument("--config", required=True)
    probe.add_argument("--once", action="store_true")
    probe.set_defaults(handler=handle_probe)

    summary = sub.add_parser("summary", help="汇总当前架构日志里的胜率")
    summary.add_argument("--log-file", default="/var/log/mynetspeeder-tcp-edge.log")
    summary.add_argument("--top", type=int, default=10)
    summary.add_argument("--json", action="store_true")
    summary.set_defaults(handler=handle_summary)
    return parser


def handle_relay_tcp(args: argparse.Namespace) -> int:
    asyncio.run(TcpRelayServer(args.token).start(args.listen_host, args.listen_port))
    return 0


def handle_relay_udp(args: argparse.Namespace) -> int:
    asyncio.run(UdpRelayServer(args.token).start(args.listen_host, args.listen_port))
    return 0


def handle_edge_tcp(args: argparse.Namespace) -> int:
    asyncio.run(TcpEdge(args.listen_host, args.listen_port, TcpConfig.load(args.config), kernel_mode=args.kernel).start())
    return 0


def handle_edge_udp(args: argparse.Namespace) -> int:
    asyncio.run(UdpEdge(args.listen_host, args.listen_port, UdpConfig.load(args.config)).start())
    return 0


def handle_probe(args: argparse.Namespace) -> int:
    async def run_probe() -> None:
        manager = TcpRelayManager(TcpConfig.load(args.config))
        await manager.start()
        await asyncio.sleep(2)
        print(json.dumps(manager.snapshot(), ensure_ascii=False, indent=2))
        if not args.once:
            while True:
                await asyncio.sleep(5)
                print(json.dumps(manager.snapshot(), ensure_ascii=False, indent=2))
    asyncio.run(run_probe())
    return 0


def handle_summary(args: argparse.Namespace) -> int:
    log_path = Path(args.log_file)
    if not log_path.exists():
        raise SystemExit(f"log file not found: {log_path}")
    parsed = parse_summary_log(log_path)
    tcp_total = parsed["tcp_total"]
    tcp_direct = parsed["tcp_direct"]
    tcp_relay = parsed["tcp_relay"]
    tcp_winners = parsed["tcp_winners"]
    tcp_targets = parsed["tcp_targets"]
    udp_total = parsed["udp_total"]
    udp_direct = parsed["udp_direct"]
    udp_relay = parsed["udp_relay"]
    udp_winners = parsed["udp_winners"]
    udp_targets = parsed["udp_targets"]
    udp_snapshot = parsed["udp_snapshot"]

    if udp_total == 0:
        sibling_logs = [
            log_path.with_name("mynetspeeder-udp-socks.log"),
            log_path.with_name("mynetspeeder-relay-tcp.log"),
            log_path.with_name("mynetspeeder-relay-udp.log"),
            log_path.with_name("mynetspeeder-socks.log"),
            log_path.with_name("mynetspeeder-relay.log"),
            log_path.with_name("mynetspeeder-edge.log"),
        ]
        for sibling in sibling_logs:
            if not sibling.exists():
                continue
            sibling_parsed = parse_summary_log(sibling)
            if sibling_parsed["udp_total"] == 0 and sibling_parsed["udp_snapshot"] is None:
                continue
            udp_total = sibling_parsed["udp_total"]
            udp_direct = sibling_parsed["udp_direct"]
            udp_relay = sibling_parsed["udp_relay"]
            udp_winners = sibling_parsed["udp_winners"]
            udp_targets = sibling_parsed["udp_targets"]
            udp_snapshot = sibling_parsed["udp_snapshot"]
            break

    tcp_ordered_targets = sorted(
        tcp_targets.items(),
        key=lambda item: sum(item[1].values()),
        reverse=True,
    )[: max(args.top, 0)]
    udp_ordered_targets = sorted(
        udp_targets.items(),
        key=lambda item: sum(item[1].values()),
        reverse=True,
    )[: max(args.top, 0)]

    result = {
        "log_file": str(log_path),
        "tcp": {
            "total": tcp_total,
            "direct": tcp_direct,
            "relay": tcp_relay,
            "winners": dict(sorted(tcp_winners.items(), key=lambda item: (-item[1], item[0]))),
            "targets": [
                {
                    "target": target,
                    "wins": dict(sorted(counts.items(), key=lambda item: (-item[1], item[0]))),
                }
                for target, counts in tcp_ordered_targets
            ],
        },
        "udp": {
            "total": udp_total,
            "direct": udp_direct,
            "relay": udp_relay,
            "winners": dict(sorted(udp_winners.items(), key=lambda item: (-item[1], item[0]))),
            "targets": [
                {
                    "target": target,
                    "wins": dict(sorted(counts.items(), key=lambda item: (-item[1], item[0]))),
                }
                for target, counts in udp_ordered_targets
            ],
            "snapshot": udp_snapshot,
        },
    }

    if udp_total == 0 and udp_snapshot is not None:
        result["udp"]["total"] = udp_snapshot["winners"]
        result["udp"]["direct"] = udp_snapshot["direct"]
        result["udp"]["relay"] = udp_snapshot["relay"]
    if not result["udp"]["winners"] and udp_snapshot is not None:
        snapshot_winners: dict[str, int] = {}
        if int(udp_snapshot["direct"]) > 0:
            snapshot_winners["direct"] = int(udp_snapshot["direct"])
        if int(udp_snapshot["relay"]) > 0:
            snapshot_winners["relay"] = int(udp_snapshot["relay"])
        result["udp"]["winners"] = snapshot_winners
    if not result["udp"]["targets"] and udp_snapshot is not None:
        snapshot_target_wins: dict[str, int] = {}
        if int(udp_snapshot["direct"]) > 0:
            snapshot_target_wins["direct"] = int(udp_snapshot["direct"])
        if int(udp_snapshot["relay"]) > 0:
            snapshot_target_wins["relay"] = int(udp_snapshot["relay"])
        result["udp"]["targets"] = [
            {
                "target": f"snapshot@{udp_snapshot['bind']}",
                "wins": snapshot_target_wins,
            }
        ]

    if args.json:
        print(json.dumps(result, ensure_ascii=False, indent=2))
        return 0

    print(f"log: {log_path}")
    for protocol in ("tcp", "udp"):
        section = result[protocol]
        print(f"{protocol}: total={section['total']} direct={section['direct']} relay={section['relay']}")
        if protocol == "udp" and section.get("snapshot"):
            snapshot = section["snapshot"]
            print(
                f"  snapshot: bind={snapshot['bind']} flows={snapshot['flows']} winners={snapshot['winners']} "
                f"direct_paths={snapshot['direct_paths']} relay_paths={snapshot['relay_paths']} dup={snapshot['dup']} "
                f"candidates={snapshot['candidates']}"
            )
        print("winners:")
        for name, count in section["winners"].items():
            print(f"  {name}: {count}")
        print("targets:")
        for item in section["targets"]:
            wins = ", ".join(f"{name}={count}" for name, count in item["wins"].items())
            print(f"  {item['target']}: {wins}")
    return 0


def main() -> int:
    parser = build_parser()
    args = parser.parse_args()
    return args.handler(args)