photoanalyzer/tests/fixtures/media_corpus.py

"""The golden media corpus: every format, orientation, profile, damage, and
metadata case this application claims to survive (US07-03, concept §18).

``CASES`` is the manifest and the authority. Each entry declares a stable logical
id (never a path), how the file is generated, and what the pipeline must do with
it — decode it, or refuse it with one precise error code. Tests parametrize over the
manifest, so a case that is added here without an expectation, or an expectation
that stops holding, fails the suite rather than quietly going untested.

Everything is generated, never committed: fixed pixel seeds, fixed EXIF strings, no
clock, no network, no personal data. Regeneration is byte-stable, which
``test_media_hardening.py`` proves by building the corpus twice and comparing
checksums — a golden corpus that drifts is not golden.
"""

from __future__ import annotations

import struct
import zlib
from dataclasses import dataclass, field
from pathlib import Path
from typing import Callable

import numpy as np
from PIL import Image

# Error codes the pipeline may answer with; ``None`` means "must render".
UNSUPPORTED = "unsupported_image"
TOO_LARGE = "image_too_large"

# One fixed capture time for every metadata case: the corpus must not depend on when
# it was generated.
CAPTURE_TIME = "2019:07:14 10:30:00"


@dataclass(frozen=True)
class MediaCase:
    id: str
    filename: str
    kind: str  # format | orientation | profile | damaged | metadata
    build: Callable[[Path], None]
    expect_error: str | None = None
    # Declared for the metadata cases: exiftool arguments applied after the pixels
    # are written, and the user fields that must survive every later stage.
    exif_args: tuple[str, ...] = ()
    preserved_fields: tuple[str, ...] = ()
    notes: str = ""
    tags: tuple[str, ...] = field(default_factory=tuple)


# ── generators ────────────────────────────────────────────────────────────────


def _pixels(width: int, height: int, seed: int, bands: int = 3) -> np.ndarray:
    return np.random.default_rng(seed).integers(0, 256, (height, width, bands), dtype=np.uint8)


def _save(path: Path, image: Image.Image, **kwargs) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    image.save(path, **kwargs)


def _jpeg(width=320, height=240, seed=1, **kwargs):
    def build(path: Path) -> None:
        _save(path, Image.fromarray(_pixels(width, height, seed)), quality=90, **kwargs)

    return build


def _oriented_jpeg(orientation: int):
    def build(path: Path) -> None:
        image = Image.fromarray(_pixels(400, 200, 5))  # landscape source
        exif = image.getexif()
        exif[274] = orientation  # 0x0112 Orientation
        _save(path, image, exif=exif, quality=90)

    return build


def _rotated_pixels(path: Path) -> None:
    """The same scene rotated in the pixels instead of in a tag."""
    _save(path, Image.fromarray(_pixels(400, 200, 5)).transpose(Image.ROTATE_90), quality=90)


def _png(alpha: bool = False, seed: int = 2):
    def build(path: Path) -> None:
        if alpha:
            _save(path, Image.fromarray(_pixels(120, 90, seed, bands=4), "RGBA"))
        else:
            _save(path, Image.fromarray(_pixels(120, 90, seed)))

    return build


def _webp(path: Path) -> None:
    _save(path, Image.fromarray(_pixels(150, 100, 3)), quality=80)


def _tiff(path: Path) -> None:
    _save(path, Image.fromarray(_pixels(140, 110, 4)))


def _grayscale(path: Path) -> None:
    _save(path, Image.fromarray(_pixels(100, 100, 6)).convert("L"), quality=90)


def _cmyk(path: Path) -> None:
    _save(path, Image.fromarray(_pixels(100, 100, 7)).convert("CMYK"), quality=90)


def _tiny(path: Path) -> None:
    _save(path, Image.fromarray(_pixels(1, 1, 8)))


def _icc_tagged(path: Path) -> None:
    """A profile-bearing image: the colour-managed decode path must run."""
    from PIL import ImageCms

    profile = bytearray(ImageCms.ImageCmsProfile(ImageCms.createProfile("sRGB")).tobytes())
    # An ICC header carries its creation timestamp at bytes 24..35. Left alone, the
    # corpus would be a different corpus every time it is generated.
    profile[24:36] = b"\x00" * 12
    _save(path, Image.fromarray(_pixels(120, 80, 9)), icc_profile=bytes(profile), quality=90)


def _broken_icc(path: Path) -> None:
    """A profile that is not a profile: a picture is still a picture."""
    _save(path, Image.fromarray(_pixels(120, 80, 10)), icc_profile=b"not-a-profile", quality=90)


def _wide_jpeg(path: Path) -> None:
    """Large enough that decoding it at full resolution is visible in memory."""
    _save(path, Image.fromarray(_pixels(4000, 3000, 11)), quality=70)


def _zero_byte(path: Path) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_bytes(b"")


def _truncated_jpeg(path: Path) -> None:
    image = Image.fromarray(_pixels(400, 300, 12))
    _save(path, image, quality=90)
    data = path.read_bytes()
    path.write_bytes(data[: len(data) // 2])  # header intact, pixels missing


def _corrupt_png(path: Path) -> None:
    image = Image.fromarray(_pixels(120, 90, 13))
    _save(path, image)
    data = bytearray(path.read_bytes())
    data[40:80] = b"\x00" * 40  # shred the compressed stream, keep the header
    path.write_bytes(bytes(data))


def _not_an_image(path: Path) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_bytes(b"This is a text file that happens to be named .jpg\n")


def _png_declaring(width: int, height: int):
    """A tiny, structurally valid PNG whose header claims an enormous picture.

    A few hundred bytes on disk, gigapixels on paper: the pipeline must refuse it
    from the declared dimensions, before a single pixel is allocated. Rewriting the
    IHDR of a real PNG (rather than hand-rolling a stub) keeps the file openable, so
    the refusal is proven to come from the size check and not from a parse failure.
    """

    def build(path: Path) -> None:
        import io

        buffer = io.BytesIO()
        Image.fromarray(_pixels(4, 4, 15)).save(buffer, format="PNG")
        data = bytearray(buffer.getvalue())
        start = 8 + 4  # PNG signature, then the IHDR length field
        struct.pack_into(">II", data, start + 4, width, height)
        ihdr = bytes(data[start : start + 4 + 13])
        struct.pack_into(">I", data, start + 4 + 13, zlib.crc32(ihdr))
        path.parent.mkdir(parents=True, exist_ok=True)
        path.write_bytes(bytes(data))

    return build


CASES: tuple[MediaCase, ...] = (
    # ── formats ───────────────────────────────────────────────────────────────
    MediaCase("jpeg", "formats/plain.jpg", "format", _jpeg()),
    MediaCase("jpeg_uppercase_ext", "formats/UPPER.JPG", "format", _jpeg(seed=14)),
    MediaCase("png", "formats/plain.png", "format", _png()),
    MediaCase("png_alpha", "formats/alpha.png", "format", _png(alpha=True)),
    MediaCase("webp", "formats/plain.webp", "format", _webp),
    MediaCase("tiff", "formats/plain.tiff", "format", _tiff),
    # ── orientation ───────────────────────────────────────────────────────────
    *(
        MediaCase(
            f"orientation_{value}",
            f"orientation/exif_{value}.jpg",
            "orientation",
            _oriented_jpeg(value),
            notes="EXIF orientation must be applied before resizing",
        )
        for value in range(1, 9)
    ),
    MediaCase("rotated_pixels", "orientation/rotated.jpg", "orientation", _rotated_pixels),
    # ── colour and size profiles ──────────────────────────────────────────────
    MediaCase("grayscale", "profiles/gray.jpg", "profile", _grayscale),
    MediaCase("cmyk", "profiles/cmyk.jpg", "profile", _cmyk),
    MediaCase("tiny", "profiles/tiny.png", "profile", _tiny),
    MediaCase("icc_tagged", "profiles/icc.jpg", "profile", _icc_tagged),
    MediaCase(
        "broken_icc",
        "profiles/broken-icc.jpg",
        "profile",
        _broken_icc,
        notes="an unusable ICC profile falls back to a plain conversion, never an error",
    ),
    MediaCase(
        "large_jpeg",
        "profiles/large.jpg",
        "profile",
        _wide_jpeg,
        notes="12 megapixels: the decode must stay near the requested size",
    ),
    # ── damaged and hostile inputs ────────────────────────────────────────────
    MediaCase("zero_byte", "damaged/empty.jpg", "damaged", _zero_byte, UNSUPPORTED),
    MediaCase("truncated_jpeg", "damaged/truncated.jpg", "damaged", _truncated_jpeg, UNSUPPORTED),
    MediaCase("corrupt_png", "damaged/corrupt.png", "damaged", _corrupt_png, UNSUPPORTED),
    MediaCase("text_as_jpeg", "damaged/text.jpg", "damaged", _not_an_image, UNSUPPORTED),
    MediaCase(
        "bomb_header",
        "damaged/bomb.png",
        "damaged",
        _png_declaring(60_000, 60_000),
        TOO_LARGE,
        notes="3.6 gigapixels declared in the header and nothing else",
    ),
    MediaCase(
        "bomb_warning_band",
        "damaged/bomb-warning.png",
        "damaged",
        _png_declaring(10_000, 10_000),
        TOO_LARGE,
        notes="inside Pillow's warn-only band; the warning is promoted to a refusal",
    ),
    # ── metadata ──────────────────────────────────────────────────────────────
    MediaCase("no_exif", "metadata/bare.jpg", "metadata", _jpeg(seed=20)),
    MediaCase(
        "user_exif",
        "metadata/user.jpg",
        "metadata",
        _jpeg(seed=21),
        exif_args=(
            "-Artist=Ada Lovelace",
            "-Copyright=(c) Ada",
            f"-DateTimeOriginal={CAPTURE_TIME}",
            "-GPSLatitude=48.137",
            "-GPSLatitudeRef=N",
            "-Rating=4",
            "-ImageDescription=A day out",
        ),
        preserved_fields=(
            "EXIF:IFD0:Artist",
            "EXIF:IFD0:Copyright",
            "EXIF:ExifIFD:DateTimeOriginal",
            "EXIF:IFD0:ImageDescription",
            "XMP:XMP-xmp:Rating",
        ),
        notes="user metadata that every stage must leave exactly as it found it",
    ),
    MediaCase(
        "prior_safety_keyword",
        "metadata/prior-safety.jpg",
        "metadata",
        _jpeg(seed=22),
        exif_args=("-Keywords+=nsfw", "-Subject+=nsfw", "-Artist=Ada Lovelace"),
        preserved_fields=("EXIF:IFD0:Artist",),
        notes="a safety decision already written by an earlier run",
    ),
    MediaCase(
        "prior_analysis_keywords",
        "metadata/prior-analysis.jpg",
        "metadata",
        _jpeg(seed=23),
        exif_args=("-Keywords+=beach", "-Keywords+=sunset", "-Subject+=beach", "-Subject+=sunset"),
        notes="analysis keywords from an earlier run; a safety write must not drop them",
    ),
    MediaCase(
        "conflicting_safety_keywords",
        "metadata/conflicting.jpg",
        "metadata",
        _jpeg(seed=24),
        exif_args=("-Keywords+=sfw", "-Keywords+=nsfw", "-Subject+=sfw", "-Subject+=nsfw"),
        notes="both safety keywords at once: mutually exclusive means one must go",
    ),
    MediaCase(
        "malformed_metadata",
        "metadata/malformed.jpg",
        "metadata",
        _jpeg(seed=25, exif=b"\x00\x01\x02not-a-valid-exif-block"),
        notes="a broken EXIF block must not stop the picture from being usable",
    ),
)

CASES_BY_ID = {case.id: case for case in CASES}


def build_corpus(root: Path, *, ids: tuple[str, ...] | None = None) -> dict[str, Path]:
    """Generate the corpus (or a named subset) under ``root``; return id → path."""
    import subprocess

    built: dict[str, Path] = {}
    for case in CASES:
        if ids is not None and case.id not in ids:
            continue
        path = root / case.filename
        case.build(path)
        if case.exif_args:
            subprocess.run(
                ["exiftool", "-m", "-overwrite_original", *case.exif_args, str(path)],
                capture_output=True,
                check=False,
            )
        built[case.id] = path
    return built