#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
EO 全段扫描 v3 (方案 B: 找所有能服务 box.shiroko.one 的 EO 边缘 IP)
两段式: 1) asyncio TCP:443 预筛全段  2) openssl 对 TCP 通的 IP 做 SNI 证书核验 (CN 含 host)
输出: eo-all-hits.json  [{ip, cn, tcp_ms}]
用法: python eo-scan-all.py --host box.shiroko.one --nets 43.174.0.0/16 43.175.0.0/16 --workers 400
"""
import argparse
import asyncio
import json
import socket
import subprocess
import time
from concurrent.futures import ThreadPoolExecutor

def iter_hosts(net: str):
    net, bits = net.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):
        v = base + i
        yield f"{(v >> 24) & 255}.{(v >> 16) & 255}.{(v >> 8) & 255}.{v & 255}"

async def tcp_probe(ip: str, port: int, timeout: float, sem: asyncio.Semaphore):
    async with sem:
        t0 = time.perf_counter()
        try:
            s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            s.setblocking(False)
            loop = asyncio.get_running_loop()
            await asyncio.wait_for(loop.sock_connect(s, (ip, port)), timeout)
            ms = (time.perf_counter() - t0) * 1000
            s.close()
            return ip, ms
        except Exception:
            try:
                s.close()
            except Exception:
                pass
            return None

def tls_cn(ip: str, host: str, timeout: float):
    try:
        r = subprocess.run(
            ["openssl", "s_client", "-connect", f"{ip}:443", "-servername", host,
             "-tlsextdebug", "-no_ign_eof"],
            capture_output=True, text=True, timeout=timeout, errors="replace",
            input="Q")
        out = r.stdout + r.stderr
        cn = ""
        for line in out.splitlines():
            ls = line.strip()
            if ls.startswith("subject="):
                cn = ls
                break
        return {"ip": ip, "cn": cn or "(no-subject)", "ok": host in cn}
    except Exception as e:
        return {"ip": ip, "cn": f"ERR:{type(e).__name__}", "ok": False}

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--host", default="box.shiroko.one")
    ap.add_argument("--nets", nargs="*",
                    default=["43.174.0.0/16", "43.175.0.0/16"])
    ap.add_argument("--workers", type=int, default=300)
    ap.add_argument("--timeout", type=float, default=2.0)
    ap.add_argument("--out", default="eo-all-hits.json")
    args = ap.parse_args()

    all_ips = []
    for n in args.nets:
        all_ips.extend(iter_hosts(n))
    print(f"[*] host={args.host} 网段={args.nets} 总IP={len(all_ips)}")

    # Stage 1: TCP 预筛
    sem = asyncio.Semaphore(args.workers)
    open_ips = []

    async def run_tcp():
        CH = args.workers * 6
        for i in range(0, len(all_ips), CH):
            chunk = all_ips[i:i + CH]
            tasks = [asyncio.create_task(tcp_probe(ip, 443, args.timeout, sem)) for ip in chunk]
            res = await asyncio.gather(*tasks)
            for r in res:
                if r:
                    open_ips.append(r)
            if (i + CH) % 20000 < CH or i + CH >= len(all_ips):
                print(f"    tcp预筛 [{i+CH}/{len(all_ips)}] open={len(open_ips)}", flush=True)
    t0 = time.perf_counter()
    asyncio.run(run_tcp())
    print(f"[*] TCP 预筛完成: {len(all_ips)} -> {len(open_ips)} open, {time.perf_counter()-t0:.0f}s")

    if not open_ips:
        print("[!] 无开放 IP")
        return

    # Stage 2: openssl 证书核验
    hits = []
    with ThreadPoolExecutor(max_workers=64) as ex:
        futs = [ex.submit(tls_cn, ip, args.host, args.timeout * 2) for ip, _ in open_ips]
        done = 0
        for fut in futures_as_completed(futs):
            done += 1
            r = fut.result()
            if r and r["ok"]:
                hits.append(r)
                print(f"  HIT {r['ip']:16s} {r['cn']}", flush=True)
            if done % 200 == 0:
                print(f"    证书核验 [{done}/{len(open_ips)}] hits={len(hits)}", flush=True)

    print(f"\n[*] 全部完成: hits={len(hits)}")
    if hits:
        with open(args.out, "w", encoding="utf-8") as f:
            json.dump(hits, f, ensure_ascii=False, indent=2)
        print(f"[*] saved: {args.out}")

def futures_as_completed(futs):
    import concurrent.futures
    for fut in concurrent.futures.as_completed(futs):
        yield fut

if __name__ == "__main__":
    main()