spectre_pdf 1.0.0

Native Rust PDF extraction engine: text, markdown for RAG, AcroForm widgets, image decoding, and encrypted PDFs. Lazy parser, persistent Document handle, no C dependencies.
Documentation
"""
Benchmark Document.open_with_password() against pymupdf authenticate.

Tests every PDF Standard Security Handler variant on a fixed corpus
generated via qpdf, comparing wall-clock decrypt-and-open + text-extract
time and verifying the extracted text exactly matches the unencrypted
baseline.

Usage:
    pip install pymupdf spectre_rs
    python scripts/bench_encrypted_vs_pymupdf.py [--corpus DIR]

The corpus is generated on first run via qpdf if not present.
"""

from __future__ import annotations

import argparse
import os
import statistics
import subprocess
import time
import urllib.request
from pathlib import Path

import fitz
import spectre_rs


SOURCE_URL = "https://www.irs.gov/pub/irs-pdf/fw9.pdf"

VARIANTS = [
    ("v1_rc4_40.pdf", b"secret", "V=1 R=2", "RC4 40-bit"),
    ("v2_rc4_128.pdf", b"secret", "V=2 R=3", "RC4 128-bit"),
    ("v4_aes_128.pdf", b"secret", "V=4 R=4", "AES-128 CBC"),
    ("v4_aes_128_empty_user.pdf", b"", "V=4 R=4", "AES-128, empty user"),
    ("v5_aes_256_r5.pdf", b"secret", "V=5 R=5", "AES-256 (Acrobat 9, withdrawn)"),
    ("v5_aes_256.pdf", b"secret", "V=5 R=6", "AES-256 (Alg 2.B)"),
    ("v5_aes_256_empty_user.pdf", b"", "V=5 R=6", "AES-256, empty user"),
]


def ensure_corpus(corpus_dir: Path) -> None:
    corpus_dir.mkdir(parents=True, exist_ok=True)
    source = corpus_dir / "source.pdf"
    if not source.exists():
        print("fetch IRS W-9 source PDF")
        urllib.request.urlretrieve(SOURCE_URL, source)
    if not (corpus_dir / "v5_aes_256.pdf").exists():
        if subprocess.run(["which", "qpdf"], capture_output=True).returncode != 0:
            raise SystemExit(
                "qpdf is required to generate the encrypted corpus. "
                "Install via `apt install qpdf` and re-run."
            )
        cmds = [
            ("v1_rc4_40.pdf",
             ["qpdf", "--allow-weak-crypto", "--encrypt", "secret", "secret",
              "40", "--", str(source), str(corpus_dir / "v1_rc4_40.pdf")]),
            ("v2_rc4_128.pdf",
             ["qpdf", "--allow-weak-crypto", "--encrypt", "secret", "secret",
              "128", "--use-aes=n", "--", str(source),
              str(corpus_dir / "v2_rc4_128.pdf")]),
            ("v4_aes_128.pdf",
             ["qpdf", "--encrypt", "secret", "secret", "128",
              "--use-aes=y", "--", str(source),
              str(corpus_dir / "v4_aes_128.pdf")]),
            ("v4_aes_128_empty_user.pdf",
             ["qpdf", "--encrypt", "", "secret", "128", "--use-aes=y",
              "--", str(source), str(corpus_dir / "v4_aes_128_empty_user.pdf")]),
            ("v5_aes_256_r5.pdf",
             ["qpdf", "--encrypt", "secret", "secret", "256", "--force-R5",
              "--", str(source), str(corpus_dir / "v5_aes_256_r5.pdf")]),
            ("v5_aes_256.pdf",
             ["qpdf", "--encrypt", "secret", "secret", "256", "--",
              str(source), str(corpus_dir / "v5_aes_256.pdf")]),
            ("v5_aes_256_empty_user.pdf",
             ["qpdf", "--encrypt", "", "secret", "256", "--",
              str(source), str(corpus_dir / "v5_aes_256_empty_user.pdf")]),
        ]
        for name, cmd in cmds:
            print(f"generate {name}")
            subprocess.run(cmd, check=True)


def time_decrypt_and_extract_spectre(b: bytes, pwd: bytes, n: int = 3) -> tuple[float, str]:
    text_out = ""
    runs = []
    for _ in range(n):
        t0 = time.perf_counter()
        d = spectre_rs.Document(b, password=pwd)
        text_out = d.text()
        runs.append((time.perf_counter() - t0) * 1000)
    return statistics.median(runs), text_out


def time_decrypt_and_extract_pymupdf(b: bytes, pwd: bytes, n: int = 3) -> tuple[float, str]:
    text_out = ""
    runs = []
    for _ in range(n):
        t0 = time.perf_counter()
        d = fitz.open(stream=b)
        d.authenticate(pwd.decode("latin-1"))
        text_out = "".join(p.get_text() for p in d).strip()
        runs.append((time.perf_counter() - t0) * 1000)
        d.close()
    return statistics.median(runs), text_out


def bench(corpus_dir: Path) -> None:
    source_bytes = (corpus_dir / "source.pdf").read_bytes()
    baseline = spectre_rs.Document(source_bytes).text().strip()
    print(f"plaintext baseline: {len(baseline)} chars")
    print("=" * 110)
    print(f"{'variant':<22} {'cipher':<25} {'pmu ms':>8} {'spec ms':>9} "
          f"{'speedup':>9} {'text match':>11}")
    print("-" * 110)
    total_s = total_p = 0.0
    matches_spec = matches_pmu = 0
    for name, pwd, version, label in VARIANTS:
        path = corpus_dir / name
        if not path.exists() or path.stat().st_size == 0:
            print(f"  {name}: SKIPPED (file missing/empty)")
            continue
        b = path.read_bytes()
        spec_ms, spec_text = time_decrypt_and_extract_spectre(b, pwd)
        pmu_ms, pmu_text = time_decrypt_and_extract_pymupdf(b, pwd)
        total_s += spec_ms
        total_p += pmu_ms
        spec_ok = spec_text.strip() == baseline
        pmu_ok = pmu_text.strip() == baseline
        matches_spec += int(spec_ok)
        matches_pmu += int(pmu_ok)
        speedup = pmu_ms / spec_ms if spec_ms else float("inf")
        tag_spec = "OK" if spec_ok else "DIFF"
        print(f"{version:<22} {label:<25} {pmu_ms:>8.2f} {spec_ms:>9.2f} "
              f"{speedup:>8.2f}x {tag_spec:>11}")
    print("-" * 110)
    print(f"TOTAL  speedup {total_p / total_s:.2f}x "
          f"({total_p:.1f}ms / {total_s:.1f}ms)")
    print(f"  spectre baseline-text match:  {matches_spec}/{len(VARIANTS)}")
    print(f"  pymupdf baseline-text match:  {matches_pmu}/{len(VARIANTS)}")

    # Negative test — wrong password should error cleanly on both.
    print()
    print("Wrong-password test (expect both libraries to reject):")
    for name in ["v4_aes_128.pdf", "v5_aes_256.pdf"]:
        path = corpus_dir / name
        if not path.exists():
            continue
        b = path.read_bytes()
        try:
            spectre_rs.Document(b, password=b"wrong-password")
            print(f"  {name}: spectre INCORRECTLY ACCEPTED wrong password")
        except Exception as e:
            print(f"  {name}: spectre rejected ({type(e).__name__})")
        d = fitz.open(stream=b)
        rc = d.authenticate("wrong-password")
        print(f"  {name}: pymupdf rc={rc}")
        d.close()


def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--corpus", default="bench-corpus/encrypted")
    args = p.parse_args()
    corpus = Path(args.corpus)
    ensure_corpus(corpus)
    bench(corpus)


if __name__ == "__main__":
    main()