scripts/fetch_extension.py

"""
Downloads the Keyword Surfer Chrome extension (.crx) from the Chrome Web Store
and unpacks it into extension/ so Playwright can load it with --load-extension.

Run at Docker build time (see Dockerfile). Failure here is non-fatal — the app
falls back to the Autocomplete/pytrends/LLM-estimate chain (modules/keywords.py)
if the extension isn't present.
"""

from __future__ import annotations

import argparse
import io
import struct
import zipfile
from pathlib import Path
from urllib.parse import urlencode

import httpx

# Verified against the live Chrome Web Store listing:
# https://chromewebstore.google.com/detail/keyword-surfer/bafijghppfhdpldihckdcadbcobikaca
KEYWORD_SURFER_EXTENSION_ID = "bafijghppfhdpldihckdcadbcobikaca"

CRX_DOWNLOAD_URL = "https://clients2.google.com/service/update2/crx"


def download_crx(extension_id: str) -> bytes:
    params = {
        "response": "redirect",
        "prodversion": "120.0.0.0",
        "acceptformat": "crx2,crx3",
        "x": f"id={extension_id}&uc",
    }
    url = f"{CRX_DOWNLOAD_URL}?{urlencode(params)}"
    resp = httpx.get(url, follow_redirects=True, timeout=30)
    resp.raise_for_status()
    return resp.content


def crx_to_zip_bytes(crx_bytes: bytes) -> bytes:
    """Strips the CRX2/CRX3 header to yield the raw ZIP payload."""
    magic = crx_bytes[0:4]
    if magic != b"Cr24":
        raise ValueError("Not a valid CRX file (bad magic number)")
    version = struct.unpack("<I", crx_bytes[4:8])[0]
    if version == 3:
        header_size = struct.unpack("<I", crx_bytes[8:12])[0]
        zip_start = 12 + header_size
    elif version == 2:
        pubkey_len, sig_len = struct.unpack("<II", crx_bytes[8:16])
        zip_start = 16 + pubkey_len + sig_len
    else:
        raise ValueError(f"Unsupported CRX version: {version}")
    return crx_bytes[zip_start:]


def unpack_extension(extension_id: str, out_dir: Path) -> None:
    print(f"Downloading extension {extension_id}...")
    crx_bytes = download_crx(extension_id)
    zip_bytes = crx_to_zip_bytes(crx_bytes)

    out_dir.mkdir(parents=True, exist_ok=True)
    with zipfile.ZipFile(io.BytesIO(zip_bytes)) as zf:
        zf.extractall(out_dir)
    print(f"Unpacked extension into {out_dir}")


def main():
    parser = argparse.ArgumentParser(description="Fetch and unpack the Keyword Surfer extension")
    parser.add_argument("--extension-id", default=KEYWORD_SURFER_EXTENSION_ID)
    parser.add_argument("--out", default="extension")
    args = parser.parse_args()

    try:
        unpack_extension(args.extension_id, Path(args.out))
    except Exception as exc:
        print(f"WARNING: failed to fetch Keyword Surfer extension: {exc}")


if __name__ == "__main__":
    main()