remote/tools/punch_probe.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""UDP 打洞探针 —— 两台不同网络的机器各跑一个，用来摸清你们线路的底牌。

回答四个问题：
  1. 这条线路的 NAT 能不能打洞？
  2. 要几轮才能打通？（轮数越少说明 NAT 越宽松）
  3. 对方的源端口有没有被 NAT 改写？（判断是不是对称 NAT）
  4. 有没有公网 IPv6？IPv6 直连通不通？

用法
----
A 机（假设公网 1.2.3.4）:
    py punch_probe.py --peer 2.3.4.5:47800

B 机（假设公网 2.3.4.5）:
    py punch_probe.py --peer 1.2.3.4:47800

IPv6（双方都有公网 v6 时）:
    py punch_probe.py --peer [2408:8207:1234::1]:47800

先查自己的公网映射（走 STUN，不需要对方）:
    py punch_probe.py --stun

被动监听（不发包，只看能不能收到对方打进来的包）:
    py punch_probe.py

说明
----
- 两台机器必须**同时**启动，脚本靠「双方同时向对方发包」来穿透 NAT。
- 先各自查公网 IP（浏览器搜 "my ip" 或 --stun），再互相填。
- 只依赖 Python 标准库。
"""

import argparse
import socket
import struct
import sys
import threading
import time

MAGIC = b"PUNCHv1|"


# ---------------------------------------------------------------- 地址解析

def parse_addr(s):
    """解析 'ip:port' 或 '[v6]:port' -> (host, port, family)"""
    s = s.strip()
    if s.startswith("["):
        end = s.index("]")
        host = s[1:end]
        port = int(s[end + 2:])
        return host, port, socket.AF_INET6
    host, _, p = s.rpartition(":")
    return host, int(p), socket.AF_INET


def fmt_addr(sockaddr):
    host, port = sockaddr[0], sockaddr[1]
    if ":" in host:
        return "[%s]:%d" % (host, port)
    return "%s:%d" % (host, port)


# ---------------------------------------------------------------- STUN

def stun_mapped(server="stun.l.google.com", port=19302, timeout=3.0, sock=None):
    """最小 STUN Binding 请求，返回 (ip, port)；失败返回 None。

    sock 不为 None 时复用它（用来测「同一个本地端口对另一个目标」的映射端口，
    这是判断 NAT 是不是对称型的关键）。
    """
    txid = b"\x5a" * 12
    req = struct.pack(">HHI12s", 0x0001, 0, 0x2112A442, txid)
    try:
        infos = socket.getaddrinfo(server, port, 0, socket.SOCK_DGRAM)
    except OSError:
        return None
    fam, sockaddr = infos[0][0], infos[0][4]
    own = sock is None
    s = socket.socket(fam, socket.SOCK_DGRAM) if own else sock
    s.settimeout(timeout)
    try:
        s.sendto(req, sockaddr)
        data, _ = s.recvfrom(2048)
    except OSError:
        return None
    finally:
        if own:
            s.close()

    if len(data) < 20:
        return None
    off = 20
    while off + 4 <= len(data):
        atype, alen = struct.unpack(">HH", data[off:off + 4])
        val = data[off + 4:off + 4 + alen]
        if atype == 0x0020 and len(val) >= 8:          # XOR-MAPPED-ADDRESS
            famv = val[1]
            xport = struct.unpack(">H", val[2:4])[0] ^ 0x2112
            if famv == 0x01 and len(val) >= 8:
                ip = bytes(a ^ b for a, b in zip(val[4:8], b"\x21\x12\xA4\x42"))
                return socket.inet_ntoa(ip), xport
            if famv == 0x02 and len(val) >= 20:
                mask = b"\x21\x12\xA4\x42" + txid
                ip = bytes(a ^ b for a, b in zip(val[4:20], mask))
                return socket.inet_ntop(socket.AF_INET6, ip), xport
        off += 4 + alen + ((4 - alen % 4) % 4)
    return None


# ---------------------------------------------------------------- 信令

def discover_peer(sock, server, room, timeout=60.0):
    """注册到房间，等服务器回推对方的公网地址。

    注意：这里**复用同一个 socket**——服务器看到的源地址，就是 NAT 为
    「本机 → 服务器」这条流分配的映射。对锥形 NAT 而言，它与「本机 → 对端」
    的映射相同，所以拿它去打洞是有效的；对对称 NAT 则不同，打不通属预期。
    """
    host, port = server[0], server[1]
    deadline = time.time() + timeout
    last_send = 0.0
    while time.time() < deadline:
        now = time.time()
        if now - last_send > 1.0:
            try:
                sock.sendto(b"REGISTER|" + room.encode("utf-8"), (host, port))
            except OSError as e:
                print("  注册发送失败：%s" % e)
                return None
            last_send = now
        try:
            data, _src = sock.recvfrom(2048)
        except socket.timeout:
            continue
        except ConnectionResetError:
            continue
        except OSError as e:
            if getattr(e, "winerror", None) == 10054:
                continue
            return None
        txt = data.decode("utf-8", "replace").strip()
        if txt.startswith("PEER|"):
            h, p, _f = parse_addr(txt[5:])
            return (h, p)
    return None


# ---------------------------------------------------------------- 主逻辑

class Probe:
    def __init__(self, bind_host, port, family):
        self.sock = socket.socket(family, socket.SOCK_DGRAM)
        self.sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        if family == socket.AF_INET6:
            try:
                self.sock.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY, 1)
            except OSError:
                pass
        self.sock.bind((bind_host, port))
        self.sock.settimeout(0.2)
        # Windows 特有坑：曾经发过包给对方，但对方端口没开 → 会收到 ICMP 端口不可达，
        # 之后下一次 recv 直接抛 WSAECONNRESET(10054)。不打这个 ioctl，
        # 接收循环会被"自己之前的失败发送"打断，之后再也收不到包。
        # （Rust 实现同样要处理：tokio recv 返回 ConnectionReset 时不能退出循环）
        if sys.platform == "win32":
            try:
                self.sock.ioctl(socket.SIO_UDP_CONNRESET, False)
            except (AttributeError, OSError):
                pass
        self.port = self.sock.getsockname()[1]
        self.stop = threading.Event()
        self.first_src = None
        self.packets = 0
        self.junk = 0

    def receiver(self, ignore=()):
        """接收线程。

        只认「以 MAGIC 开头」且「来自对端而非信令服务器」的包。
        信令服务器的 PEER/WAIT 回复绝不能算作打通——否则客户端反复注册引来的
        回复会造成假阳性（已踩过：把服务器的 PEER 回复误判成打洞成功）。
        """
        while not self.stop.is_set():
            try:
                data, src = self.sock.recvfrom(2048)
            except socket.timeout:
                continue
            except ConnectionResetError:
                # ICMP 端口不可达留下的错误，忽略并继续收（绝不能 break）
                continue
            except OSError as e:
                if getattr(e, "winerror", None) == 10054:
                    continue
                break

            # ---- 过滤：非打洞包一律丢弃
            if not data.startswith(MAGIC) or (ignore and src[:2] in ignore):
                self.junk += 1
                continue

            if self.first_src is None:
                self.first_src = src
                print("\n  ✅ 收到对方的包！")
                print("     对方源地址（NAT 之后）= %s" % fmt_addr(src))
                if src[1] != self.port:
                    print("     对方源端口 = %d（本机端口 %d）；若双方约定端口相同，"
                          "说明 NAT 改写了端口" % (src[1], self.port))
                else:
                    print("     对方源端口与约定端口一致 → NAT 保留了端口")
            self.packets += 1
            if self.packets <= 3:
                print("     [%s] 收到 %d 字节" % (fmt_addr(src), len(data)))

    def send_rounds(self, peer_sockaddr, rounds, interval):
        for i in range(1, rounds + 1):
            if self.stop.is_set():
                return
            payload = MAGIC + struct.pack(">d", time.time())
            try:
                self.sock.sendto(payload, peer_sockaddr)
            except OSError as e:
                print("  发送失败：%s" % e)
                return
            if i == 1 or i % 5 == 0 or self.first_src is not None:
                mark = "✅" if self.first_src is not None else "  "
                print("  %s 第 %2d/%d 轮已发出" % (mark, i, rounds))
            if self.first_src is not None:
                return
            time.sleep(interval)


def main():
    ap = argparse.ArgumentParser(add_help=True, description="UDP 打洞探针")
    ap.add_argument("--peer", help="对方地址 ip:port 或 [v6]:port；不给则只被动监听")
    ap.add_argument("--port", type=int, default=47800, help="本地端口（默认 47800）")
    ap.add_argument("--rounds", type=int, default=25, help="发包轮数（默认 25）")
    ap.add_argument("--interval", type=float, default=0.2, help="轮间隔秒（默认 0.2）")
    ap.add_argument("--grace", type=float, default=2.5,
                    help="发完后继续监听多少秒等回包（默认 2.5）")
    ap.add_argument("--server", help="信令服务器 ip:port（与 --room 配合，自动配对）")
    ap.add_argument("--room", help="房间名；同一房间的两个客户端会被配对")
    ap.add_argument("--wait", type=float, default=60.0,
                    help="等对方加入房间的秒数（默认 60）")
    ap.add_argument("--stun", action="store_true", help="只查询自己的公网映射后退出")
    ap.add_argument("--stun-port", type=int, default=0,
                    help="配合 --stun：用指定本地端口发 STUN（与打洞端口一致时才能对比映射）")
    args = ap.parse_args()

    # ---- 只查 STUN
    if args.stun:
        s = None
        if args.stun_port:
            s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            s.bind(("0.0.0.0", args.stun_port))
            print("[STUN] 用本机端口 %d 查询（与打洞同端口）..." % args.stun_port)
        else:
            print("[STUN] 查询公网映射（临时端口）...")
        r = stun_mapped(sock=s)
        if s:
            s.close()
        if r:
            print("  对 Google STUN 的映射 = %s:%d" % r)
            if args.stun_port:
                print("  本机端口              = %d" % args.stun_port)
                print("  下一步：对比信令服务器日志里记录的源端口")
                print("    两者相同 → 端点无关映射（锥形 NAT）→ 打洞大概率成功")
                print("    两者不同 → 对称 NAT → 打洞困难，需要端口预测或中继")
        else:
            print("  查询失败（STUN 19302 可能被封；可稍后重试）")
        return 0

    # ---- 决定地址族与模式
    peer = None
    if args.peer:
        host, pport, family = parse_addr(args.peer)
        peer = (host, pport)
        bind_host = "::" if family == socket.AF_INET6 else "0.0.0.0"
        mode = "双向打洞（手工填地址）"
    elif args.server and args.room:
        _h, _p, family = parse_addr(args.server)
        bind_host = "::" if family == socket.AF_INET6 else "0.0.0.0"
        mode = "信令自动配对 + 打洞"
    else:
        family = socket.AF_INET
        bind_host = "0.0.0.0"
        mode = "被动监听"

    probe = Probe(bind_host, args.port, family)
    print("=" * 60)
    print(" UDP 打洞探针  模式=%s" % mode)
    print(" 本机监听      %s:%d" % (bind_host, probe.port))
    if peer:
        print(" 发往          %s  （%s）" % (fmt_addr(peer + (0, 0)),
                                          "IPv6" if family == socket.AF_INET6 else "IPv4"))
    print("=" * 60)

    # ---- 信令配对：不用自己查公网 IP，服务器看到的就是你的公网映射
    if args.server and args.room and peer is None:
        srv = parse_addr(args.server)[:2]
        print("\n向信令服务器 %s 注册 room=%s ..." % (fmt_addr(srv), args.room))
        found = discover_peer(probe.sock, srv, args.room, timeout=args.wait)
        if found is None:
            print("  ❌ 超时：没等到对方加入同一个房间")
            probe.sock.close()
            return 1
        peer = found
        print("  ✅ 配对成功，对方公网地址 = %s" % fmt_addr(found))

    # 信令服务器的回复（PEER/WAIT）绝不能算作打通，接收线程必须把它排除掉
    ignore = set()
    if args.server:
        _sh, _sp, _sf = parse_addr(args.server)
        ignore.add((_sh, _sp))

    t = threading.Thread(target=probe.receiver, args=(ignore,), daemon=True)
    t.start()

    try:
        if peer:
            print("\n开始同时发包（若对方也在发包，第 1~3 轮通常就通）...")
            probe.send_rounds(peer, args.rounds, args.interval)
            # 对方可能比我晚启动：发完后继续监听，否则会把"晚到的成功"误判成失败
            if probe.first_src is None:
                print("\n发包结束，继续监听 %.1f 秒等对方回包 ..." % args.grace)
                waited = 0.0
                while waited < args.grace and probe.first_src is None:
                    time.sleep(0.1)
                    waited += 0.1
        else:
            print("\n被动监听中，Ctrl+C 退出 ...")
            while not probe.stop.is_set():
                time.sleep(0.3)
    except KeyboardInterrupt:
        pass
    finally:
        probe.stop.set()
        time.sleep(0.3)
        probe.sock.close()

    print("\n" + "-" * 60)
    if probe.first_src is not None:
        print("结论：✅ 打洞成功")
        print("      对方实际源地址 = %s" % fmt_addr(probe.first_src))
        print("      共收到 %d 个包" % probe.packets)
    else:
        print("结论：❌ 本轮没有收到对方的打洞包")
        if probe.junk:
            print("  （另外收到 %d 个非打洞包并已忽略：那是信令服务器的回复，不算打通）"
                  % probe.junk)
        print("  可能原因（按概率排序）：")
        print("   1. 一端或两端是对称型 NAT / 运营商级 NAT(CGNAT) → 需要端口预测或中继")
        print("   2. 一方根本没发包（对方没启动、或地址填错/端口填错）")
        print("   3. UDP 被防火墙拦截（企业网络常见）")
        print("   4. 双方没有真正同时发包（时间差超过 NAT 表项寿命）")
        print("  下一步：两端各跑一次 --stun，看看是不是落在 100.64~100.127 网段（CGNAT）")
    print("-" * 60)
    return 0 if probe.first_src is not None else 1


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