watermarks-remover/tests/test_clean_image.py

"""Tests for PNG/JPEG/WebP metadata strip."""

from __future__ import annotations

import struct
import sys
import zlib
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
SCRIPTS = ROOT / "skills" / "remove-ai-marks" / "scripts"
sys.path.insert(0, str(SCRIPTS))

from image_meta import (  # noqa: E402
    clean_image,
    detect_format,
    inspect_image,
    inspect_webp,
    strip_jpeg,
    strip_png,
    strip_webp,
)


def _png_chunk(ctype: bytes, payload: bytes) -> bytes:
    crc = zlib.crc32(ctype)
    crc = zlib.crc32(payload, crc) & 0xFFFFFFFF
    return struct.pack(">I", len(payload)) + ctype + payload + struct.pack(">I", crc)


def _minimal_png_with_text() -> bytes:
    """1x1 IHDR + tEXt with c2pa marker + IDAT + IEND (not a valid image decode, structure OK)."""
    sig = b"\x89PNG\r\n\x1a\n"
    # IHDR: 1x1 RGB 8bit
    ihdr = struct.pack(">IIBBBBB", 1, 1, 8, 2, 0, 0, 0)
    # minimal empty IDAT may fail decoders; enough for our chunk walker
    idat = zlib.compress(b"\x00\x00\x00")
    text = b"Comment\x00c2pa test contentcredentials"
    return (
        sig
        + _png_chunk(b"IHDR", ihdr)
        + _png_chunk(b"tEXt", text)
        + _png_chunk(b"IDAT", idat)
        + _png_chunk(b"IEND", b"")
    )


def _minimal_jpeg_with_app11() -> bytes:
    """Minimal JPEG: SOI, APP0 JFIF, APP11 with c2pa, SOS stub, EOI."""
    app0 = b"JFIF\x00\x01\x01\x00\x00\x01\x00\x01\x00\x00"
    app0_seg = b"\xff\xe0" + struct.pack(">H", len(app0) + 2) + app0
    app11 = b"JUMB" + b"c2pa-manifest-fake"
    app11_seg = b"\xff\xeb" + struct.pack(">H", len(app11) + 2) + app11
    sos_payload = b"\x03\x01\x00\x02\x11\x03\x11\x00\x3f\x00"
    sos = b"\xff\xda" + struct.pack(">H", len(sos_payload) + 2) + sos_payload
    entropy = b"\x00\x00"  # dummy
    return b"\xff\xd8" + app0_seg + app11_seg + sos + entropy + b"\xff\xd9"


def _webp_chunk(fourcc: bytes, payload: bytes) -> bytes:
    padding = b"\x00" if len(payload) & 1 else b""
    return fourcc + struct.pack("<I", len(payload)) + payload + padding


def _minimal_webp(*chunks: tuple[bytes, bytes]) -> bytes:
    body = b"WEBP" + b"".join(
        _webp_chunk(fourcc, payload) for fourcc, payload in chunks
    )
    return b"RIFF" + struct.pack("<I", len(body)) + body


def test_strip_png_removes_text_c2pa(tmp_path: Path):
    data = _minimal_png_with_text()
    cleaned, actions = strip_png(data)
    assert b"c2pa" not in cleaned.lower() or b"tEXt" not in cleaned
    assert any("drop" in a for a in actions)
    # structural: still starts with PNG sig and has IEND
    assert cleaned.startswith(b"\x89PNG")
    assert b"IEND" in cleaned


def test_strip_jpeg_removes_app11():
    data = _minimal_jpeg_with_app11()
    cleaned, actions = strip_jpeg(data)
    assert b"c2pa-manifest-fake" not in cleaned
    assert any("APP11" in a or "drop" in a for a in actions)
    assert cleaned.startswith(b"\xff\xd8")


def test_clean_image_roundtrip(tmp_path: Path):
    src = tmp_path / "t.png"
    src.write_bytes(_minimal_png_with_text())
    dest = tmp_path / "t.cleaned.png"
    result = clean_image(src, dest)
    assert dest.is_file()
    assert result["bytes_out"] > 0


def test_webp_c2pa_and_xmp_are_detected_and_removed(tmp_path: Path):
    data = _minimal_webp(
        (b"VP8X", b"\x04" + b"\x00" * 9),
        (b"VP8 ", b"image-data"),
        (b"XMP ", b"generator=OpenAI"),
        (b"C2PA", b"jumb c2pa manifest"),
    )
    assert detect_format(data) == "webp"
    has_c2pa, has_ai, findings = inspect_webp(data)
    assert has_c2pa and has_ai
    assert "WebP C2PA chunk" in findings

    cleaned, actions = strip_webp(data)
    assert any("C2PA" in action for action in actions)
    assert any("XMP" in action for action in actions)
    assert struct.unpack("<I", cleaned[4:8])[0] == len(cleaned) - 8
    assert cleaned[20] & 0x04 == 0
    assert inspect_webp(cleaned)[0:2] == (False, False)

    src = tmp_path / "marked.webp"
    dest = tmp_path / "marked.cleaned.webp"
    src.write_bytes(data)
    result = clean_image(src, dest)
    assert result["format"] == "webp"
    assert not result["still_has_c2pa"]
    assert not result["still_has_ai_metadata"]


def test_webp_image_payload_text_is_not_a_metadata_false_positive():
    data = _minimal_webp((b"VP8 ", b"ordinary pixels mentioning c2pa"))
    has_c2pa, has_ai, findings = inspect_webp(data)
    assert not has_c2pa
    assert not has_ai
    assert findings == []


def test_truncated_webp_is_reported_and_not_rewritten():
    data = b"RIFF\x10\x00\x00\x00WEBPC2PA\x10\x00\x00\x00short"
    has_c2pa, has_ai, findings = inspect_webp(data)
    assert not has_c2pa
    assert not has_ai
    assert any("truncated" in finding for finding in findings)
    try:
        strip_webp(data)
    except ValueError as error:
        assert "malformed WebP" in str(error)
    else:
        raise AssertionError("truncated WebP should not be rewritten")


def test_webp_size_mismatch_is_not_rewritten():
    data = _minimal_webp((b"VP8 ", b"image-data"))
    malformed = data[:4] + struct.pack("<I", len(data)) + data[8:]
    assert any("size mismatch" in finding for finding in inspect_webp(malformed)[2])

    try:
        strip_webp(malformed)
    except ValueError as error:
        assert "size mismatch" in str(error)
    else:
        raise AssertionError("WebP with a size mismatch should not be rewritten")