foto-duplicate-finder/tools/make-edge-case-fixtures.py

#!/usr/bin/env python3
"""Generate hostile / edge-case image fixtures for Foto Dup Finder.

Every file here is something a real photo library can contain: an interrupted download, a
half-written JPEG, a crafted "decompression bomb", a photo whose pixels are fine but whose
EXIF says it is rotated, two byte-identical copies, and so on.

They are written to fixtures/foto-dup-edge/ inside the repository. That directory is
git-ignored — the fixtures are regenerated on demand, so no binary test data is committed:

    python3 tools/make-edge-case-fixtures.py
    adb push fixtures/foto-dup-edge /sdcard/DCIM/foto-dup-edge

Then scan in the app and check: nothing crashes, unreadable files are reported rather than
silently dropped, the EXIF variants of the same picture are *not* reported as duplicates of
each other (they normalise to the same pixels), and dup-a/dup-b *are*.
"""
import random
import struct
import zlib
from pathlib import Path

from PIL import Image

REPO = Path(__file__).resolve().parent.parent
OUT = REPO / "fixtures" / "foto-dup-edge"
random.seed(7)


def write(name: str, data: bytes) -> Path:
    path = OUT / name
    path.write_bytes(data)
    return path


def gradient(width: int = 320, height: int = 240) -> Image.Image:
    """A structured base image — distinct enough from the solids to matter."""
    img = Image.new("RGB", (width, height))
    px = img.load()
    for y in range(height):
        for x in range(width):
            px[x, y] = ((60 + x) % 256, (90 + y) % 256, 200)
    return img


def png_chunk(kind: bytes, data: bytes) -> bytes:
    return (
        struct.pack(">I", len(data))
        + kind
        + data
        + struct.pack(">I", zlib.crc32(kind + data) & 0xFFFFFFFF)
    )


def heic_section(base: Image.Image) -> None:
    """HEIC samples, encoded locally so they carry no third-party licence.

    Android can decode HEIC from API 28, so these exercise that path — and because we encode
    them ourselves from our own images, they can be shipped/pushed freely.
    """
    try:
        import pillow_heif

        pillow_heif.register_heif_opener()
    except ImportError:
        print("\nHEIC fixtures skipped — needs an encoder. Generate them with:")
        print("  python3 -m venv /tmp/heicvenv")
        print("  /tmp/heicvenv/bin/pip install pillow-heif")
        print("  /tmp/heicvenv/bin/python tools/make-edge-case-fixtures.py")
        return

    # The iPhone default format: control, plus a byte-identical copy for the exact-duplicate path.
    basic = OUT / "heic-basic.heic"
    base.save(basic, "HEIF", quality=90)
    write("heic-dup.heic", basic.read_bytes())

    # Bigger than a thumbnail, so the subsampled fallback decode gets exercised.
    gradient(2400, 1800).save(OUT / "heic-large.heic", "HEIF", quality=85)

    # NOTE: HEIC EXIF *orientation* cannot be set with this encoder. Both `exif=<bytes>` and
    # `exif=<Image.Exif>` are written, but the orientation tag comes back as 1 regardless, which
    # would make a fixture that lies about what it tests. That one case therefore has to be
    # checked with a real camera/iPhone HEIC — see docs/EDGE_CASE_TESTING.md.

    # One scene encoded both ways: is a codec conversion still worth reporting?
    scene = gradient(800, 600)
    scene.save(OUT / "scene.jpg", "JPEG", quality=92)
    scene.save(OUT / "scene.heic", "HEIF", quality=90)

    # Near-identical solids in HEIC too — the band-index stress case.
    for index in (0, 1):
        Image.new("RGB", (400, 300), (20 + index, 40, 60)).save(
            OUT / f"heic-solid-{index}.heic",
            "HEIF",
            quality=90,
        )


def main() -> None:
    OUT.mkdir(parents=True, exist_ok=True)
    base = gradient()

    # A distinct control scene, so the duplicate pair below is unambiguous. (Previously this was
    # byte-identical to dup-a/dup-b, which produced a confusing group of three.)
    control = gradient(240, 180)
    control = control.transpose(Image.FLIP_LEFT_RIGHT)
    write("valid.jpg", _jpeg_bytes(control))

    good = _jpeg_bytes(base)

    # 2. Interrupted download / half-written file: valid header, no end-of-image marker.
    write("truncated.jpg", good[: len(good) * 4 // 10])

    # 3. Right header, garbage body.
    write("garbage-body.jpg", good[:120] + bytes(random.getrandbits(8) for _ in range(4096)))

    # 4. Zero bytes. MediaStore normally filters these, but a sideloaded file can slip through.
    write("zero-bytes.jpg", b"")

    # 5. A PNG wearing a .jpg extension — the MIME type lies.
    png_bytes = _png_bytes(base)
    write("renamed-png.jpg", png_bytes)

    # 6. A valid JPEG with no extension at all.
    write("no-extension", good)

    # 7. Decompression bomb: the header declares 60000 x 60000 (14 GB decoded) but carries
    #    almost no pixel data. A naive decoder allocates until it dies.
    ihdr = struct.pack(">IIBBBBB", 60000, 60000, 8, 2, 0, 0, 0)
    write(
        "decompression-bomb.png",
        b"\x89PNG\r\n\x1a\n"
        + png_chunk(b"IHDR", ihdr)
        + png_chunk(b"IDAT", zlib.compress(b"\x00" * 64))
        + png_chunk(b"IEND", b""),
    )

    # 8. The same picture written with every EXIF orientation value. All eight are the same
    #    photo and must normalise to the same fingerprint.
    for orientation in range(1, 9):
        exif = Image.Exif()
        exif[0x0112] = orientation
        path = OUT / f"exif-orientation-{orientation}.jpg"
        base.save(path, "JPEG", quality=90, exif=exif)

    # 9. Byte-identical copies — the exact-duplicate path must find these.
    write("dup-a.jpg", good)
    write("dup-b.jpg", good)

    # 10. Odd but legal encodings that must not break the hasher.
    base.convert("L").save(OUT / "grayscale.jpg", "JPEG", quality=90)
    base.convert("RGBA").save(OUT / "with-alpha.png", "PNG")
    Image.new("RGB", (1, 1), (255, 0, 0)).save(OUT / "tiny-1px.jpg", "JPEG")

    # 11. Near-identical solids: colour-only similarity bait, and a stress case for the banded
    #     index (all of these may share hash bands).
    for index in range(6):
        Image.new("RGB", (400, 300), (20 + index, 40, 60)).save(
            OUT / f"solid-{index}.png", "PNG"
        )

    total = sorted(OUT.iterdir())
    heic_section(base)
    total = sorted(OUT.iterdir())
    print(f"wrote {len(total)} fixtures to {OUT}")
    for path in total:
        print(f"  {path.stat().st_size:>9,}  {path.name}")


def _jpeg_bytes(img: Image.Image) -> bytes:
    import io

    buffer = io.BytesIO()
    img.save(buffer, "JPEG", quality=92)
    return buffer.getvalue()


def _png_bytes(img: Image.Image) -> bytes:
    import io

    buffer = io.BytesIO()
    img.save(buffer, "PNG")
    return buffer.getvalue()


if __name__ == "__main__":
    main()