util_zip.py

"""Safe extraction of an uploaded SERFF filing ZIP.

Hardened for a public Space: rejects path traversal and absolute members, caps
entry count and total uncompressed size (zip-bomb guard), then returns the main
filing PDF (a top-level `*.pdf`, preferring one whose name matches the tracking #).
"""
from __future__ import annotations

import re
import zipfile
from pathlib import Path

MAX_ENTRIES = 200
MAX_TOTAL_BYTES = 200 * 1024 * 1024  # 200 MB uncompressed
TRACKING_RE = re.compile(r"[A-Z]{3,4}-\d{6,}")


def _is_safe(name: str) -> bool:
    """Reject absolute paths and parent-dir traversal."""
    if name.startswith(("/", "\\")) or ".." in Path(name).parts:
        return False
    return True


def safe_unzip(zip_path: Path, dest: Path) -> Path:
    """Extract `zip_path` into `dest` safely; return the chosen main filing PDF."""
    dest.mkdir(parents=True, exist_ok=True)
    with zipfile.ZipFile(zip_path) as zf:
        infos = zf.infolist()
        assert len(infos) <= MAX_ENTRIES, f"zip has too many entries ({len(infos)})"
        total = sum(i.file_size for i in infos)
        assert total <= MAX_TOTAL_BYTES, f"zip too large uncompressed ({total} bytes)"
        for info in infos:  # bounded by MAX_ENTRIES
            if info.is_dir() or not _is_safe(info.filename):
                continue
            target = dest / info.filename
            target.parent.mkdir(parents=True, exist_ok=True)
            with zf.open(info) as src, target.open("wb") as out:
                out.write(src.read())
    return _pick_main_pdf(dest)


def _pick_main_pdf(root: Path) -> Path:
    """Choose the main filing PDF: a top-level *.pdf, preferring the tracking-# one."""
    top_pdfs = sorted(p for p in root.glob("*.pdf"))
    if not top_pdfs:
        any_pdfs = sorted(root.rglob("*.pdf"))
        assert any_pdfs, "no PDF found in the uploaded zip"
        return any_pdfs[0]
    tracked = [p for p in top_pdfs if TRACKING_RE.search(p.stem)]
    return tracked[0] if tracked else top_pdfs[0]


def tracking_from_name(pdf_path: Path) -> str:
    """Best-effort SERFF tracking number from the PDF filename."""
    m = TRACKING_RE.search(pdf_path.stem)
    return m.group(0) if m else pdf_path.stem