scripts/download_web_collections.py

#!/usr/bin/env python3
"""Download bounded document/code collections from approved author/university pages."""

from __future__ import annotations

import argparse
import hashlib
import json
import re
import ssl
import subprocess
import urllib.error
import urllib.parse
import urllib.request
from datetime import datetime, timezone
from html.parser import HTMLParser
from pathlib import Path

import certifi


PROJECT_ROOT = Path(__file__).resolve().parents[1]
DEFAULT_MANIFEST = PROJECT_ROOT / "data" / "sources" / "web_collections_manifest.json"
DEFAULT_OUTPUT = PROJECT_ROOT / "data" / "raw" / "web_collections"
TLS_CONTEXT = ssl.create_default_context(cafile=certifi.where())
USER_AGENT = "controlai-source-fetcher/0.1"
ALLOWED_SUFFIXES = {".pdf", ".m", ".py", ".ipynb", ".zip"}


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

    def handle_starttag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None:
        if tag.lower() != "a":
            return
        href = dict(attrs).get("href")
        if href:
            self.links.append(href)


def fetch_bytes(url: str) -> bytes:
    request = urllib.request.Request(url, headers={"User-Agent": USER_AGENT})
    try:
        with urllib.request.urlopen(request, context=TLS_CONTEXT) as response:
            return response.read()
    except urllib.error.URLError:
        result = subprocess.run(
            [
                "curl", "--fail", "--location", "--silent", "--show-error",
                "--proto", "=https", "--user-agent", USER_AGENT, url,
            ],
            check=True,
            stdout=subprocess.PIPE,
        )
        return result.stdout


def discover_files(collection: dict) -> list[str]:
    if collection.get("files"):
        return collection["files"]
    html = fetch_bytes(collection["landing_url"]).decode("utf-8", errors="replace")
    parser = LinkParser()
    parser.feed(html)
    pattern = re.compile(collection["link_pattern"], re.IGNORECASE)
    return sorted(dict.fromkeys(link for link in parser.links if pattern.search(link)))


def safe_url(base_url: str, link: str) -> str:
    url = urllib.parse.urljoin(base_url, link)
    base = urllib.parse.urlparse(base_url)
    parsed = urllib.parse.urlparse(url)
    if parsed.scheme != "https" or parsed.hostname != base.hostname:
        raise ValueError(f"Link escaped approved HTTPS host: {url}")
    if Path(parsed.path).suffix.lower() not in ALLOWED_SUFFIXES:
        raise ValueError(f"Unsupported collection file type: {url}")
    return url


def slug(text: str) -> str:
    value = re.sub(r"[^a-z0-9]+", "_", text.lower()).strip("_")
    return value[:100]


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as stream:
        for block in iter(lambda: stream.read(1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest()


def download_file(url: str, destination: Path) -> None:
    if destination.exists():
        return
    data = fetch_bytes(url)
    if destination.suffix.lower() == ".pdf" and not data.startswith(b"%PDF-"):
        raise ValueError(f"Expected PDF but received different content: {url}")
    partial = destination.with_suffix(destination.suffix + ".part")
    partial.write_bytes(data)
    partial.replace(destination)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--manifest", type=Path, default=DEFAULT_MANIFEST)
    parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT)
    parser.add_argument("--collection", action="append", default=[])
    parser.add_argument("--download", action="store_true")
    args = parser.parse_args()

    manifest = json.loads(args.manifest.read_text(encoding="utf-8"))
    collections = manifest["collections"]
    if args.collection:
        requested = set(args.collection)
        collections = [item for item in collections if item["collection_id"] in requested]
        missing = requested - {item["collection_id"] for item in collections}
        if missing:
            raise ValueError(f"Unknown collection: {', '.join(sorted(missing))}")
    else:
        collections = [item for item in collections if item.get("enabled", True)]

    args.output.mkdir(parents=True, exist_ok=True)
    lock_path = args.output / "source_lock.json"
    existing = json.loads(lock_path.read_text(encoding="utf-8")) if lock_path.exists() else []
    records = {item["source_id"]: item for item in existing}
    failures = []

    for collection in collections:
        files = discover_files(collection)
        expected = collection.get("expected_count")
        if expected is not None and len(files) != expected:
            raise ValueError(
                f"{collection['collection_id']}: expected {expected} files, discovered {len(files)}"
            )
        print(f"{collection['collection_id']}: {len(files)} approved files")
        for link in files:
            url = safe_url(collection["base_url"], link)
            original_name = Path(urllib.parse.urlparse(url).path).name
            filename = f"{collection['collection_id']}__{original_name}"
            source_id = f"{collection['collection_id']}__{slug(original_name)}"
            destination = args.output / filename
            print(f"- {url}")
            if not args.download:
                continue
            try:
                download_file(url, destination)
            except (OSError, ValueError, subprocess.CalledProcessError) as error:
                failures.append(
                    {
                        "collection_id": collection["collection_id"],
                        "url": url,
                        "error": str(error),
                    }
                )
                print(f"  [failed] {error}")
                continue
            records[source_id] = {
                "source_id": source_id,
                "collection_id": collection["collection_id"],
                "title": collection["title"],
                "authors": collection.get("authors", []),
                "filename": filename,
                "url": url,
                "landing_url": collection["landing_url"],
                "license": collection["license"],
                "corpus_tier": collection.get("corpus_tier"),
                "coverage": collection.get("coverage", []),
                "bytes": destination.stat().st_size,
                "sha256": sha256(destination),
                "downloaded_at": datetime.now(timezone.utc).isoformat(),
            }

    if args.download:
        lock_path.write_text(
            json.dumps([records[key] for key in sorted(records)], indent=2) + "\n",
            encoding="utf-8",
        )
        print(f"Wrote checksums and provenance: {lock_path}")
        failure_path = args.output / "download_failures.json"
        failure_path.write_text(json.dumps(failures, indent=2) + "\n", encoding="utf-8")
        print(f"Failures: {len(failures)} ({failure_path})")
        if failures:
            raise SystemExit(1)


if __name__ == "__main__":
    main()