"""网络层验证码图片采集器（CDP websocket 常驻监听）。

原理：浏览器级别的 Target.setAutoAttach 自动 attach 所有新 target
（含验证码 iframe），Network 域监听全部请求响应，凡 mime 为 image/*
或 URL 含 captcha/verify 的响应体直接落盘——不经过 DOM/wasm 渲染层，
字节与服务器下发的完全一致。

用法（验证码触发前启动，后台常驻）:
    .venv/bin/python try/captcha_net_capture.py            # 默认监听 10 分钟
    .venv/bin/python try/captcha_net_capture.py --minutes 30

输出: try/captures/net_<时间戳>_<序号>.<ext>（按 Content-Type 扩展名）
      + try/captures/net_<时间戳>_get.json（captcha/get 响应原文）
"""
import argparse
import base64
import json
import sys
import time
import urllib.request
from pathlib import Path

try:
    import websockets
except ImportError:
    print("需要 websockets: .venv/bin/pip install websockets", file=sys.stderr)
    sys.exit(1)

CDP_HTTP = "http://127.0.0.1:9222"
OUTDIR = Path(__file__).resolve().parent / "captures"

INTEREST_URL = ("captcha", "verify", "rmc", "secsdk")
IMAGE_MIME = {"image/png": "png", "image/jpeg": "jpg", "image/webp": "webp", "image/gif": "gif", "image/bmp": "bmp"}


class Collector:
    def __init__(self, ws):
        self.ws = ws
        self.msg_id = 1
        self.sessions = {}          # sessionId -> None
        self.pending = {}           # msg_id -> (url, mime) 等待 getResponseBody
        self.finished = {}          # requestId -> (url, mime) 待取响应体
        self.n_saved = 0

    async def send(self, method, params=None, session_id=None):
        mid = self.msg_id
        self.msg_id += 1
        msg = {"id": mid, "method": method}
        if params:
            msg["params"] = params
        if session_id:
            msg["sessionId"] = session_id
        await self.ws.send(json.dumps(msg))
        return mid

    async def recv_loop(self, deadline):
        import asyncio
        saved_any_at = time.time()
        while time.time() < deadline:
            try:
                raw = await asyncio.wait_for(self.ws.recv(), timeout=1.0)
            except asyncio.TimeoutError:
                # 超时窗口内把已完成请求的响应体取回
                await self.drain_bodies()
                continue
            except websockets.ConnectionClosed:
                print("[conn closed]", flush=True)
                break
            try:
                m = json.loads(raw)
            except json.JSONDecodeError:
                continue
            method = m.get("method", "")

            if method == "Target.attachedToTarget":
                # 新 target（页面/iframe）出现，对其启用 Network
                sid = m["params"]["sessionId"]
                ti = m["params"]["targetInfo"]
                self.sessions[sid] = ti.get("url", "")
                print(f"[attach] {ti.get('type')} {ti.get('url', '')[:80]}", flush=True)
                await self.send("Network.enable", session_id=sid)

            elif method == "Network.requestWillBeSent":
                u = m["params"]["request"]["url"]
                if any(k in u for k in INTEREST_URL):
                    print(f"[req ] {m['params']['request']['method']} {u[:110]}", flush=True)

            elif method == "Network.responseReceived":
                p = m["params"]
                u = p["response"]["url"]
                mime = (p["response"].get("mimeType") or "").split(";")[0].strip()
                interesting = mime in IMAGE_MIME or any(k in u for k in INTEREST_URL)
                if interesting:
                    self.finished[p["requestId"]] = (u, mime)
                    print(f"[resp] {p['response']['status']} {mime or '?'} {u[:100]}", flush=True)

            elif method == "Network.loadingFinished":
                rid = m["params"]["requestId"]
                if rid in self.finished:
                    # 立即拉响应体（需带 sessionId）
                    info = self.finished.pop(rid)
                    sid = m.get("sessionId")
                    if sid:
                        mid = await self.send("Network.getResponseBody", {"requestId": rid}, session_id=sid)
                        self.pending[mid] = info

            elif "id" in m and m["id"] in self.pending and "result" in m:
                await self.save_body(m, self.pending.pop(m["id"]))

        await self.drain_bodies()

    _body_seq = 100000

    async def next_body_id(self):
        Collector._body_seq += 1
        return Collector._body_seq

    async def drain_bodies(self):
        """把已完成请求的响应体真正拉回（发 getResponseBody，响应由 recv_loop 处理）。"""
        for rid in list(self.finished.keys()):
            info = self.finished.pop(rid)
            # 找到该 request 所在 session 的 sessionId（任意已 attach session 均可试）
            for sid in self.sessions:
                mid = await self.send("Network.getResponseBody", {"requestId": rid}, session_id=sid)
                self.pending[mid] = info
                break

    async def save_body(self, msg, info):
        url, mime = info
        result = msg.get("result", {})
        body = result.get("body", "")
        if not body:
            return
        OUTDIR.mkdir(parents=True, exist_ok=True)
        ts = time.strftime("%H%M%S")
        if result.get("base64Encoded"):
            # 二进制：按 mime 保存图片
            ext = IMAGE_MIME.get(mime, "bin")
            self.n_saved += 1
            path = OUTDIR / f"net_{ts}_{self.n_saved}.{ext}"
            path.write_bytes(base64.b64decode(body))
            print(f"[save] {path.name} ({len(body)} b64, {mime}) <- {url[:80]}", flush=True)
        else:
            # 文本：captcha/get + captcha/verify 的 JSON 响应全存
            if "captcha/get" in url or "captcha/verify" in url or (body.startswith("{") and ("question" in body or "captcha" in body)):
                self.n_saved += 1
                tag = "get" if "captcha/get" in url else "verify"
                path = OUTDIR / f"net_{ts}_{tag}.json"
                path.write_text(body, encoding="utf-8")
                print(f"[save] {path.name} ({len(body)} chars) <- {url[:80]}", flush=True)
                print(f"       head: {body[:200]}", flush=True)


async def main(minutes: float):
    import asyncio
    try:
        ver = json.load(urllib.request.urlopen(f"{CDP_HTTP}/json/version", timeout=5))
    except (OSError, ValueError) as e:
        print(f"无法连接 CDP {CDP_HTTP}（Chrome 需带 --remote-debugging-port=9222 启动）: {e}", file=sys.stderr)
        sys.exit(1)
    ws_url = ver["webSocketDebuggerUrl"]
    deadline = time.time() + minutes * 60
    async with websockets.connect(ws_url, max_size=64 * 1024 * 1024) as ws:
        col = Collector(ws)
        # 浏览器级 autoAttach：所有新 target（含未来弹出的验证码 iframe）自动 attach
        await col.send("Target.setAutoAttach", {
            "autoAttach": True, "waitForDebuggerOnStart": False, "flatten": True,
        })
        # 补充 attach 已存在的 target（验证码 iframe 可能已经弹出）
        try:
            targets = json.load(urllib.request.urlopen(f"{CDP_HTTP}/json/list", timeout=5))
        except (OSError, ValueError):
            targets = []
        for t in targets:
            if t.get("type") in ("page", "iframe", "webview") and t.get("id"):
                await col.send("Target.attachToTarget", {"targetId": t["id"], "flatten": True})
        print(f"[listen] auto-attach + 已存在 target attach({len(targets)})，监听 {minutes} 分钟，Ctrl+C 退出", flush=True)
        await col.recv_loop(deadline)
    print(f"[done] 共保存 {col.n_saved} 个文件到 {OUTDIR}", flush=True)


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--minutes", type=float, default=10)
    args = ap.parse_args()
    import asyncio
    asyncio.run(main(args.minutes))