remove-ai-watermarks/tests/test_metadata_record.py

"""Tests for the portable metadata record.

The contract is one property: a verdict built from the record equals the verdict
built from the file. Everything here exists to pin that, or to pin a placement the
record could plausibly drop.
"""

from __future__ import annotations

import json
import struct
import zlib
from pathlib import Path

import numpy as np
import pytest
from PIL import Image

from remove_ai_watermarks.identify import (
    ProvenanceReport,
    evidence_from_metadata_record,
    identify,
    identify_from_evidence,
)
from remove_ai_watermarks.metadata_record import HEAD_WINDOW, collect_metadata_record

FIXTURES = Path(__file__).resolve().parent.parent / "data" / "fixtures" / "provenance"
COMPARED = ("is_ai_generated", "platform", "confidence", "ai_source_kind", "ai_from_metadata")


def _verdict_via_record(path: Path) -> ProvenanceReport:
    """The contractor's path: collect, serialize, judge -- without the file."""
    record = json.loads(json.dumps(collect_metadata_record(path)))
    return identify_from_evidence(evidence_from_metadata_record(record, path=path))


def _assert_same_verdict(path: Path) -> None:
    via_record = _verdict_via_record(path)
    via_file = identify(path, check_visible=False, check_invisible=False)

    for field in COMPARED:
        assert getattr(via_record, field) == getattr(via_file, field), field
    assert sorted(s.name for s in via_record.signals) == sorted(s.name for s in via_file.signals)
    assert sorted(via_record.watermarks) == sorted(via_file.watermarks)


def _noise_png(path: Path, size: tuple[int, int] = (700, 700)) -> Path:
    """A PNG whose IDAT is incompressible, so the file exceeds the head window."""
    rng = np.random.default_rng(0)
    Image.fromarray(rng.integers(0, 255, (size[1], size[0], 3), dtype=np.uint8)).save(path)
    return path


def _insert_png_chunk(path: Path, chunk_type: bytes, payload: bytes) -> None:
    """Splice a chunk in just before IEND, i.e. AFTER the whole pixel stream."""
    data = path.read_bytes()
    end = data.rindex(b"IEND") - 4
    chunk = struct.pack(">I", len(payload)) + chunk_type + payload
    chunk += struct.pack(">I", zlib.crc32(chunk_type + payload) & 0xFFFFFFFF)
    path.write_bytes(data[:end] + chunk + data[end:])


class TestRecordReproducesTheFileVerdict:
    """The whole point of the record. Every tracked provenance fixture is compared,
    which covers C2PA, the China TC260 label, the xAI EXIF pair and the IPTC tag."""

    @pytest.mark.skipif(not FIXTURES.is_dir(), reason="provenance fixtures not present")
    @pytest.mark.parametrize("name", sorted(p.name for p in FIXTURES.iterdir()) if FIXTURES.is_dir() else [])
    def test_fixture(self, name: str):
        _assert_same_verdict(FIXTURES / name)

    @pytest.mark.skipif(not FIXTURES.is_dir(), reason="provenance fixtures not present")
    def test_the_fixtures_actually_exercise_several_signals(self):
        """Guards the test above: if the fixture set ever narrows to one signal
        family, equality across it stops meaning much."""
        found = {
            signal.name
            for path in FIXTURES.iterdir()
            for signal in identify(path, check_visible=False, check_invisible=False).signals
        }
        assert len(found) >= 4, found


class TestPlacementsTheRecordCouldDrop:
    def test_a_trailer_after_eoi_survives(self, tmp_path: Path):
        """Samsung Galaxy AI appends its marker past the JPEG EOI, and the value it is
        gated on can sit further back still. A record that stopped at the last
        structural marker would report no signal at all."""
        path = tmp_path / "edited.jpg"
        Image.fromarray(np.zeros((64, 64, 3), dtype=np.uint8)).save(path, "JPEG")
        with path.open("ab") as handle:
            handle.write(b'PhotoEditor_Re_Edit_Data{"genAIType":17}')

        assert identify(path, check_visible=False, check_invisible=False).confidence == "medium"
        _assert_same_verdict(path)

    def test_a_metadata_chunk_past_the_head_window_survives(self, tmp_path: Path):
        """A PNG encoder may put the label chunk after the pixels. The file is larger
        than the record's head window, so only the seek-past-IDAT reader finds it."""
        path = _noise_png(tmp_path / "late.png")
        assert path.stat().st_size > HEAD_WINDOW
        # An XMP packet in the namespaced TC260 form, which is how the label travels
        # when it is not a bare ``AIGC`` keyword chunk.
        label = json.dumps({"Label": "1", "ContentProducer": "001191110102MACQD9K64010000"})
        xmp = (
            b'<x:xmpmeta xmlns:x="adobe:ns:meta/"><rdf:RDF><rdf:Description '
            b'xmlns:TC260="http://www.tc260.org.cn/ns/AIGC/1.0/"><TC260:AIGC>'
            + label.encode()
            + b"</TC260:AIGC></rdf:Description></rdf:RDF></x:xmpmeta>"
        )
        _insert_png_chunk(path, b"iTXt", b"XML:com.adobe.xmp\x00\x00\x00\x00\x00" + xmp)

        assert "aigc" in {s.name for s in identify(path, check_visible=False, check_invisible=False).signals}
        _assert_same_verdict(path)

    def test_an_exif_pair_survives(self, tmp_path: Path):
        """xAI is recognized from an (ImageDescription, Artist) PAIR, and both are read
        by tag NAME. A record carrying only raw bytes loses it."""
        import piexif

        path = tmp_path / "grok.jpg"
        Image.fromarray(np.zeros((64, 64, 3), dtype=np.uint8)).save(path, "JPEG")
        exif = {
            "0th": {
                piexif.ImageIFD.ImageDescription: b"Signature: " + b"A" * 80,
                piexif.ImageIFD.Artist: b"3f2504e0-4f89-11d3-9a0c-0305e82c3301",
            }
        }
        piexif.insert(piexif.dump(exif), str(path))

        assert "xai_signature" in {s.name for s in identify(path, check_visible=False, check_invisible=False).signals}
        _assert_same_verdict(path)


class TestRecordShape:
    def test_the_record_survives_json(self, tmp_path: Path):
        record = collect_metadata_record(_noise_png(tmp_path / "plain.png"))

        assert json.loads(json.dumps(record))["container"] == "png"

    def test_pixels_are_not_carried(self, tmp_path: Path):
        """The reason the record walks regions instead of shipping the head: a record
        that carried the pixel stream would be the size of the image."""
        path = _noise_png(tmp_path / "big.png", size=(1200, 1200))
        record = collect_metadata_record(path)

        idat = path.read_bytes()
        start = idat.index(b"IDAT") + 4
        assert idat[start : start + 512] not in json.dumps(record).encode()
        assert len(json.dumps(record)) < path.stat().st_size // 10

    def test_an_unreadable_container_still_yields_a_record(self, tmp_path: Path):
        path = tmp_path / "junk.bin"
        path.write_bytes(b"\x00\x01\x02not an image at all" * 100)

        record = collect_metadata_record(path)

        assert record["container"] == "unknown"
        assert json.dumps(record)  # serializable, and no exception on the way here
        _assert_same_verdict(path)

    def test_a_missing_file_does_not_raise(self, tmp_path: Path):
        """Collection runs over whatever a caller hands it, including a path that
        vanished between listing and reading."""
        record = collect_metadata_record(tmp_path / "gone.png")

        assert record["container"] == "unknown"
        assert record["metadata_base64"] == ""


def test_a_webp_record_matches(tmp_path: Path):
    """RIFF has its own walk; a chunk kept or dropped wrongly shows up here."""
    path = tmp_path / "image.webp"
    Image.fromarray(np.zeros((64, 64, 3), dtype=np.uint8)).save(path, "WEBP", xmp=b"<x:xmpmeta>plain</x:xmpmeta>")

    assert collect_metadata_record(path)["container"] == "webp"
    _assert_same_verdict(path)


def test_the_record_never_reopens_the_source(tmp_path: Path, monkeypatch):
    """The reason the record exists: the verdict must run where the file is not."""
    path = FIXTURES / "chatgpt-1.png" if (FIXTURES / "chatgpt-1.png").exists() else _noise_png(tmp_path / "x.png")
    record = json.loads(json.dumps(collect_metadata_record(path)))

    def fail_if_called(*args, **kwargs):
        raise AssertionError("the record path must not open the source file")

    monkeypatch.setattr("builtins.open", fail_if_called)
    monkeypatch.setattr(Path, "open", fail_if_called)
    monkeypatch.setattr(Image, "open", fail_if_called)

    report = identify_from_evidence(evidence_from_metadata_record(record, path=path))

    assert isinstance(report, ProvenanceReport)