onesvm-browser-server/bench/site-matrix/lp_probe.py

#!/usr/bin/env python3
"""Drive lightpanda CDP fetches from the host. Writes raw JSON. No secrets."""
from __future__ import annotations

import argparse
import json
import subprocess
import sys
import time
from pathlib import Path

from classify import classify_intercept, extractable, title_ok


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--targets", required=True)
    ap.add_argument("--group", required=True, choices=["foreign", "domestic"])
    ap.add_argument("--cdp", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--cdp-fetch", required=True)
    ap.add_argument("--resolved", default="")
    ap.add_argument("--timeout", type=int, default=20)
    ap.add_argument("--slug", default="")
    args = ap.parse_args()

    raw_dir = Path(args.out)
    raw_dir.mkdir(parents=True, exist_ok=True)
    data = json.loads(Path(args.targets).read_text(encoding="utf-8"))
    sites = [s for s in data["sites"] if s["group"] == args.group]
    if args.slug:
        sites = [s for s in sites if s["slug"] == args.slug]
    resolved = {}
    if args.resolved and Path(args.resolved).exists():
        resolved = json.loads(Path(args.resolved).read_text(encoding="utf-8"))

    rows = []
    for site in sites:
        url = resolved.get(site["slug"]) or site["url"]
        print(f"[lp_probe] start id={site['id']} {site['slug']} {url}", flush=True)
        t0 = time.perf_counter()
        try:
            proc = subprocess.run(
                ["node", args.cdp_fetch, args.cdp, url, str(args.timeout)],
                capture_output=True,
                text=True,
                timeout=args.timeout + 15,
            )
            line = (proc.stdout or "").strip().splitlines()[-1] if proc.stdout.strip() else "{}"
            d = json.loads(line)
        except subprocess.TimeoutExpired:
            d = {"ok": False, "error": "subprocess timeout", "http_status": 0, "title": "", "text": "", "textLen": 0}
        except Exception as exc:  # noqa: BLE001
            d = {"ok": False, "error": f"{type(exc).__name__}: {exc}", "http_status": 0, "title": "", "text": "", "textLen": 0}
        elapsed = round(time.perf_counter() - t0, 3)
        title = d.get("title") or ""
        text = d.get("text") or ""
        html_head = d.get("htmlHead") or ""
        status = int(d.get("http_status") or 0)
        final_url = d.get("finalUrl") or url
        hops = len(d.get("redirect_chain") or [])
        intercept = classify_intercept(
            http_status=status,
            final_url=final_url,
            start_url=url,
            title=title,
            text=text,
            html_head=html_head,
            headers={},
            error=d.get("error") or d.get("errorText"),
            redirect_hops=hops,
        )
        # prefer CDP vendor when it is more specific
        vendor = d.get("challenge_vendor") or intercept
        if vendor in {"cloudflare", "waf", "captcha", "paywall", "redirect", "empty"} and intercept == "none":
            intercept = vendor
        if vendor == "cloudflare":
            intercept = "cloudflare"
        tok = title_ok(title, site.get("title_needles") or [], intercept)
        if not tok and intercept == "none":
            blob = (title + "\n" + text[:1500]).lower()
            tok = any(n.lower() in blob for n in (site.get("title_needles") or []))
        text_len = int(d.get("textLen") or len(text))
        rec = {
            "id": site["id"],
            "slug": site["slug"],
            "group": site["group"],
            "channel": "lightpanda",
            "start_url": site["url"],
            "used_url": url,
            "http_status": status,
            "final_url": final_url,
            "title": title[:240],
            "title_ok": bool(tok),
            "text_len": text_len,
            "html_len": int(d.get("htmlLen") or 0),
            "extractable": extractable(intercept, text_len, tok),
            "intercept": intercept,
            "elapsed_s": elapsed,
            "redirect_chain": d.get("redirect_chain") or [],
            "error": d.get("error") or d.get("errorText"),
            "text_head": text[:1200],
            "cdp_ok": bool(d.get("ok")),
            "cdp_errors": d.get("cdp_errors") or [],
            "loadFired": d.get("loadFired"),
            "scenario": site.get("scenario"),
            "known_risk": site.get("known_risk"),
        }
        (raw_dir / f"lp-{site['id']:02d}-{site['slug']}.json").write_text(
            json.dumps(rec, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
        )
        excerpt = {
            "title": title[:200],
            "final_url": final_url,
            "http_status": status,
            "html_head": html_head[:2000],
            "text_head": text[:1500],
            "cdp": {k: d.get(k) for k in ("ok", "error", "errorText", "cdp_mode", "challenge_vendor", "header_bits")},
        }
        (raw_dir / f"lp-{site['id']:02d}-{site['slug']}.excerpt.json").write_text(
            json.dumps(excerpt, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
        )
        rows.append(rec)
        print(
            f"[lp_probe] done id={site['id']} status={status} intercept={intercept} "
            f"text_len={text_len} t={elapsed}",
            flush=True,
        )

    outp = raw_dir / f"lp-{args.group}.jsonl"
    outp.write_text("".join(json.dumps(r, ensure_ascii=False) + "\n" for r in rows), encoding="utf-8")
    print(f"[lp_probe] wrote {outp}", flush=True)
    return 0


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