mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-08-09 23:50:40 +02:00
237 lines
8.5 KiB
Python
237 lines
8.5 KiB
Python
"""Tests for the metadata-only forensic collector."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import base64
|
|
import json
|
|
import zlib
|
|
from typing import TYPE_CHECKING
|
|
|
|
import piexif
|
|
import pytest
|
|
from PIL import Image
|
|
from PIL.PngImagePlugin import PngInfo
|
|
|
|
from remove_ai_watermarks.forensic_metadata import (
|
|
FORENSIC_METADATA_RECORD_TYPE,
|
|
FORENSIC_METADATA_SCHEMA_VERSION,
|
|
SUPPORTED_EXTENSIONS,
|
|
_b64,
|
|
_decode_exif_value,
|
|
_jpeg_forensics_bytes,
|
|
_png_text_decode,
|
|
_safe_str,
|
|
apple_live_photo_id,
|
|
collect_forensic_metadata,
|
|
read_full_exif,
|
|
read_isobmff_inventory,
|
|
read_isobmff_provenance_path,
|
|
read_jpeg_segments,
|
|
read_pil_info,
|
|
read_png_chunks,
|
|
read_png_late_metadata_path,
|
|
read_webp_chunks,
|
|
sha256_of,
|
|
sniff_format,
|
|
xattr_quarantine,
|
|
xattr_where_from,
|
|
)
|
|
|
|
if TYPE_CHECKING:
|
|
from pathlib import Path
|
|
|
|
|
|
def _jpeg(path: Path) -> Path:
|
|
Image.new("RGB", (48, 32), (20, 80, 160)).save(path, "JPEG", quality=87)
|
|
return path
|
|
|
|
|
|
def _png_chunk(chunk_type: bytes, payload: bytes) -> bytes:
|
|
crc = zlib.crc32(chunk_type + payload).to_bytes(4, "big")
|
|
return len(payload).to_bytes(4, "big") + chunk_type + payload + crc
|
|
|
|
|
|
def test_supported_extensions_are_media_not_documents():
|
|
assert {".jpg", ".png", ".webp", ".heic", ".mp4"}.issubset(SUPPORTED_EXTENSIONS)
|
|
assert ".pdf" not in SUPPORTED_EXTENSIONS
|
|
|
|
|
|
def test_json_helpers_and_format_sniffer():
|
|
class BadString:
|
|
def __str__(self):
|
|
raise RuntimeError("no string")
|
|
|
|
assert _safe_str("ok") == "ok"
|
|
assert "BadString" in _safe_str(BadString())
|
|
assert _b64(b"abc") == base64.b64encode(b"abc").decode("ascii")
|
|
assert _b64(b"x" * 20, cap=4) == "eHh4eA==...TRUNCATED(20 bytes total)"
|
|
assert _decode_exif_value(b"ascii") == "ascii"
|
|
assert _decode_exif_value(b"\xff").startswith("hex:")
|
|
assert _decode_exif_value((1, b"two")) == [1, "two"]
|
|
assert sniff_format(b"\x89PNG\r\n\x1a\n") == "png"
|
|
assert sniff_format(b"\xff\xd8\xff\xe0") == "jpeg"
|
|
assert sniff_format(b"RIFF....WEBP") == "webp"
|
|
assert sniff_format(b"....ftypheic").startswith("isobmff:")
|
|
assert sniff_format(b"unknown").startswith("unknown:")
|
|
|
|
|
|
def test_png_text_and_container_metadata_are_preserved(tmp_path: Path):
|
|
info = PngInfo()
|
|
info.add_text("parameters", "Steps: 20, Model: SDXL", zip=True)
|
|
path = tmp_path / "workflow.png"
|
|
Image.new("RGB", (32, 32)).save(path, pnginfo=info)
|
|
trailer = b'<TC260:AIGC>{"Label":"1"}</TC260:AIGC>'
|
|
path.write_bytes(path.read_bytes() + trailer)
|
|
|
|
record = collect_forensic_metadata(path)
|
|
|
|
assert record["schema_version"] == FORENSIC_METADATA_SCHEMA_VERSION == 1
|
|
assert record["record_type"] == FORENSIC_METADATA_RECORD_TYPE == "forensic_metadata"
|
|
assert record["content_format"] == "png"
|
|
assert any(chunk.get("type") == "zTXt" for chunk in record["png_chunks"])
|
|
assert record["png_post_iend_bytes"] == len(trailer)
|
|
assert base64.b64decode(record["png_post_iend_base64"]) == trailer
|
|
assert "Steps: 20" in json.dumps(record)
|
|
assert json.loads(json.dumps(record, allow_nan=False)) == record
|
|
|
|
|
|
def test_png_text_decoders_and_direct_chunk_reader(tmp_path: Path):
|
|
assert "hello" in _png_text_decode("tEXt", b"key\x00hello")
|
|
compressed = b"prompt\x00\x00" + zlib.compress(b"workflow")
|
|
assert "workflow" in _png_text_decode("zTXt", compressed)
|
|
assert "value" in _png_text_decode("iTXt", b"key\x00\x00\x00\x00\x00value")
|
|
|
|
path = tmp_path / "plain.png"
|
|
Image.new("RGB", (8, 8)).save(path)
|
|
chunks, trailer = read_png_chunks(path.read_bytes())
|
|
|
|
assert chunks[0]["type"] == "IHDR"
|
|
assert trailer == b""
|
|
|
|
|
|
def test_jpeg_exif_segments_encoder_and_trailer(tmp_path: Path):
|
|
path = tmp_path / "camera.jpg"
|
|
exif = piexif.dump(
|
|
{
|
|
"0th": {
|
|
piexif.ImageIFD.Make: b"Camera Corp",
|
|
piexif.ImageIFD.Software: b"Camera Firmware",
|
|
},
|
|
"Exif": {},
|
|
"GPS": {},
|
|
"1st": {},
|
|
}
|
|
)
|
|
Image.new("RGB", (64, 48)).save(path, "JPEG", exif=exif, quality=82)
|
|
trailer = b'PhotoEditor_Re_Edit_Data{"genAIType":1}'
|
|
path.write_bytes(path.read_bytes() + trailer)
|
|
|
|
record = collect_forensic_metadata(path)
|
|
|
|
assert record["exif"]["0th"]["Make"] == "Camera Corp"
|
|
assert record["jpeg"]["post_eoi_bytes"] == len(trailer)
|
|
assert base64.b64decode(record["jpeg"]["post_eoi_base64"]) == trailer
|
|
assert record["jpeg_forensics"]["quant_tables"]
|
|
assert sha256_of(path.read_bytes()) == record["sha256"]
|
|
|
|
|
|
def test_direct_exif_pil_and_jpeg_readers(tmp_path: Path):
|
|
path = _jpeg(tmp_path / "plain.jpg")
|
|
|
|
exif, thumbnail = read_full_exif(path)
|
|
pil, iptc, exif_blob = read_pil_info(path)
|
|
segments = read_jpeg_segments(path.read_bytes())
|
|
|
|
assert isinstance(exif, dict)
|
|
assert thumbnail is None
|
|
assert pil["width"] == 48
|
|
assert pil["height"] == 32
|
|
assert isinstance(iptc, dict)
|
|
assert exif_blob is None or isinstance(exif_blob, bytes)
|
|
assert isinstance(segments["segments"], list)
|
|
assert _jpeg_forensics_bytes(path.read_bytes())["quant_tables"]
|
|
assert _jpeg_forensics_bytes(b"not a jpeg") == {}
|
|
|
|
|
|
def test_webp_inventory_keeps_metadata_but_not_frame_pixels(tmp_path: Path):
|
|
path = tmp_path / "image.webp"
|
|
xmp = b"<x:xmpmeta>metadata</x:xmpmeta>"
|
|
Image.new("RGB", (32, 32), (30, 40, 50)).save(path, "WEBP", xmp=xmp)
|
|
|
|
chunks = read_webp_chunks(path.read_bytes())
|
|
|
|
xmp_chunk = next(chunk for chunk in chunks if chunk["type"] == "XMP ")
|
|
assert xmp_chunk["text"] == xmp.decode()
|
|
assert all("base64" not in chunk for chunk in chunks if chunk["type"] in {"VP8 ", "VP8L", "ANMF"})
|
|
|
|
|
|
def test_isobmff_inventory_and_streaming_provenance(tmp_path: Path):
|
|
path = tmp_path / "signed.mp4"
|
|
ftyp = b"\x00\x00\x00\x18ftypmp42\x00\x00\x00\x00mp42isom"
|
|
payload = b"jumb c2pa trainedAlgorithmicMedia"
|
|
uuid_box = (8 + len(payload)).to_bytes(4, "big") + b"uuid" + payload
|
|
path.write_bytes(ftyp + b"\x00\x00\x00\x08mdat" + uuid_box)
|
|
|
|
inventory = read_isobmff_inventory(path.read_bytes())
|
|
streamed = read_isobmff_provenance_path(path)
|
|
|
|
assert "ftyp" in inventory["boxes"]
|
|
assert base64.b64decode(inventory["provenance_boxes"][0]["base64"]) == payload
|
|
assert base64.b64decode(streamed["provenance_boxes"][0]["base64"]) == payload
|
|
|
|
|
|
def test_oversized_path_keeps_bounded_windows_and_late_png_metadata(tmp_path: Path, monkeypatch):
|
|
path = tmp_path / "late.png"
|
|
Image.new("RGB", (16, 16)).save(path)
|
|
source = path.read_bytes()
|
|
iend = source.rfind(b"\x00\x00\x00\x00IEND")
|
|
padding = _png_chunk(b"vpAg", b"\x00" * ((1 << 20) + 1))
|
|
metadata = b'AIGC\x00{"Label":"1"}'
|
|
path.write_bytes(source[:iend] + padding + _png_chunk(b"tEXt", metadata) + source[iend:])
|
|
monkeypatch.setattr("remove_ai_watermarks.forensic_metadata._MAX_FULL_READ", 1)
|
|
|
|
record = collect_forensic_metadata(path)
|
|
|
|
assert record["oversized"]["head_scanned_bytes"] == path.stat().st_size
|
|
assert base64.b64decode(record["raw_metadata_windows"]["head_base64"])
|
|
assert base64.b64decode(record["png_late_metadata_chunks"][0]["base64"]) == metadata
|
|
|
|
|
|
def test_collection_registers_optional_heif_and_missing_file_raises(tmp_path: Path, monkeypatch):
|
|
registered = False
|
|
|
|
def mark_registered():
|
|
nonlocal registered
|
|
registered = True
|
|
|
|
monkeypatch.setattr("remove_ai_watermarks.image_io._register_heif", mark_registered)
|
|
collect_forensic_metadata(_jpeg(tmp_path / "plain.jpg"))
|
|
|
|
assert registered is True
|
|
with pytest.raises(FileNotFoundError):
|
|
collect_forensic_metadata(tmp_path / "missing.jpg")
|
|
|
|
|
|
@pytest.mark.parametrize("schema_version", [2, True, 1.0])
|
|
def test_collection_rejects_unsupported_output_schema_before_reading(tmp_path: Path, schema_version: object):
|
|
with pytest.raises(ValueError, match="Unsupported forensic metadata schema"):
|
|
collect_forensic_metadata(
|
|
tmp_path / "missing.jpg",
|
|
schema_version=schema_version, # type: ignore[arg-type]
|
|
)
|
|
|
|
|
|
def test_xattrs_and_live_photo_probe_are_safe_on_plain_file(tmp_path: Path):
|
|
path = _jpeg(tmp_path / "plain.jpg")
|
|
|
|
assert xattr_where_from(path) == [] or isinstance(xattr_where_from(path), list)
|
|
assert xattr_quarantine(path) is None or isinstance(xattr_quarantine(path), str)
|
|
assert apple_live_photo_id(path.read_bytes()) is None
|
|
|
|
|
|
def test_late_png_reader_soft_fails_on_non_png(tmp_path: Path):
|
|
path = tmp_path / "plain.bin"
|
|
path.write_bytes(b"not png")
|
|
|
|
assert read_png_late_metadata_path(path) == []
|