multinetspeeder/cli.py

from __future__ import annotations

import argparse
import asyncio
import json

from . import __version__
from .config import Config
from .relay_server import RelayServer
from .relay_client import RelayManager
from .transparent_edge import TransparentEdge


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.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)
    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)).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 main() -> int:
    parser = build_parser()
    args = parser.parse_args()
    return args.handler(args)