Files
pdf/scripts/fetch_corpus.py
2026-06-22 15:18:47 +05:30

127 lines
4.7 KiB
Python

"""Download a large PDF corpus for fuzzing into a gitignored directory.
The committed `corpus/` holds a small, curated set of fixtures used by the
rendering-regression baseline. Fuzzing wants *volume and variety* instead, so we
pull hundreds of real-world PDFs from public test suites into `corpus/fuzz/`,
which is gitignored.
python scripts/fetch_corpus.py # default: ~600 from pdf.js
python scripts/fetch_corpus.py --limit 200
python scripts/fetch_corpus.py --source pdfium # GoogleTest pdfium corpus
Network failures are tolerated: whatever downloads is usable, and re-running only
fetches what's missing. Only standard-library modules are used.
"""
from __future__ import annotations
import argparse
import hashlib
import json
import sys
import urllib.error
import urllib.request
from pathlib import Path
_RAW_BASE = "https://raw.githubusercontent.com/mozilla/pdf.js/master/test/pdfs"
ROOT = Path(__file__).resolve().parents[1]
DEST = ROOT / "corpus" / "fuzz"
SOURCES = {
"pdfjs": "https://api.github.com/repos/mozilla/pdf.js/contents/test/pdfs?ref=master",
"pdfium": "https://api.github.com/repos/PDFium/pdfium/contents/testing/resources?ref=main",
}
_HEADERS = {"User-Agent": "pdfengine-fetch-corpus", "Accept": "application/vnd.github+json"}
def _get(url: str, raw: bool = False) -> bytes:
req = urllib.request.Request(url, headers=_HEADERS)
with urllib.request.urlopen(req, timeout=30) as resp:
return resp.read()
def _fetch_from_manifest(manifest_path: Path, dest: Path) -> int:
"""Download exactly the files in the pinned manifest and verify their SHA-256,
so every developer and CI run gets a byte-identical corpus."""
manifest = json.loads(manifest_path.read_text())
files = manifest.get("files", [])
print(f"Manifest: {len(files)} pinned files -> {dest} (verifying SHA-256)")
got = existed = failed = mismatch = 0
for entry in files:
name, want = entry["name"], entry["sha256"]
out = dest / name
if out.exists() and hashlib.sha256(out.read_bytes()).hexdigest() == want:
existed += 1
continue
try:
data = _get(f"{_RAW_BASE}/{name}", raw=True)
except (urllib.error.URLError, OSError) as exc:
print(f" fail {name} ({exc})")
failed += 1
continue
have = hashlib.sha256(data).hexdigest()
if have != want:
print(f" MISMATCH {name}: expected {want[:12]}..., got {have[:12]}... (skipped)")
mismatch += 1
continue
out.write_bytes(data)
got += 1
if got % 50 == 0:
print(f" ... {got} verified")
total = len(list(dest.glob("*.pdf")))
print(f"\nDone. +{got} new, {existed} already present & verified, "
f"{failed} failed, {mismatch} hash-mismatch.")
print(f"Corpus now holds {total} PDFs at {dest}")
return 0 if mismatch == 0 else 1
def main() -> int:
ap = argparse.ArgumentParser(description=__doc__)
ap.add_argument("--source", choices=sorted(SOURCES), default="pdfjs")
ap.add_argument("--limit", type=int, default=600, help="max files to download")
ap.add_argument("--dest", type=Path, default=DEST)
ap.add_argument("--manifest", type=Path, default=None,
help="reproduce the exact pinned corpus from a manifest (verifies SHA-256)")
args = ap.parse_args()
args.dest.mkdir(parents=True, exist_ok=True)
if args.manifest:
return _fetch_from_manifest(args.manifest, args.dest)
print(f"Listing {args.source} corpus ...")
try:
listing = json.loads(_get(SOURCES[args.source]))
except (urllib.error.URLError, json.JSONDecodeError) as exc:
print(f"ERROR: could not list corpus ({exc}). Check your network / GitHub rate limit.")
return 2
pdfs = [e for e in listing if e.get("name", "").lower().endswith(".pdf") and e.get("download_url")]
print(f"Found {len(pdfs)} PDFs; downloading up to {args.limit} into {args.dest} ...")
got = failed = existed = 0
for entry in pdfs[: args.limit]:
out = args.dest / entry["name"]
if out.exists() and out.stat().st_size > 0:
existed += 1
continue
try:
out.write_bytes(_get(entry["download_url"], raw=True))
got += 1
if got % 25 == 0:
print(f" ... {got} downloaded")
except (urllib.error.URLError, OSError) as exc:
print(f" fail {entry['name']} ({exc})")
failed += 1
total = len(list(args.dest.glob("*.pdf")))
print(f"\nDone. +{got} new, {existed} already present, {failed} failed.")
print(f"Corpus now holds {total} PDFs at {args.dest}")
return 0
if __name__ == "__main__":
sys.exit(main())