#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
EO 优选 IP 免费探测脚本 (方案 B 免费替代 wetest.vip)
原理: EdgeOne 边缘节点(腾讯云香港 43.174.x/43.175.x 段)对任意 EO 接入域名做 SNI 路由。
      若某 IP:443 能对 box.shiroko.one 完成 TLS 握手, 则该 IP 可作为境内 A 记录直连目标。
用法:
  python eo-probe-ips.py                     # 扫默认段列表(已知优选 IP 所在 /24)
  python eo-probe-ips.py --host box.shiroko.one --out eo-hits.json
判定: TCP 连通 + TLS 握手成功(SNI=目标host) -> hit; 记录握手 RTT。
"""
import argparse
import asyncio
import json
import socket
import ssl
import sys
import time

# 已知/疑似 EO 香港优质 IP 所在段(来自社区 itdog 截图统计 + 教程)
DEFAULT_CIDRS = [
    "43.174.150.0/24", "43.174.78.0/24", "43.174.79.0/24", "43.174.57.0/24",
    "43.175.130.0/24", "43.175.132.0/24", "43.175.151.0/24", "43.175.165.0/24",
    "43.175.185.0/24", "43.175.150.0/24", "43.174.152.0/24", "43.174.153.0/24",
    "43.175.131.0/24", "43.175.133.0/24",
]

def cidr_hosts(cidr: str):
    net, bits = cidr.split("/")
    bits = int(bits)
    a, b, c, d = map(int, net.split("."))
    base = (a << 24) + (b << 16) + (c << 8) + d
    count = 1 << (32 - bits)
    for i in range(1, count - 1):  # skip network & broadcast
        v = base + i
        yield f"{(v >> 24) & 255}.{(v >> 16) & 255}.{(v >> 8) & 255}.{v & 255}"

async def probe(ip: str, host: str, timeout: float, sem: asyncio.Semaphore):
    """TCP connect + TLS handshake (SNI=host). Return dict or None."""
    async with sem:
        loop = asyncio.get_running_loop()
        t0 = time.perf_counter()
        try:
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.settimeout(timeout)
            # async connect
            await loop.sock_connect(sock, (ip, 443))
            tcp_ms = (time.perf_counter() - t0) * 1000
            ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
            ctx.check_hostname = False
            ctx.verify_mode = ssl.CERT_NONE
            ctx.set_alpn_protocols(["http/1.1"])
            tls_sock = ctx.wrap_socket(sock, server_hostname=host)
            tls_sock.settimeout(timeout)
            tls_ms = (time.perf_counter() - t0) * 1000
            ver = tls_sock.version()
            alpn = tls_sock.selected_alpn_protocol()
            tls_sock.close()
            sock.close()
            return {"ip": ip, "tcp_ms": round(tcp_ms, 1),
                    "handshake_ms": round(tls_ms, 1), "tls": ver, "alpn": alpn}
        except Exception as e:
            try:
                sock.close()
            except Exception:
                pass
            return None

async def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--host", default="box.shiroko.one")
    ap.add_argument("--cidrs", nargs="*", default=None)
    ap.add_argument("--concurrency", type=int, default=256)
    ap.add_argument("--timeout", type=float, default=2.5)
    ap.add_argument("--out", default="eo-hits.json")
    args = ap.parse_args()

    cidrs = args.cidrs or DEFAULT_CIDRS
    ips = []
    for c in cidrs:
        ips.extend(cidr_hosts(c))
    print(f"[*] 目标 host: {args.host}  段数: {len(cidrs)}  待测 IP: {len(ips)}  并发: {args.concurrency}")

    sem = asyncio.Semaphore(args.concurrency)
    hits = []
    done = 0
    t_start = time.perf_counter()

    # chunk to avoid building all tasks at once
    CH = args.concurrency * 4
    for i in range(0, len(ips), CH):
        chunk = ips[i:i + CH]
        tasks = [asyncio.create_task(probe(ip, args.host, args.timeout, sem)) for ip in chunk]
        results = await asyncio.gather(*tasks)
        for r in results:
            if r:
                hits.append(r)
        done += len(chunk)
        if done % (len(ips) // 10 or 1) == 0 or done == len(ips):
            el = time.perf_counter() - t_start
            print(f"    [{done}/{len(ips)}] hits={len(hits)}  elapsed={el:.1f}s")

    hits.sort(key=lambda x: x["handshake_ms"])
    print(f"\n[*] 完成: {done} IP 用时 {time.perf_counter()-t_start:.1f}s, 命中 {len(hits)}")
    for h in hits[:40]:
        print(f"    {h['ip']:16s} tcp={h['tcp_ms']:6.1f}ms  hs={h['handshake_ms']:6.1f}ms  tls={h['tls']}  alpn={h['alpn']}")

    if hits and args.out:
        with open(args.out, "w", encoding="utf-8") as f:
            json.dump(hits, f, ensure_ascii=False, indent=2)
        print(f"[*] 结果已存: {args.out}")

if __name__ == "__main__":
    asyncio.run(main())