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 import Config
from .relay_server import RelayServer
from .relay_client import RelayManager
from .transparent_edge import TransparentEdge


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+)"
)


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 = sub.add_parser("relay", help="在子节点 VPS 上启动 relay")
    relay.add_argument("--listen-host", default="0.0.0.0")
    relay.add_argument("--listen-port", type=int, default=9009)
    relay.add_argument("--token", required=True)
    relay.set_defaults(handler=handle_relay)

    edge = sub.add_parser("edge", help="在当前主 VPS 上启动透明 direct 出站加速")
    edge.add_argument("--listen-host", default="127.0.0.1")
    edge.add_argument("--listen-port", type=int, default=19080)
    edge.add_argument("--config", required=True)
    edge.add_argument("--enable-udp", action="store_true")
    edge.set_defaults(handler=handle_edge)

    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-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(args: argparse.Namespace) -> int:
    asyncio.run(RelayServer(args.token).start(args.listen_host, args.listen_port))
    return 0


def handle_edge(args: argparse.Namespace) -> int:
    asyncio.run(TransparentEdge(args.listen_host, args.listen_port, Config.load(args.config), enable_udp=args.enable_udp).start())
    return 0


def handle_probe(args: argparse.Namespace) -> int:
    async def run_probe() -> None:
        manager = RelayManager(Config.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}")

    total = 0
    direct = 0
    relay = 0
    winners: dict[str, int] = defaultdict(int)
    targets: dict[str, dict[str, int]] = defaultdict(lambda: defaultdict(int))

    for line in log_path.read_text(errors="replace").splitlines():
        match = WIN_RE.search(line)
        if not match:
            continue
        total += 1
        winner = match.group("winner")
        host = match.group("host")
        port = match.group("port")
        key = f"{host}:{port}"
        winners[winner] += 1
        targets[key][winner] += 1
        if winner == "direct":
            direct += 1
        else:
            relay += 1

    ordered_targets = sorted(
        targets.items(),
        key=lambda item: sum(item[1].values()),
        reverse=True,
    )[: max(args.top, 0)]

    result = {
        "log_file": str(log_path),
        "total": total,
        "direct": direct,
        "relay": relay,
        "winners": dict(sorted(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 ordered_targets
        ],
    }

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

    print(f"log: {log_path}")
    print(f"total: {total} direct: {direct} relay: {relay}")
    print("winners:")
    for name, count in result["winners"].items():
        print(f"  {name}: {count}")
    print("targets:")
    for item in result["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)