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

#!/usr/bin/env python3
"""Batch HTTP fetch + trafilatura extract. Writes raw JSON/HTML heads. No secrets."""
from __future__ import annotations

import argparse
import gzip
import json
import ssl
import sys
import time
import traceback
from pathlib import Path
from urllib.error import HTTPError, URLError
from urllib.parse import urljoin, urlparse
from urllib.request import (
    HTTPRedirectHandler,
    ProxyHandler,
    Request,
    build_opener,
    urlopen,
)
from html.parser import HTMLParser
import zlib

import trafilatura
from trafilatura.metadata import extract_metadata
from trafilatura.settings import DEFAULT_CONFIG
from copy import deepcopy

from classify import classify_intercept, extractable, title_ok

TIMEOUT_S = 20
MAX_BYTES = 2_000_000
UA = (
    "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 "
    "(KHTML, like Gecko) Chrome/151.0.7922.109 Safari/537.36"
)

_CFG = deepcopy(DEFAULT_CONFIG)
_CFG["DEFAULT"]["DOWNLOAD_TIMEOUT"] = str(TIMEOUT_S)
_CFG["DEFAULT"]["MAX_FILE_SIZE"] = str(MAX_BYTES)


class _HrefParser(HTMLParser):
    def __init__(self) -> None:
        super().__init__()
        self.hrefs: list[str] = []

    def handle_starttag(self, tag, attrs):
        if tag != "a":
            return
        for k, v in attrs:
            if k == "href" and v:
                self.hrefs.append(v)


class _Tracker(HTTPRedirectHandler):
    def __init__(self):
        super().__init__()
        self.chain: list[dict] = []

    def redirect_request(self, req, fp, code, msg, headers, newurl):
        self.chain.append({"status": int(code), "from": req.full_url, "to": newurl})
        return super().redirect_request(req, fp, code, msg, headers, newurl)


def _decode(raw: bytes, headers: dict) -> str:
    enc = (headers.get("Content-Encoding") or headers.get("content-encoding") or "").lower()
    body = raw
    try:
        if "gzip" in enc:
            body = gzip.decompress(raw)
        elif "deflate" in enc:
            body = zlib.decompress(raw, -zlib.MAX_WBITS)
    except Exception:
        body = raw
    ctype = headers.get("Content-Type") or headers.get("content-type") or ""
    charset = "utf-8"
    if "charset=" in ctype.lower():
        charset = ctype.lower().split("charset=", 1)[1].split(";")[0].strip().strip('"')
    return body.decode(charset, errors="replace")


def fetch(url: str, timeout: int, proxy: str | None) -> dict:
    tracker = _Tracker()
    handlers = [tracker]
    if proxy:
        handlers.insert(0, ProxyHandler({"http": proxy, "https": proxy}))
    opener = build_opener(*handlers)
    req = Request(
        url,
        headers={
            "User-Agent": UA,
            "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8",
            "Accept-Language": "en-US,en;q=0.9,zh-CN;q=0.8",
            "Accept-Encoding": "gzip, deflate",
        },
        method="GET",
    )
    ctx = ssl.create_default_context()
    try:
        resp = opener.open(req, timeout=timeout)
        raw = resp.read(MAX_BYTES + 1)
        headers = {k: v for k, v in resp.headers.items()}
        html = _decode(raw[:MAX_BYTES], headers)
        return {
            "http_status": int(resp.status),
            "final_url": resp.geturl() or url,
            "headers": headers,
            "html": html,
            "redirect_chain": tracker.chain,
            "error": None,
        }
    except HTTPError as exc:
        raw = b""
        try:
            raw = exc.read(MAX_BYTES)
        except Exception:
            raw = b""
        headers = {k: v for k, v in (exc.headers.items() if exc.headers else [])}
        html = _decode(raw, headers) if raw else ""
        return {
            "http_status": int(exc.code),
            "final_url": exc.geturl() if hasattr(exc, "geturl") and exc.geturl() else url,
            "headers": headers,
            "html": html,
            "redirect_chain": tracker.chain,
            "error": f"HTTPError {exc.code}",
        }
    except URLError as exc:
        return {
            "http_status": 0,
            "final_url": url,
            "headers": {},
            "html": "",
            "redirect_chain": tracker.chain,
            "error": f"URLError {exc.reason}",
        }
    except Exception as exc:  # noqa: BLE001
        return {
            "http_status": 0,
            "final_url": url,
            "headers": {},
            "html": "",
            "redirect_chain": tracker.chain,
            "error": f"{type(exc).__name__}: {exc}",
        }
    finally:
        try:
            ctx  # noqa: B018
        except Exception:
            pass


def discover_article(listing_url: str, patterns: list[str], proxy: str | None) -> str | None:
    got = fetch(listing_url, TIMEOUT_S, proxy)
    html = got.get("html") or ""
    if not html:
        return None
    p = _HrefParser()
    try:
        p.feed(html)
    except Exception:
        return None
    base = got.get("final_url") or listing_url
    for href in p.hrefs:
        abs_url = urljoin(base, href)
        if abs_url.split("#", 1)[0] == listing_url.rstrip("/"):
            continue
        path = urlparse(abs_url).path or ""
        if any(pat in abs_url or pat in path for pat in patterns):
            # skip assets / same listing
            if any(abs_url.endswith(ext) for ext in (".css", ".js", ".png", ".jpg", ".svg", ".ico")):
                continue
            if "javascript:" in abs_url:
                continue
            return abs_url.split("#", 1)[0]
    return None


def extract_text(html: str, url: str) -> tuple[str, str]:
    if not html:
        return "", ""
    md = trafilatura.extract(
        html,
        output_format="markdown",
        include_comments=False,
        include_tables=True,
        favor_precision=True,
        config=_CFG,
        url=url,
    ) or ""
    title = ""
    try:
        meta = extract_metadata(html, default_url=url)
        title = (meta.title if meta and getattr(meta, "title", None) else "") or ""
    except Exception:
        title = ""
    if not title:
        m = __import__("re").search(r"<title[^>]*>(.*?)</title>", html, __import__("re").I | __import__("re").S)
        if m:
            title = __import__("re").sub(r"\s+", " ", m.group(1)).strip()
    return title, md


def run_one(site: dict, proxy: str | None, raw_dir: Path) -> dict:
    url = site["url"]
    t0 = time.perf_counter()
    discovered = None
    if site.get("discover_from") and site.get("article_patterns"):
        discovered = discover_article(site["discover_from"], site["article_patterns"], proxy)
        if discovered:
            url = discovered
    got = fetch(url, TIMEOUT_S, proxy)
    title, markdown = extract_text(got.get("html") or "", got.get("final_url") or url)
    html = got.get("html") or ""
    html_head = html[:16000]
    text_len = len(markdown)
    hops = len(got.get("redirect_chain") or [])
    intercept = classify_intercept(
        http_status=got.get("http_status") or 0,
        final_url=got.get("final_url") or url,
        start_url=url,
        title=title,
        text=markdown,
        html_head=html_head,
        headers=got.get("headers") or {},
        error=got.get("error"),
        redirect_hops=hops,
    )
    tok = title_ok(title, site.get("title_needles") or [], intercept)
    # if title missed needles but page text has them and intercept is none, still mark title from extract meta
    if not tok and intercept == "none":
        blob = (title + "\n" + markdown[:1500]).lower()
        tok = any(n.lower() in blob for n in (site.get("title_needles") or []))
    elapsed = round(time.perf_counter() - t0, 3)
    rec = {
        "id": site["id"],
        "slug": site["slug"],
        "group": site["group"],
        "channel": "trafilatura-http",
        "start_url": site["url"],
        "used_url": url,
        "discovered_url": discovered,
        "http_status": got.get("http_status") or 0,
        "final_url": got.get("final_url") or url,
        "title": (title or "")[:240],
        "title_ok": bool(tok),
        "text_len": text_len,
        "html_len": len(html),
        "extractable": extractable(intercept, text_len, tok),
        "intercept": intercept,
        "elapsed_s": elapsed,
        "redirect_chain": got.get("redirect_chain") or [],
        "error": got.get("error"),
        "header_keys": sorted((got.get("headers") or {}).keys()),
        "text_head": markdown[:1200],
        "scenario": site.get("scenario"),
        "known_risk": site.get("known_risk"),
    }
    slug = site["slug"]
    (raw_dir / f"http-{site['id']:02d}-{slug}.json").write_text(
        json.dumps(rec, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
    )
    (raw_dir / f"http-{site['id']:02d}-{slug}.html.head").write_text(html_head, encoding="utf-8", errors="replace")
    return rec


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--targets", required=True)
    ap.add_argument("--group", required=True, choices=["foreign", "domestic"])
    ap.add_argument("--out", required=True)
    ap.add_argument("--proxy", default="")
    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]
    proxy = args.proxy or None
    print(f"[http_probe] group={args.group} n={len(sites)} proxy={'yes' if proxy else 'no'}", flush=True)
    rows = []
    for site in sites:
        print(f"[http_probe] start id={site['id']} {site['slug']} {site['url']}", flush=True)
        try:
            rec = run_one(site, proxy, raw_dir)
        except Exception as exc:  # noqa: BLE001
            traceback.print_exc()
            rec = {
                "id": site["id"],
                "slug": site["slug"],
                "group": site["group"],
                "channel": "trafilatura-http",
                "start_url": site["url"],
                "used_url": site["url"],
                "http_status": 0,
                "final_url": site["url"],
                "title": "",
                "title_ok": False,
                "text_len": 0,
                "html_len": 0,
                "extractable": False,
                "intercept": "empty",
                "elapsed_s": 0,
                "error": f"{type(exc).__name__}: {exc}",
            }
            (raw_dir / f"http-{site['id']:02d}-{site['slug']}.json").write_text(
                json.dumps(rec, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
            )
        rows.append(rec)
        print(
            f"[http_probe] done id={site['id']} status={rec.get('http_status')} "
            f"intercept={rec.get('intercept')} text_len={rec.get('text_len')} t={rec.get('elapsed_s')}",
            flush=True,
        )
    outp = raw_dir / f"http-{args.group}.jsonl"
    outp.write_text("".join(json.dumps(r, ensure_ascii=False) + "\n" for r in rows), encoding="utf-8")
    # resolved urls for the other channel
    resolved = {r["slug"]: r.get("used_url") or r.get("start_url") for r in rows}
    rp = raw_dir / f"resolved-{args.group}.json"
    rp.write_text(json.dumps(resolved, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    print(f"[http_probe] wrote {outp}", flush=True)
    return 0


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