diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index cf3d1f8..a10d52f 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -45,6 +45,27 @@ jobs: - name: Run tests run: uv run pytest -q + transport-contract-python-314: + name: versioned transport contracts py3.14 + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v7 + - uses: astral-sh/setup-uv@v7 + with: + python-version: "3.14" + - name: Create clean runtime contract environment + run: uv venv --python 3.14 .venv-contract + - name: Install runtime contract dependencies + run: >- + uv pip install --python .venv-contract/bin/python + ".[pixels,heif]" "numpy>=2,<3" "pytest>=8" + - name: Run versioned transport contract tests + run: >- + .venv-contract/bin/python -m pytest -q + tests/test_forensic_metadata.py + tests/test_metadata_record.py + tests/test_pixel_evidence.py + video-e2e: name: video full-clip end-to-end runs-on: ubuntu-latest diff --git a/.gitignore b/.gitignore index 567e892..598c94d 100644 --- a/.gitignore +++ b/.gitignore @@ -20,8 +20,7 @@ Thumbs.db *.swp *.swo -# SynthID corpus reference fills (synthetic black/white calibration tiles, -# regenerable; the labeled pos/neg/cleaned images ARE tracked, see README) +# Generated calibration inputs data/synthid_corpus/refs/ # Reference materials @@ -35,11 +34,7 @@ yolov8n.pt .claude/settings.local.json .claude/scheduled_tasks.lock -# Visible-watermark alpha calibration. The solid black/gray/white CAPTURES are -# committed (content-free: a solid colour + the watermark; the source for -# scripts/visible_alpha_solve.py so the alpha assets are reproducible). The -# synthetic seeds (regenerable) and any real-content validation download (a real -# generated scene, kept local for privacy) are NOT committed. +# Generated calibration inputs and local evaluation artifacts data/doubao_capture/seeds/ data/jimeng_capture/seeds/ data/jimeng_capture/captures/jimeng_content_*.png @@ -48,10 +43,8 @@ data/gemini_capture/captures/gemini_content_*.png data/samsung_capture/seeds/ data/samsung_capture/captures/samsung_content_* -# Leftover GFPGAN weights dir from the retired face-restore experiments -# (GFPGAN wrote RetinaFace/parsing weights to a CWD ./gfpgan/weights/ working -# dir on first use). Runtime artifact, never committed. +# Runtime model artifacts gfpgan/ -# Local-only working data for analysis (not a committed corpus; never tracked) +# Local evaluation data .local-eval/ diff --git a/docs/module-internals.md b/docs/module-internals.md index 108c869..c827690 100644 --- a/docs/module-internals.md +++ b/docs/module-internals.md @@ -388,9 +388,10 @@ metadata extraction from verdict logic: - `extract_provenance_evidence` reads the supported metadata signals into `ProvenanceEvidence`. - `evidence_from_metadata_record` normalizes an externally collected nested - metadata record into the same evidence type without file access. Diagnostic - values under `error` and `kind` are excluded from evidence while nested raw - bytes remain available through encoded binary fields. + metadata record into the same evidence type without file access. Versioned native + records accept only source-derived fields; filenames, hashes, timings, errors, + prior verdicts, and pixel results cannot become evidence. Unknown native schema + versions and other record types are rejected. - The vendor registries are matched over `_metadata_region(head)`, not the whole scan buffer: they see the container's metadata and not its coded pixels. The tokens are raw substrings and the shortest are four and five bytes, so over a megabyte of @@ -427,19 +428,44 @@ defects found while establishing that equality are the reason each rule exists: returns the FIRST candidate carrying a known token, so preserving candidate order is part of verdict equivalence. -Pixel forensics are deliberately absent: nothing in the provenance path reads them. +The transport is independently versioned as `provenance_metadata` schema 1. Native +records require the exact integer schema version and a `complete` status. Source +read failures are explicit error records and cannot be judged. WebP walks the full +declared RIFF container by seeking over `VP8`, `VP8L`, `ALPH`, and `ANMF`, so late +XMP/C2PA remains visible without shipping coded frames or parsing appended trailer +bytes as chunks. + +Pixel forensics are deliberately absent: the provenance path does not read them. Verdict equivalence is checked over tracked fixtures and a separate local evaluation corpus. -### Experimental pixel forensics +### Broad forensic metadata + +[`forensic_metadata.py`](../src/remove_ai_watermarks/forensic_metadata.py) owns the +wide metadata-only inspection record: hashes and timestamps, full EXIF/IPTC, C2PA, +container inventories, bounded binary metadata, and embedded-thumbnail forensics. +It is a separate `forensic_metadata` record type and is deliberately rejected by the +provenance normalizer. Integration code publishes the strict +`ProvenanceReport.to_dict()` alongside it rather than letting operational fields or +derived results influence detection. + +### Pixel forensics [`pixel_evidence.py`](../src/remove_ai_watermarks/pixel_evidence.py) measures six families of scale-robust pixel statistics (block-DCT histograms and Benford deviation, FFT band energies and CFA peaks, high-pass residual, error level, -gradient, colour) in a single decode, sharing the intermediate maps between them. +gradient, color) in a single decode, sharing the intermediate maps between them. -It has no consumer. Nothing in the package reads it -- not the verdict, not removal, -not the CLI -- and the shape is unstable until something does. +It remains independent of verdict and removal. `PixelEvidence.to_dict()` is the +versioned service boundary: it omits the local path, exposes complete/partial/error +status, keeps exception details in logs, and can include opt-in per-stage timings. + +The provenance metadata collector, broad forensic collector, provenance report, +and pixel report all accept an explicit output `schema_version`. Package releases +may add an output schema while retaining older serializers, so a rolling consumer +can keep requesting the version it already understands. Within one schema, changes +are additive; existing fields, types, meanings, signal names, and watermark labels +remain stable. Unsupported selections raise before a different shape is returned. `artifacts=True` additionally returns the spatial layer: a perceptual hash, a 128px JPEG thumbnail, and coarse ELA, residual and phase maps. Those identify the source diff --git a/docs/python-api.md b/docs/python-api.md index 1944a28..8839c7d 100644 --- a/docs/python-api.md +++ b/docs/python-api.md @@ -194,13 +194,20 @@ import json from remove_ai_watermarks.identify import identify_metadata_record from remove_ai_watermarks.metadata_record import collect_metadata_record -record = collect_metadata_record(Path("input.png")) # reads the file +record = collect_metadata_record(Path("input.png"), schema_version=1) # reads the file blob = json.dumps(record) # ship it anywhere report = identify_metadata_record(json.loads(blob), path=Path("input.png")) # reads nothing -payload = report.to_dict() # versioned JSON contract +payload = report.to_dict(schema_version=1) # versioned JSON contract ``` +The collection record has `record_type="provenance_metadata"`, +`schema_version=1`, and a `status`. A vanished or unreadable source produces an +`error` record with structured `issues`; `identify_metadata_record` rejects that +record instead of turning a collection failure into an unknown-image verdict. +Unknown schema versions, non-integer aliases, and native records without a +`complete` collection status are rejected explicitly. + The verdict is the same one `identify(path, check_visible=False, check_invisible=False)` returns for that file. That equality is the record's whole contract and is verified over the tracked provenance fixtures and a separate local @@ -208,9 +215,18 @@ evaluation corpus. `ProvenanceReport.to_dict()` is the stable service boundary: adds a `schema_version`, contains only JSON-safe values, and deliberately omits the local source path. -A record carries metadata regions, never pixels: marker segments before the JPEG -scan, every PNG chunk except `IDAT`, RIFF chunks except the coded image, the -ISOBMFF provenance boxes, the container's trailer, the parsed EXIF tags the +Package and transport versions evolve independently. Long-lived consumers should +request the schema they implement, as above, instead of assuming the installed +package's latest schema. Within schema 1, existing fields, types, meanings, +`signals[].name` values, and `watermarks[]` labels remain compatible; releases may +add fields that consumers must ignore. A breaking change requires a new schema while +the schema 1 serializer remains available for rolling upgrades. Asking a release for +an unsupported schema raises `ValueError` rather than silently returning another +shape. + +A record carries metadata regions, not the primary coded-pixel stream: marker +segments before the JPEG scan, every PNG chunk except `IDAT`, RIFF chunks except the +coded image, the ISOBMFF provenance boxes, the container's trailer, the parsed EXIF tags the verdict reads by name, PIL's info mapping, and the C2PA manifest store. Record size is bounded by those metadata regions and trailers; images with large embedded manifests naturally produce larger records. @@ -236,13 +252,54 @@ evidence = evidence_from_metadata_record(record, path=Path("input.png")) report = identify_from_evidence(evidence) ``` -The normalizer recursively preserves text and byte values. It also decodes +Unversioned third-party records are normalized recursively for compatibility. The +normalizer preserves text and byte values. It also decodes strings prefixed with `hex:` and fields named `base64` or ending in -`_base64`. Diagnostic values under `error` and `kind` are ignored because they -describe the collector rather than the source file. Pass a C2PA manifest-store -dictionary in `record["c2pa_store"]`, or through the explicit +`_base64`. Diagnostic, transport, timing, hash, provenance-result, and pixel-result +subtrees are ignored because they describe the collector or a derived result rather +than the source file. A versioned portable record is stricter still: only +`metadata_base64`, `tail_base64`, `pil`, `exif`, and `c2pa_store` are accepted as +source evidence. Other `record_type` values are rejected, so do not pass a broad +forensic inspection record to this API. Pass a C2PA manifest-store dictionary in +`record["c2pa_store"]`, or through the explicit `c2pa_manifest_store` argument. +### Broad metadata inspection + +`collect_forensic_metadata` provides the wide metadata-only record used by forensic +inspection and migration adapters. It preserves hashes and timestamps, full EXIF and +IPTC, C2PA, container inventories, bounded raw metadata payloads, and embedded +thumbnail forensics. It does not calculate a provenance verdict or pixel statistics. + +```python +from remove_ai_watermarks.forensic_metadata import collect_forensic_metadata + +record = collect_forensic_metadata(Path("input.png"), schema_version=1) +assert record["record_type"] == "forensic_metadata" +``` + +This record is intentionally not accepted by `identify_metadata_record`. Collect the +small strict provenance record separately and publish the resulting +`ProvenanceReport.to_dict()` as the detector contract. + +### Pixel evidence + +`extract_pixel_evidence` decodes once and calculates the DCT, FFT, residual, ELA, +gradient, and color families. Its versioned `to_dict()` result has a semantic +`status`: `complete`, `partial` when an individual family failed, or `error` when the +source could not be decoded. Transported errors contain only the exception class, so +local paths stay in the caller's logs rather than crossing the service boundary. + +```python +from remove_ai_watermarks.pixel_evidence import extract_pixel_evidence + +pixels = extract_pixel_evidence(Path("input.png"), artifacts=False, timings=True) +payload = pixels.to_dict(schema_version=1) +``` + +Timings and spatial artifacts are opt-in. Artifacts include image-identifying data +such as a thumbnail and perceptual hash; aggregate feature families do not. + `identify_from_evidence` does not reopen the source file by default: it evaluates metadata only, and registered visible marks and pixel-backed invisible watermarks remain in the path-based `identify` call. diff --git a/docs/release-and-distribution.md b/docs/release-and-distribution.md index e034c9d..9e1b17d 100644 --- a/docs/release-and-distribution.md +++ b/docs/release-and-distribution.md @@ -71,10 +71,10 @@ The source distribution uses an explicit allowlist for `/src`, `/LICENSE`, `[tool.hatch.build.targets.sdist]` in `pyproject.toml`. It also defensively excludes `/data`, `/tmp`, and `/.sc`. Keep both controls: calibration captures, test corpora, generated research outputs, and local session state do not belong -in the published package archive. `.gitignore` covers `tmp/` and `.sc/`, so those -never reach a commit, but ignore rules are not the build boundary -- hatchling may -include untracked files, and `data/` is deliberately tracked, so the sdist exclude -is the only control keeping it out of the archive. +in the published package archive. Hatchling always adds the root `.gitignore` to +the sdist, so keep its comments generic and free of local operational context. +Ignore rules are not the build boundary: `data/` is deliberately tracked, while +the sdist configuration keeps it and the other excluded paths out of the archive. ## Build backend @@ -97,6 +97,15 @@ write access. ## Release verification +Forensic transports are versioned independently from the package. Before publishing +a change to provenance metadata, provenance reports, broad forensic metadata, or +pixel evidence, run their schema 1 contract tests. Additive fields are compatible; +renaming a field, changing its type or meaning, changing a signal name or watermark +label, or removing a field requires a new output schema. Add the new serializer +without removing schema 1 so long-lived consumers can update separately. A package +release must never silently substitute its latest schema when a caller explicitly +requests an older supported one. + After publication, verify: - both wheel and source distribution exist on PyPI; diff --git a/pyproject.toml b/pyproject.toml index 4afc71f..c93695e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,6 +37,7 @@ classifiers = [ "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.13", + "Programming Language :: Python :: 3.14", "Topic :: Multimedia :: Graphics", "Topic :: Multimedia :: Graphics :: Graphics Conversion", "Topic :: Scientific/Engineering :: Artificial Intelligence", diff --git a/src/remove_ai_watermarks/_internal/constants.py b/src/remove_ai_watermarks/_internal/constants.py index 2ff85ce..46b370b 100644 --- a/src/remove_ai_watermarks/_internal/constants.py +++ b/src/remove_ai_watermarks/_internal/constants.py @@ -20,6 +20,9 @@ AI_KEYWORDS = _tokens( PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n" C2PA_CHUNK_TYPE = b"caBX" +PNG_METADATA_CHUNKS = frozenset({b"tEXt", b"iTXt", b"zTXt", b"eXIf", b"iCCP"}) +RIFF_METADATA_CHUNKS = frozenset({b"EXIF", b"XMP ", b"ICCP", b"C2PA"}) +RIFF_CODED_IMAGE_CHUNKS = frozenset({b"VP8 ", b"VP8L", b"ALPH", b"ANMF"}) C2PA_SIGNATURES = tuple( token.encode() for token in _tokens("c2pa|C2PA|jumb|jumd|JUMBF|jumbf|cbor|contentcreds|digid|assertions|manifest") ) diff --git a/src/remove_ai_watermarks/_internal/isobmff.py b/src/remove_ai_watermarks/_internal/isobmff.py index 03f5865..e192c55 100644 --- a/src/remove_ai_watermarks/_internal/isobmff.py +++ b/src/remove_ai_watermarks/_internal/isobmff.py @@ -64,7 +64,7 @@ _AI_LABEL_MARKERS: tuple[bytes, ...] = AIGC_MARKERS + IPTC_AI_MARKERS + IPTC_AI_ # blanked in place (see ``blank_ai_xmp_packets``). _XMP_PACKET_RE = re.compile(rb"<\?xpacket begin=.*?<\?xpacket end=[^>]*?\?>", re.DOTALL) _STREAM_COPY_BYTES = 1024 * 1024 -_STREAM_SCAN_BYTES = 4 * 1024 * 1024 +STREAM_SCAN_BYTES = 4 * 1024 * 1024 # TC260-PG-20257A stores an MP4/MOV label as an ``AIGC`` key in @@ -133,7 +133,7 @@ def _read_box_header( return end, box_type, payload_off -def _iter_file_boxes( +def iter_file_boxes( stream: BinaryIO, start: int, end: int, @@ -193,17 +193,17 @@ def _tc260_aigc_regions( Each tuple is ``(key_start, key_end, value_start, value_end, value)``. """ regions: list[tuple[int, int, int, int, bytes]] = [] - for _moov_start, moov_end, moov_type, moov_payload in _iter_file_boxes(stream, 0, file_size): + for _moov_start, moov_end, moov_type, moov_payload in iter_file_boxes(stream, 0, file_size): if moov_type != b"moov": continue - for _udta_start, udta_end, udta_type, udta_payload in _iter_file_boxes( + for _udta_start, udta_end, udta_type, udta_payload in iter_file_boxes( stream, moov_payload, moov_end, ): if udta_type != b"udta": continue - for _meta_start, meta_end, meta_type, meta_payload in _iter_file_boxes( + for _meta_start, meta_end, meta_type, meta_payload in iter_file_boxes( stream, udta_payload, udta_end, @@ -212,7 +212,7 @@ def _tc260_aigc_regions( continue keys: dict[int, tuple[int, int]] = {} ilst_boxes: list[tuple[int, int]] = [] - for _child_start, child_end, child_type, child_payload in _iter_file_boxes( + for _child_start, child_end, child_type, child_payload in iter_file_boxes( stream, meta_payload + 4, meta_end, @@ -224,7 +224,7 @@ def _tc260_aigc_regions( if not keys: continue for ilst_payload, ilst_end in ilst_boxes: - for _item_start, item_end, item_type, item_payload in _iter_file_boxes( + for _item_start, item_end, item_type, item_payload in iter_file_boxes( stream, ilst_payload, ilst_end, @@ -233,7 +233,7 @@ def _tc260_aigc_regions( key_span = keys.get(index) if key_span is None: continue - for _data_start, data_end, data_type, data_payload in _iter_file_boxes( + for _data_start, data_end, data_type, data_payload in iter_file_boxes( stream, item_payload, item_end, @@ -425,7 +425,7 @@ def strip_isobmff_media_file( source: str | Path, output: str | Path, *, - max_box_scan: int = _STREAM_SCAN_BYTES, + max_box_scan: int = STREAM_SCAN_BYTES, ) -> tuple[int, int]: """Stream-copy an MP4/MOV/M4A while removing supported AI metadata. diff --git a/src/remove_ai_watermarks/_internal/schema.py b/src/remove_ai_watermarks/_internal/schema.py new file mode 100644 index 0000000..8253e9d --- /dev/null +++ b/src/remove_ai_watermarks/_internal/schema.py @@ -0,0 +1,16 @@ +"""Shared validation for versioned JSON transport contracts.""" + +from collections.abc import Collection + + +def require_schema_version( + value: object, + *, + contract: str, + supported: Collection[int], +) -> int: + """Return an explicitly supported integer schema version or raise.""" + if type(value) is not int or value not in supported: + versions = ", ".join(str(version) for version in sorted(supported)) + raise ValueError(f"Unsupported {contract} schema: {value!r}; supported versions: {versions}") + return value diff --git a/src/remove_ai_watermarks/forensic_metadata.py b/src/remove_ai_watermarks/forensic_metadata.py new file mode 100644 index 0000000..c49a930 --- /dev/null +++ b/src/remove_ai_watermarks/forensic_metadata.py @@ -0,0 +1,861 @@ +"""Collect JSON-safe metadata and container forensics for one media file. + +The collector is deliberately evidence-only: it preserves raw EXIF, IPTC, C2PA, +container metadata, encoder structure, hashes, timestamps, and bounded binary +payloads without deciding whether the content is AI-generated. Provenance verdicts +and pixel statistics are separate library stages. +""" + +import base64 +import contextlib +import hashlib +import io +import json +import os +import plistlib +import re +import struct +import zlib +from pathlib import Path +from typing import Any, cast + +import piexif +from PIL import Image +from PIL.IptcImagePlugin import getiptcinfo + +from remove_ai_watermarks import image_io +from remove_ai_watermarks._internal.constants import ( + PNG_METADATA_CHUNKS, + RIFF_CODED_IMAGE_CHUNKS, + RIFF_METADATA_CHUNKS, +) +from remove_ai_watermarks._internal.isobmff import ( + C2PA_BOX_TYPES, + STREAM_SCAN_BYTES, + iter_file_boxes, +) +from remove_ai_watermarks._internal.schema import require_schema_version +from remove_ai_watermarks.metadata import QUICK_SCAN_BYTES +from remove_ai_watermarks.metadata_record import HEAD_WINDOW + +__all__ = [ + "FORENSIC_METADATA_RECORD_TYPE", + "FORENSIC_METADATA_SCHEMA_VERSION", + "SUPPORTED_EXTENSIONS", + "collect_forensic_metadata", +] + +SUPPORTED_EXTENSIONS = { + ".png", + ".jpg", + ".jpeg", + ".webp", + ".heic", + ".heif", + ".avif", + ".tif", + ".tiff", + ".bmp", + ".gif", + # video/px containers: no pixel decode, but C2PA reads them (Sora/Veo + # carry C2PA manifests) and the byte scans still apply + ".mp4", + ".mov", + ".m4v", + ".jxl", +} + +FORENSIC_METADATA_SCHEMA_VERSION = 1 +FORENSIC_METADATA_RECORD_TYPE = "forensic_metadata" + +_B64_CAP = 1 << 20 # 1 MB safety ceiling per embedded blob +_TEXT_CAP = 1 << 20 # decoded PNG text ceiling per chunk +# Preserve enough top-level ISOBMFF uuid/jumb payload data for downstream +# provenance algorithms without requiring them to reopen the source file. +_PROVENANCE_B64_CAP = STREAM_SCAN_BYTES +_RAW_SCAN_HEAD = HEAD_WINDOW +_RAW_SCAN_TAIL = QUICK_SCAN_BYTES + + +def _safe_str(v: Any) -> str: + try: + return str(v) + except Exception: + return repr(v) + + +def _b64(b: bytes, *, cap: int = _B64_CAP) -> str: + """Legacy base64 value, with an explicit marker when the payload is capped.""" + encoded = base64.b64encode(b[:cap]).decode("ascii") + return encoded + f"...TRUNCATED({len(b)} bytes total)" if len(b) > cap else encoded + + +def _decode_exif_value(v: Any) -> Any: + """Make a piexif value JSON-safe; bytes are kept in full as hex.""" + if isinstance(v, bytes): + if len(v) <= 64: + try: + return v.decode("utf-8", "strict") + except (UnicodeDecodeError, ValueError): + return f"hex:{v.hex()}" + return f"hex:{v.hex()}" + if isinstance(v, tuple | list): + sequence = cast("list[Any] | tuple[Any, ...]", v) + return [_decode_exif_value(item) for item in sequence] + return v + + +def read_full_exif( + path: Path, exif_blob: bytes | None = None, data: bytes | None = None +) -> tuple[dict[str, Any], bytes | None]: + """All EXIF IFDs with decoded tag names (piexif, no re-encode), plus the + raw embedded-thumbnail bytes for the caller's own thumbnail forensics. + + ``exif_blob`` is the PIL-exposed EXIF blob (PNG/WebP/HEIC path) so the + caller's single Image.open is not repeated here. ``data`` is the + already-read file bytes so piexif does not re-read the file.""" + try: + exif: dict[str, Any] = piexif.load(data) if data is not None else piexif.load(str(path)) + except Exception: + if not exif_blob: + return {}, None + try: + exif = piexif.load(exif_blob) + except Exception as exc: + return {"error": _safe_str(exc)}, None + out: dict[str, Any] = {} + thumbnail: bytes | None = None + for ifd, tags in exif.items(): + if ifd == "thumbnail": + thumbnail = tags if isinstance(tags, bytes) else None + out["thumbnail"] = f"{len(tags)} bytes" if isinstance(tags, bytes) else None + continue + if not isinstance(tags, dict): + continue + all_tag_names = cast("dict[str, dict[int, dict[str, Any]]]", getattr(piexif, "TAGS", {})) + tag_names = all_tag_names.get(ifd, {}) + decoded: dict[str, Any] = {} + for tag, value in cast("dict[int, Any]", tags).items(): + name = str(tag_names.get(tag, {}).get("name", f"tag_{tag}")) + if name == "MakerNote" and isinstance(value, bytes): + # full hex, no cap: measured on real uploads, Apple is ~2 KB + # but Canon reaches 28 KB and Sony 38 KB (AF data, serials, + # embedded previews) -- a cap would silently drop exactly the + # camera-original evidence this scan exists to preserve + decoded[name] = f"hex:{value.hex()}" + else: + decoded[name] = _decode_exif_value(value) + out[ifd] = decoded + return out, thumbnail + + +def _png_text_decode(ctype: str, body: bytes) -> str: + """Decode a tEXt/zTXt/iTXt chunk, inflating zlib where used. + + The compressed forms are where ComfyUI / Automatic1111 hide the + generation workflow and prompt, so skipping the inflate would drop + the strongest AI-provenance text a PNG can carry.""" + if ctype == "tEXt": + suffix = b"...TRUNCATED" if len(body) > _TEXT_CAP else b"" + return (body[:_TEXT_CAP] + suffix).decode("utf-8", "replace") + if ctype == "zTXt": + nul = body.find(b"\x00") + if nul == -1: + return body[:_TEXT_CAP].decode("utf-8", "replace") + keyword = body[:nul].decode("latin-1", "replace") + # body[nul+1] = compression method (0 = zlib) + try: + inflater = zlib.decompressobj() + decoded = inflater.decompress(body[nul + 2 :], _TEXT_CAP + 1) + suffix = "...TRUNCATED" if len(decoded) > _TEXT_CAP else "" + text = decoded[:_TEXT_CAP].decode("utf-8", "replace") + suffix + except zlib.error: + text = body[:_TEXT_CAP].decode("utf-8", "replace") + return f"{keyword}\x00{text}" + # iTXt: keyword\0 compflag(1) compmethod(1) lang\0 translated\0 text + parts = body.split(b"\x00", 1) + if len(parts) < 2: + return body[:_TEXT_CAP].decode("utf-8", "replace") + keyword = parts[0].decode("latin-1", "replace") + rest = parts[1] + if len(rest) < 2: + return body[:_TEXT_CAP].decode("utf-8", "replace") + compflag = rest[0] + tail = rest[2:] + for _ in range(2): # skip language tag and translated keyword + nul = tail.find(b"\x00") + if nul == -1: + return body[:_TEXT_CAP].decode("utf-8", "replace") + tail = tail[nul + 1 :] + if compflag: + with contextlib.suppress(zlib.error): + inflater = zlib.decompressobj() + tail = inflater.decompress(tail, _TEXT_CAP + 1) + if len(tail) > _TEXT_CAP: + tail = tail[:_TEXT_CAP] + b"...TRUNCATED" + return f"{keyword}\x00{tail.decode('utf-8', 'replace')}" + + +def read_png_chunks(data: bytes) -> tuple[list[dict[str, Any]], bytes]: + """Every PNG chunk in order (type, length; text chunks decoded and + inflated, binary chunks as base64) plus the post-IEND trailer bytes.""" + chunks: list[dict[str, Any]] = [] + post_iend = b"" + try: + pos = 8 + while pos + 12 <= len(data): + length = struct.unpack(">I", data[pos : pos + 4])[0] + ctype = data[pos + 4 : pos + 8].decode("latin-1") + body = data[pos + 8 : pos + 8 + length] + entry: dict[str, Any] = {"type": ctype, "length": length} + if ctype in ("tEXt", "zTXt", "iTXt"): + entry["text"] = _png_text_decode(ctype, body) + if entry["text"].startswith("XML:com.adobe.xmp"): + entry["kind"] = "xmp" + elif ctype == "tIME" and length == 7: + y, mo, d, h, mi, s = struct.unpack(">HBBBBB", body) + entry["time"] = f"{y:04d}-{mo:02d}-{d:02d}T{h:02d}:{mi:02d}:{s:02d}Z" + elif ctype == "gAMA" and length == 4: + entry["gamma"] = struct.unpack(">I", body)[0] / 100000 + elif ctype == "sRGB" and length == 1: + entry["rendering_intent"] = body[0] + elif ctype == "iCCP": + nul = body.find(b"\x00") + if nul > 0: + entry["profile_name"] = body[:nul].decode("latin-1", "replace") + entry["base64"] = _b64(body) + elif ctype == "iDOT": + # present in iOS/macOS screenshots + entry["apple_screenshot_marker"] = True + elif ctype in ("IHDR", "IDAT"): + pass # pixel-data / header chunks: length is signal enough + elif length: + entry["base64"] = _b64(body) + chunks.append(entry) + pos += 12 + length + if ctype == "IEND": + post_iend = data[pos:] + break + except Exception as exc: + chunks.append({"error": _safe_str(exc)}) + return chunks, post_iend + + +def _set_jpeg_trailer(result: dict[str, Any], data: bytes, eoi: int) -> None: + """Preserve bytes after JPEG EOI for Samsung Galaxy AI detection.""" + trailer = data[eoi + 2 :] + result["post_eoi_bytes"] = len(trailer) + if trailer: + result["post_eoi_base64"] = _b64(trailer) + + +def read_jpeg_segments(data: bytes) -> dict[str, Any]: + """Every JPEG APP segment in order, plus post-EOI trailer size. + + XMP APP1 segments are kept as full text; every other segment body is + kept as full base64 (1 MB ceiling per segment). + """ + result: dict[str, Any] = {"segments": [], "post_eoi_bytes": 0} + try: + pos = 2 + while pos + 4 <= len(data): + if data[pos] != 0xFF: + break + marker = data[pos + 1] + if marker == 0xD9: # EOI + _set_jpeg_trailer(result, data, pos) + break + if marker == 0xDA: # SOS: entropy-coded data follows + eoi = data.rfind(b"\xff\xd9") + if eoi != -1: + _set_jpeg_trailer(result, data, eoi) + break + if not (0xE0 <= marker <= 0xEF): + length = struct.unpack(">H", data[pos + 2 : pos + 4])[0] + pos += 2 + length + continue + length = struct.unpack(">H", data[pos + 2 : pos + 4])[0] + body = data[pos + 4 : pos + 2 + length] + name = f"APP{marker - 0xE0}" + entry: dict[str, Any] = {"marker": name, "length": length} + # Adobe JPEG XMP APP1 magic (namespace URI in the packet, not a request). + if body.startswith(b"http://ns.adobe.com/xap/1.0/\x00"): # NOSONAR + entry["kind"] = "xmp" + entry["text"] = body[29:].decode("utf-8", "replace") + elif name == "APP2" and body.startswith(b"MPF\x00"): + # Multi-Picture Format: Ultra HDR gain map, Samsung dual shot + entry["kind"] = "mpf" + entry["base64"] = _b64(body) + elif name == "APP2" and body.startswith(b"ICC_PROFILE"): + entry["kind"] = "icc" + entry["base64"] = _b64(body) + elif name == "APP2" and body.startswith(b"FPXR"): + entry["kind"] = "flashpix" + entry["base64"] = _b64(body) + elif name == "APP11": + entry["kind"] = "c2pa_or_jumbf" + # the parsed manifest is in c2pa_store, but the raw JUMBF + # also carries assertion thumbnails the JSON may omit + entry["base64"] = _b64(body) + elif body.startswith(b"Exif\x00\x00"): + entry["kind"] = "exif" + entry["base64"] = _b64(body) + elif body.startswith(b"Photoshop 3.0\x00"): + entry["kind"] = "iptc_iim" + entry["base64"] = _b64(body) + else: + entry["base64"] = _b64(body) + result["segments"].append(entry) + pos += 2 + length + except Exception as exc: + result["error"] = _safe_str(exc) + return result + + +def read_pil_info(path: Path) -> tuple[dict[str, Any], dict[str, Any], bytes | None]: + """One Image.open serving all PIL-derived data: container basics, + img.info passthrough (XMP, comments), the IPTC-IIM dataset, and the + raw EXIF blob (for the caller's piexif parse on PNG/WebP/HEIC).""" + out: dict[str, Any] = {} + iptc: dict[str, Any] = {} + exif_blob: bytes | None = None + try: + with Image.open(path) as img: + out["format"] = img.format + out["mode"] = img.mode + out["width"], out["height"] = img.size + out["n_frames"] = getattr(img, "n_frames", 1) + dpi = img.info.get("dpi") + if dpi: + out["dpi"] = [round(float(d), 2) for d in dpi] + icc = img.info.get("icc_profile") + if icc: + out["icc_profile"] = { + "length": len(icc), + # header: profile class, color space, PCS (bytes 12-24) + "header_hex": icc[12:24].hex() if len(icc) >= 24 else "", + "base64": _b64(icc), + } + blob = img.info.get("exif") + if isinstance(blob, bytes): + exif_blob = blob + try: + info = getiptcinfo(img) + except Exception: + info = None + if info: + iptc = {f"{k[0]}:{k[1]}": _decode_exif_value(v) for k, v in info.items()} + for key, value in img.info.items(): + if key in ("icc_profile", "exif", "dpi"): + continue + if isinstance(value, bytes): + try: + out[f"info:{key}"] = value.decode("utf-8", "strict") + except (UnicodeDecodeError, ValueError): + out[f"info:{key}"] = f"base64:{_b64(value)}" + else: + out[f"info:{key}"] = _safe_str(value) + except Exception as exc: + out["error"] = _safe_str(exc) + return out, iptc, exif_blob + + +def read_c2pa_store(path: Path) -> dict[str, Any]: + """Full C2PA manifest store through the package's cached reader.""" + from remove_ai_watermarks._internal.c2pa import read_manifest_store_json + + raw = read_manifest_store_json(path) + if raw is None: + return {} + try: + value: Any = json.loads(raw) + return ( + cast("dict[str, Any]", value) + if isinstance(value, dict) + else {"error": "C2PA manifest store is not an object"} + ) + except (TypeError, ValueError) as exc: + return {"error": _safe_str(exc)} + + +def sniff_format(head: bytes) -> str: + if head.startswith(b"\x89PNG"): + return "png" + if head.startswith(b"\xff\xd8"): + return "jpeg" + if head.startswith(b"RIFF") and head[8:12] == b"WEBP": + return "webp" + if head[:6] in (b"GIF87a", b"GIF89a"): + return "gif" + if head.startswith(b"BM"): + return "bmp" + if head.startswith((b"II*\x00", b"MM\x00*")): + return "tiff" + if head[4:8] == b"ftyp": + return f"isobmff:{head[8:12].decode('latin-1', 'replace')}" + return f"unknown:{head[:16].hex()}" + + +# --- JPEG encoder structure (metadata layer) --- + + +def _jpeg_forensics_bytes(data: bytes) -> dict[str, Any]: + """Structure-level JPEG forensics: DQT tables (encoder fingerprint), + SOF type (baseline/progressive) + chroma subsampling, DHT Huffman + tables (custom = optimizing encoder), per-scan spectral selection + (progressive scan script), JFIF/Adobe app markers, COM, DRI.""" + out: dict[str, Any] = {} + try: + if not data.startswith(b"\xff\xd8"): + return out + pos = 2 + scans: list[dict[str, int]] = [] + dqt: dict[str, list[int]] = {} + dht: list[str] = [] + comments: list[str] = [] + while pos + 4 <= len(data): + if data[pos] != 0xFF: + break + marker = data[pos + 1] + if marker in (0xD8, 0x01) or 0xD0 <= marker <= 0xD7: + pos += 2 + continue + if marker == 0xD9: + break + length = struct.unpack(">H", data[pos + 2 : pos + 4])[0] + body = data[pos + 4 : pos + 2 + length] + if marker == 0xDB: # DQT + off = 0 + while off < len(body): + tid = body[off] & 0x0F + prec = body[off] >> 4 + n = 128 if prec else 64 + vals = list(body[off + 1 : off + 1 + n]) + if prec: # 16-bit entries + vals = [struct.unpack(">H", bytes(vals[i : i + 2]))[0] for i in range(0, len(vals) - 1, 2)] + dqt[str(tid)] = vals[:64] + off += 1 + n + elif marker == 0xC4: # DHT: custom tables mean an optimizing encoder + dht.append(body.hex()) + elif marker == 0xDD and len(body) >= 2: # DRI + out["restart_interval"] = struct.unpack(">H", body[:2])[0] + elif marker == 0xE0 and body.startswith(b"JFIF\x00") and len(body) >= 12: + out["jfif"] = { + "version": f"{body[5]}.{body[6]}", + "density_units": body[7], + "x_density": struct.unpack(">H", body[8:10])[0], + "y_density": struct.unpack(">H", body[10:12])[0], + } + elif marker == 0xEE and body.startswith(b"Adobe") and len(body) >= 12: + out["adobe_transform"] = body[11] + elif marker in (0xC0, 0xC1, 0xC2) and len(body) >= 6: + out["progressive"] = marker == 0xC2 + out["precision_bits"] = body[0] + out["sof_height"] = struct.unpack(">H", body[1:3])[0] + out["sof_width"] = struct.unpack(">H", body[3:5])[0] + comps: list[dict[str, int]] = [] + for i in range(body[5]): + c = body[6 + i * 3 : 9 + i * 3] + if len(c) == 3: + comps.append({"h": c[1] >> 4, "v": c[1] & 0x0F, "tq": c[2]}) + if len(comps) >= 3: + lum = comps[0] + subs = {1: "4:4:4", 2: "4:2:2"}.get(lum["h"] * lum["v"]) + out["subsampling"] = subs or f"{lum['h']}x{lum['v']}" + elif marker == 0xFE: # COM + comments.append(body.decode("utf-8", "replace")[:2000]) + elif marker == 0xDA: + # SOS spectral selection: the progressive scan script + # differs across libjpeg / mozjpeg / Photoshop + if len(body) >= 3: + ns = body[0] + tail = body[1 + ns * 2 :] + if len(tail) >= 3: + scans.append({"ss": tail[0], "se": tail[1], "ah": tail[2] >> 4, "al": tail[2] & 0x0F}) + # skip entropy-coded data to the next marker + end = data.find(b"\xff\xd9", pos) + nxt = data.find(b"\xff", pos + 2) + while nxt != -1 and nxt + 1 < len(data) and data[nxt + 1] == 0x00: + nxt = data.find(b"\xff", nxt + 2) + if nxt == -1 or (end != -1 and nxt >= end): + break + pos = nxt + continue + pos += 2 + length + if dqt: + out["quant_tables"] = dqt + if dht: + out["huffman_tables_hex"] = dht + if comments: + out["comments"] = comments + if scans: + out["scan_count"] = len(scans) + out["scan_script"] = scans + except Exception as exc: + out["error"] = _safe_str(exc) + return out + + +def read_webp_chunks(data: bytes) -> list[dict[str, Any]]: + """WebP RIFF chunk inventory (VP8X/VP8/VP8L/EXIF/XMP/ICCP/ANIM...).""" + chunks: list[dict[str, Any]] = [] + try: + pos = 12 + declared_end = 8 + struct.unpack("= 12 else len(data) + container_end = min(len(data), declared_end) + while pos + 8 <= container_end: + chunk_type = data[pos : pos + 4] + ctype = chunk_type.decode("latin-1") + length = struct.unpack(" list[dict[str, Any]]: + """Stream metadata chunks after ``window`` while seeking over coded frames.""" + chunks: list[dict[str, Any]] = [] + try: + file_size = path.stat().st_size + with open(path, "rb") as handle: + header = handle.read(12) + if len(header) < 12 or not header.startswith(b"RIFF") or header[8:12] != b"WEBP": + return chunks + container_end = min(file_size, 8 + struct.unpack("= window: + handle.seek(start) + body = handle.read(min(safe_length, _B64_CAP)) + entry: dict[str, Any] = { + "type": chunk_type.decode("latin-1"), + "length": length, + "base64": _b64(body), + } + if len(body) < safe_length: + entry["truncated"] = True + chunks.append(entry) + position = start + safe_length + (safe_length & 1) + except (OSError, struct.error) as exc: + chunks.append({"error": _safe_str(exc)}) + return chunks + + +def sha256_of(data: bytes) -> str: + return hashlib.sha256(data).hexdigest() + + +def xattr_where_from(path: Path) -> list[str]: + """macOS download-source URLs (kMDItemWhereFroms), empty elsewhere.""" + try: + getter = cast("Any", os.getxattr) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] + raw = cast("bytes", getter(path, "com.apple.metadata:kMDItemWhereFroms")) + value = plistlib.loads(raw) + values = cast("list[Any]", value) if isinstance(value, list) else [value] + return [str(item) for item in values] + except (AttributeError, OSError, ValueError): + return [] + + +def xattr_quarantine(path: Path) -> str | None: + """macOS quarantine string: flags; timestamp; downloading agent (Safari, + Telegram, Chrome...). Presence alone means 'came from the internet'.""" + try: + getter = cast("Any", os.getxattr) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] + raw = cast("bytes", getter(path, "com.apple.quarantine")) + return raw.decode("utf-8", "replace")[:500] + except (AttributeError, OSError): + return None + + +def read_isobmff_inventory(data: bytes) -> dict[str, Any]: + """HEIC/AVIF/MOV box inventory: top-level boxes plus the meta item + types (Exif, mime=XMP, auxl depth/gain-map, aae Apple-edits plist, + irot derived images). Strong phone-provenance signal.""" + out: dict[str, Any] = {} + try: + stream = io.BytesIO(data) + + def boxes(start: int, end: int) -> list[tuple[str, int, int]]: + return [ + (box_type.decode("latin-1"), payload_offset, box_end) + for _, box_end, box_type, payload_offset in iter_file_boxes(stream, start, end) + ] + + top = boxes(0, len(data)) + out["boxes"] = [t for t, _, _ in top] + provenance_boxes: list[dict[str, Any]] = [] + for t, s, e in top: + if t.encode("latin-1") in C2PA_BOX_TYPES: + provenance_boxes.append( + {"type": t, "length": e - s, "base64": _b64(data[s:e], cap=_PROVENANCE_B64_CAP)} + ) + if t == "moov": + for ct, cs, ce in boxes(s, e): + if ct == "mvhd" and ce - cs >= 24: + # full box + creation/modification times (1904 epoch) + version = data[cs] + base = cs + 4 + creation = struct.unpack(">I", data[base : base + 4])[0] if version == 0 else None + if creation: + out["mvhd_creation_time"] = creation - 2082844800 + elif t == "meta": + # full box: 4 bytes version/flags, then child boxes + for ct, cs, ce in boxes(s + 4, e): + if ct == "iinf": + # full box + entry count, then infe entries + count = struct.unpack(">H", data[cs + 4 : cs + 6])[0] + out["meta_item_count"] = count + item_types: list[str] = [] + for it, is_, ie in boxes(cs + 6, ce): + if it == "infe" and ie - is_ >= 8: + # infe full box: version(1)+flags(3), then + # v2: item_ID(2)+protection(2)+item_type(4) + # v3: item_ID(4)+protection(2)+item_type(4) + version = data[is_] + off = is_ + 4 + (4 if version == 3 else 2) + 2 + if off + 4 <= ie: + item_types.append(data[off : off + 4].decode("latin-1", "replace")) + if item_types: + out["meta_item_types"] = sorted(set(item_types)) + elif ct == "iprp": + out["has_iprp"] = True + for pt, ps, pe in boxes(cs, ce): + if pt == "ipco": + props = [t for t, _, _ in boxes(ps, pe)] + out["ipco_properties"] = props + # auxC holds the auxiliary image type URN + for box_type, qs, qe in boxes(ps, pe): + if box_type == "auxC": + out["auxc_types"] = ( + data[qs + 4 : qe].split(b"\x00")[0].decode("latin-1", "replace") + ) + elif ct == "iref": + out["has_iref"] = True + if provenance_boxes: + out["provenance_boxes"] = provenance_boxes + # QuickTime metadata keys (©mak/©mod/©swr) for the MOV side of + # Live Photos: tolerant printable-string grab after each atom + qt: dict[str, str] = {} + for atom, key in ((b"\xa9mak", "make"), (b"\xa9mod", "model"), (b"\xa9swr", "software")): + idx = data.find(atom) + if idx != -1: + m = re.search(rb"[ -~]{4,80}", data[idx + 4 : idx + 200]) + if m: + qt[key] = m.group(0).decode("ascii", "replace") + if qt: + out["quicktime"] = qt + except Exception as exc: + out["error"] = _safe_str(exc) + return out + + +def read_isobmff_provenance_path(path: Path) -> dict[str, Any]: + """Stream top-level ISOBMFF boxes and preserve provenance payloads. + + This is the large-file counterpart to :func:`read_isobmff_inventory`. + It seeks over media payloads instead of loading them into memory. + """ + out: dict[str, Any] = {"boxes": []} + provenance_boxes: list[dict[str, Any]] = [] + collected = 0 + try: + file_size = path.stat().st_size + with open(path, "rb") as f: + for _, box_end, box_type_raw, payload_offset in iter_file_boxes(f, 0, file_size): + box_type = box_type_raw.decode("latin-1") + out["boxes"].append(box_type) + payload_length = box_end - payload_offset + if box_type_raw in C2PA_BOX_TYPES and collected < _PROVENANCE_B64_CAP: + to_read = min(payload_length, _PROVENANCE_B64_CAP - collected) + f.seek(payload_offset) + payload = f.read(to_read) + entry: dict[str, Any] = { + "type": box_type, + "length": payload_length, + "base64": _b64(payload, cap=_PROVENANCE_B64_CAP), + } + if to_read < payload_length: + entry["truncated"] = True + provenance_boxes.append(entry) + collected += len(payload) + except (OSError, struct.error) as exc: + out["error"] = _safe_str(exc) + if provenance_boxes: + out["provenance_boxes"] = provenance_boxes + return out + + +def read_png_late_metadata_path(path: Path, window: int = _RAW_SCAN_HEAD) -> list[dict[str, Any]]: + """Stream PNG metadata chunks whose payload starts after ``window``.""" + chunks: list[dict[str, Any]] = [] + try: + file_size = path.stat().st_size + with open(path, "rb") as f: + if f.read(8) != b"\x89PNG\r\n\x1a\n": + return chunks + pos = 8 + while pos + 12 <= file_size: + f.seek(pos) + header = f.read(8) + if len(header) < 8: + break + length, chunk_type = struct.unpack(">I4s", header) + data_start = pos + 8 + safe_length = max(0, min(length, file_size - data_start)) + if chunk_type in PNG_METADATA_CHUNKS and data_start >= window: + body = f.read(min(safe_length, _B64_CAP)) + entry: dict[str, Any] = { + "type": chunk_type.decode("latin-1"), + "length": length, + "base64": _b64(body), + } + if len(body) < safe_length: + entry["truncated"] = True + chunks.append(entry) + pos = data_start + safe_length + 4 + if chunk_type == b"IEND": + break + except (OSError, struct.error) as exc: + chunks.append({"error": _safe_str(exc)}) + return chunks + + +def apple_live_photo_id(head: bytes) -> str | None: + """Apple Live Photo content identifier (links the still to its MOV). + + The UUID sits in the Apple MakerNote (tag 17) of the still and in the + MOV metadata; a raw head scan finds it in either container.""" + # the UUID string sits next to "content.identifier" in the MOV, but in + # the STILL it is a bare UUID inside the Apple MakerNote (whose header + # is "Apple iOS"), so gate on either marker + if b"content.identifier" not in head and b"com.apple.quicktime" not in head and b"Apple iOS" not in head: + return None + m = re.search(rb"[0-9A-Fa-f]{8}-[0-9A-Fa-f]{4}-[0-9A-Fa-f]{4}-[0-9A-Fa-f]{4}-[0-9A-Fa-f]{12}", head) + return m.group(0).decode("ascii") if m else None + + +_MAX_FULL_READ = 256 << 20 # files bigger than this are scanned head-only +_HEAD_READ = 4 << 20 + + +def _sha256_stream(path: Path) -> str: + h = hashlib.sha256() + with open(path, "rb") as f: + for block in iter(lambda: f.read(1 << 20), b""): + h.update(block) + return h.hexdigest() + + +def collect_forensic_metadata( + path: Path, + *, + schema_version: int = FORENSIC_METADATA_SCHEMA_VERSION, +) -> dict[str, Any]: + """Collect the versioned, metadata-only forensic record for ``path``. + + This broad inspection record is not provenance-detector input. Use + :func:`remove_ai_watermarks.metadata_record.collect_metadata_record` for the + strict record accepted by ``identify_metadata_record``. Long-lived consumers + should request the schema they implement; unsupported versions raise before the + source is read. + """ + schema_version = require_schema_version( + schema_version, + contract="forensic metadata", + supported=(1,), + ) + image_io._register_heif() # pyright: ignore[reportPrivateUsage] + stat = path.stat() + oversized = stat.st_size > _MAX_FULL_READ + if oversized: + data = None + with open(path, "rb") as f: + head = f.read(_HEAD_READ) + else: + data = path.read_bytes() + head = data + record: dict[str, Any] = { + "schema_version": schema_version, + "record_type": FORENSIC_METADATA_RECORD_TYPE, + "file": str(path), + "name": path.name, + "extension": path.suffix.lower(), + "size_bytes": stat.st_size, + "mtime": stat.st_mtime, + "birthtime": getattr(stat, "st_birthtime", None), + "sha256": _sha256_stream(path) if data is None else sha256_of(data), + "content_format": sniff_format(head), + } + if oversized: + # Preserve the same bounded byte windows used by downstream provenance + # algorithms while path-based readers (PIL, piexif, C2PA) run normally. + record["oversized"] = {"head_scanned_bytes": len(head)} + record["raw_metadata_windows"] = {"head_base64": _b64(head[:_RAW_SCAN_HEAD])} + if stat.st_size > _RAW_SCAN_TAIL: + with open(path, "rb") as f: + f.seek(-_RAW_SCAN_TAIL, 2) + record["raw_metadata_windows"]["tail_base64"] = _b64(f.read()) + where_from = xattr_where_from(path) + if where_from: + record["download_source_urls"] = where_from + quarantine = xattr_quarantine(path) + if quarantine: + record["quarantine"] = quarantine + live_photo_id = apple_live_photo_id(head[: 2 << 20]) + if live_photo_id: + record["live_photo_content_id"] = live_photo_id + record["pil"], record["iptc"], exif_blob = read_pil_info(path) + record["exif"], thumbnail = read_full_exif(path, exif_blob, data) + record["c2pa_store"] = read_c2pa_store(path) + if data is not None: + fmt = record["content_format"] + if fmt == "png": + record["png_chunks"], post_iend = read_png_chunks(data) + if post_iend: + record["png_post_iend_bytes"] = len(post_iend) + record["png_post_iend_base64"] = _b64(post_iend) + elif fmt == "jpeg": + record["jpeg"] = read_jpeg_segments(data) + record["jpeg_forensics"] = _jpeg_forensics_bytes(data) + elif fmt == "webp": + record["webp_chunks"] = read_webp_chunks(data) + elif fmt.startswith("isobmff"): + record["isobmff"] = read_isobmff_inventory(data) + elif record["content_format"] == "png": + late_chunks = read_png_late_metadata_path(path) + if late_chunks: + record["png_late_metadata_chunks"] = late_chunks + elif record["content_format"] == "webp": + late_chunks = read_webp_late_metadata_path(path) + if late_chunks: + record["webp_late_metadata_chunks"] = late_chunks + elif record["content_format"].startswith("isobmff"): + record["isobmff"] = read_isobmff_provenance_path(path) + if thumbnail: + record["has_exif_thumbnail"] = True + # the embedded thumbnail is its own JPEG; after an edit its encoder + # forensics commonly MISMATCH the main image (classic tamper tell) + thumb_forensics = _jpeg_forensics_bytes(thumbnail) + thumb_forensics["base64"] = _b64(thumbnail) + record["exif_thumbnail_forensics"] = thumb_forensics + return record diff --git a/src/remove_ai_watermarks/identify.py b/src/remove_ai_watermarks/identify.py index 3f84a38..2dfc324 100644 --- a/src/remove_ai_watermarks/identify.py +++ b/src/remove_ai_watermarks/identify.py @@ -40,6 +40,7 @@ from remove_ai_watermarks._internal.constants import ( C2PA_IDENTITY_AI_ORGS, C2PA_ISSUERS, ) +from remove_ai_watermarks._internal.schema import require_schema_version from remove_ai_watermarks.metadata import ( AI_METADATA_KEYS, AIGC_MARKERS, @@ -176,7 +177,32 @@ def _external_metadata(value: Any) -> tuple[list[tuple[str, Any]], bytes]: """Index nested metadata and recover common encoded binary values in one pass.""" pairs: list[tuple[str, Any]] = [] parts: list[bytes] = [] - diagnostic_keys = {"error", "kind"} + diagnostic_keys = { + "artifacts", + "birthtime", + "color", + "content_format", + "dct", + "ela", + "error", + "extension", + "fft", + "file", + "filename", + "full", + "gradient", + "kind", + "mtime", + "name", + "noise", + "path", + "pixel", + "provenance", + "sha256", + "signals", + "size_bytes", + "timing_ms", + } def visit(item: Any) -> None: if isinstance(item, dict): @@ -184,7 +210,6 @@ def _external_metadata(value: Any) -> tuple[list[tuple[str, Any]], bytes]: for key, nested in mapping.items(): key_text = str(key) pairs.append((key_text, nested)) - parts.append(key_text.encode("utf-8", "replace")) if key_text.lower() in diagnostic_keys: continue if isinstance(nested, str) and (key_text == "base64" or key_text.endswith("_base64")): @@ -248,17 +273,69 @@ def _external_exif_generator(pairs: list[tuple[str, Any]], scan: bytes) -> str | return generator_from_metadata(candidates, scan) +def _metadata_source_kind(info: dict[str, Any], scan: bytes) -> str | None: + """Normalize the source type wherever it is carried: C2PA or IPTC/XMP. + + A composite marker contains ``TrainedAlgorithmicMedia`` as a substring, so it + is removed before looking for a standalone full-generation marker. When a file + genuinely carries both kinds, full generation wins. + """ + structured = info.get("ai_source_kind") + without_composites = scan.replace(b"compositeWithTrainedAlgorithmicMedia", b"").replace(b"compositeSynthetic", b"") + generated = structured == "generated" or any( + marker in without_composites for marker in (b"trainedAlgorithmicMedia", b"TrainedAlgorithmicMedia") + ) + if generated: + return "generated" + if structured == "enhanced" or any( + marker in scan for marker in (b"compositeWithTrainedAlgorithmicMedia", b"compositeSynthetic") + ): + return "enhanced" + return None + + def evidence_from_metadata_record( record: dict[str, Any], *, path: Path, c2pa_manifest_store: str | dict[str, Any] | None = None ) -> ProvenanceEvidence: """Normalize an externally collected metadata record into provenance evidence. - The record may contain arbitrary nested dictionaries and lists. Text, bytes, - hexadecimal values prefixed with ``hex:``, and fields named ``base64`` or - ending in ``_base64`` are included in the shared byte scan. No source file is - opened. + Unversioned external records may contain arbitrary nested dictionaries and + lists. Versioned native records accept only the source-derived fields emitted by + ``collect_metadata_record``; other native record types and unknown schema + versions are rejected. No source file is opened. """ - pairs, scan = _external_metadata(record) + from remove_ai_watermarks.metadata_record import METADATA_RECORD_SCHEMA_VERSION, METADATA_RECORD_TYPE + + # Records produced by ``collect_metadata_record`` are a versioned transport + # contract. Only their source-derived fields are evidence: the filename, + # container label and schema bookkeeping describe the collector and must never + # become detector input. Shape-detect the pre-versioned form as well so records + # emitted by 0.26 remain safe and readable. + record_type = record.get("record_type") + if record_type not in (None, METADATA_RECORD_TYPE): + raise ValueError(f"Unsupported metadata record type: {record_type!r}") + if record_type == METADATA_RECORD_TYPE: + require_schema_version( + record.get("schema_version"), + contract="provenance metadata", + supported=(METADATA_RECORD_SCHEMA_VERSION,), + ) + status = record.get("status") + if status == "error": + raise ValueError("Provenance metadata collection failed") + if status != "complete": + raise ValueError(f"Unsupported provenance metadata collection status: {status!r}") + is_portable_record = record_type == METADATA_RECORD_TYPE or { + "container", + "metadata_base64", + "tail_base64", + }.issubset(record) + evidence_record = ( + {key: record[key] for key in ("metadata_base64", "tail_base64", "pil", "exif") if key in record} + if is_portable_record + else record + ) + pairs, scan = _external_metadata(evidence_record) store = c2pa_manifest_store if store is None: candidate = record.get("c2pa_store") @@ -365,13 +442,13 @@ class ProvenanceReport: is_ai_generated: bool | None # True / False is never asserted; None = unknown platform: str | None confidence: str # "high" | "medium" | "none" - # Coarse AI-origin kind from the C2PA digital-source-type, so a caller can - # branch on full generation vs an AI-touched real photo: + # Coarse AI-origin kind from a C2PA or standalone IPTC/XMP digital-source-type, + # so a caller can branch on full generation vs an AI-touched real photo: # "generated" -- digitalSourceType trainedAlgorithmicMedia (fully AI). # "enhanced" -- compositeWithTrainedAlgorithmicMedia (real content with an # AI-composited region; scrub the AI region, keep the photo). - # None -- no C2PA AI source-type (verdict, if AI, came from another - # signal: IPTC, AIGC, local gen params, xAI, ...). + # None -- no AI digital-source-type (verdict, if AI, came from another + # signal: AIGC, local gen params, xAI, ...). ai_source_kind: str | None = None # True when the AI verdict rests on a metadata or embedded-invisible signal # (C2PA AI issuer / SynthID proxy, IPTC, AIGC, local gen params, EXIF/xAI, or @@ -390,14 +467,24 @@ class ProvenanceReport: # inconsistent -- a strong tell of spoofed, transplanted, or laundered metadata. integrity_clashes: list[str] = field(default_factory=list[str]) - def to_dict(self) -> dict[str, Any]: + def to_dict( + self, + *, + schema_version: int = PROVENANCE_REPORT_SCHEMA_VERSION, + ) -> dict[str, Any]: """Return the versioned, JSON-safe verdict contract. ``path`` is deliberately omitted. It is extraction context, not part of the verdict, and local filesystem paths should not cross a service boundary. + Request an explicit schema for a long-lived transport consumer. """ + schema_version = require_schema_version( + schema_version, + contract="provenance report", + supported=(1,), + ) return { - "schema_version": PROVENANCE_REPORT_SCHEMA_VERSION, + "schema_version": schema_version, "is_ai_generated": self.is_ai_generated, "platform": self.platform, "confidence": self.confidence, @@ -997,12 +1084,7 @@ def _identify_from_evidence( # _internal.c2pa._populate_registry_fields (covers PNG + any container the c2pa-python # reader handles); fall back to a raw head scan for the non-PNG raw-blob path # where extract_c2pa_info returns {}. Full generation wins when both appear. - c2pa_source_kind = info.get("ai_source_kind") - if c2pa_source_kind is None: - if b"trainedAlgorithmicMedia" in head: - c2pa_source_kind = "generated" - elif b"compositeWithTrainedAlgorithmicMedia" in head: - c2pa_source_kind = "enhanced" + source_kind = _metadata_source_kind(info, head) # An identity-AI issuer (a pure-generator brand like Dreamina) asserts AI even # without a digitalSourceType -- some ByteDance/Dreamina manifests ship no # trainedAlgorithmicMedia, so the registered generator name is the only signal. @@ -1010,7 +1092,7 @@ def _identify_from_evidence( # does not reopen the incidental-mention problem the common-word issuers have. issuer_blob = " ".join(issuers) c2pa_identity_ai = has_c2pa and any(org in issuer_blob for org in C2PA_IDENTITY_AI_ORGS) - c2pa_is_ai = c2pa_source_kind is not None or c2pa_identity_ai + c2pa_is_ai = source_kind is not None or c2pa_identity_ai # Generator string (for the signal detail): structured for PNG, CBOR-scanned # for other containers. Best-effort -- some manifests key it as # `claim_generator_info` (Pixel), so this can be None even when a device is @@ -1066,7 +1148,7 @@ def _identify_from_evidence( # for its own callers; the verdict no longer depends on which extractor ran. synthid = meta.get("synthid_watermark") # The literal byte checks mirror `metadata.synthid_source` exactly rather than - # reusing the derived `has_c2pa` / `c2pa_source_kind` above, which are broader: + # reusing the derived `has_c2pa` / `source_kind` above, which are broader: # the file path's answer must not move. trained_source = b"trainedAlgorithmicMedia" in head or b"TrainedAlgorithmicMedia" in head if not synthid and trained_source and c2pa_marker_in(head) and (vendors := synthid_vendors_in(region)): @@ -1254,9 +1336,9 @@ def _identify_from_evidence( is_ai_generated=is_ai, platform=platform, confidence=confidence, - # Only meaningful when the AI verdict actually came from the C2PA source - # type; a non-C2PA AI signal (IPTC/AIGC/local gen) leaves it None. - ai_source_kind=c2pa_source_kind if (is_ai and has_c2pa) else None, + # Meaningful for the same digitalSourceType whether carried by C2PA or a + # standalone IPTC/XMP label. Other AI signals leave it None. + ai_source_kind=source_kind if (is_ai and (has_c2pa or iptc)) else None, ai_from_metadata=ai_from_metadata, watermarks=watermarks, signals=signals, diff --git a/src/remove_ai_watermarks/metadata.py b/src/remove_ai_watermarks/metadata.py index 0c097f3..eeae2e7 100644 --- a/src/remove_ai_watermarks/metadata.py +++ b/src/remove_ai_watermarks/metadata.py @@ -18,6 +18,11 @@ if TYPE_CHECKING: from collections.abc import Callable, Iterable from pathlib import Path +from remove_ai_watermarks._internal.constants import ( + PNG_METADATA_CHUNKS, + RIFF_METADATA_CHUNKS, +) + logger = logging.getLogger(__name__) # Smaller scan_head window for the cheap marker checks (has_ai_metadata, @@ -236,11 +241,6 @@ def _is_ai_value(value: str) -> bool: return any(token in value_lower for token in AI_GENERATOR_TOKENS) -# PNG ancillary chunks that can carry provenance metadata (XMP, EXIF, text). -# Never IDAT -- that is the compressed pixel stream. -_PNG_META_CHUNKS: frozenset[bytes] = frozenset({b"tEXt", b"iTXt", b"zTXt", b"eXIf", b"iCCP"}) - - def _png_late_metadata(image_path: Path, window: int) -> bytes: """Payloads of PNG metadata chunks that start *beyond* the first ``window`` bytes, found by seeking past the (large) ``IDAT`` pixel stream. @@ -272,7 +272,7 @@ def _png_late_metadata(image_path: Path, window: int) -> bytes: # Clamp the attacker-controlled 32-bit length to the bytes that # actually remain, so a malformed huge length can't allocate GBs. safe_length = max(0, min(length, file_size - data_start)) - if chunk_type in _PNG_META_CHUNKS and data_start >= window: + if chunk_type in PNG_METADATA_CHUNKS and data_start >= window: f.seek(data_start) out += f.read(safe_length) # Advance by the CLAMPED length: a malformed/inflated `length` that @@ -285,10 +285,6 @@ def _png_late_metadata(image_path: Path, window: int) -> bytes: return bytes(out) -# RIFF/WebP chunks that carry metadata rather than coded pixels. -_RIFF_META_CHUNKS: frozenset[bytes] = frozenset({b"EXIF", b"XMP ", b"ICCP", b"C2PA"}) - - def _riff_late_metadata(image_path: Path, window: int, *, max_total: int = 4 * 1024 * 1024) -> bytes: """Payloads of RIFF metadata chunks that start *beyond* the first ``window`` bytes, found by stepping over the (large) coded-image chunk. @@ -311,8 +307,13 @@ def _riff_late_metadata(image_path: Path, window: int, *, max_total: int = 4 * 1 return b"" f.seek(0, 2) file_size = f.tell() + f.seek(4) + declared_size = f.read(4) + if len(declared_size) < 4: + return b"" + container_end = min(file_size, 8 + struct.unpack("= window: + safe_length = max(0, min(length, container_end - start)) + if chunk_type in RIFF_METADATA_CHUNKS and start >= window: f.seek(start) out += f.read(min(safe_length, max_total - len(out))) position = start + safe_length + (safe_length & 1) # chunks are word-aligned @@ -1603,4 +1604,5 @@ QUICK_SCAN_BYTES = _QUICK_SCAN_BYTES SAMSUNG_EDITOR_MARKER = _SAMSUNG_EDITOR_MARKER read_file_tail = _read_file_tail png_late_metadata = _png_late_metadata +riff_late_metadata = _riff_late_metadata exif_text = _exif_text diff --git a/src/remove_ai_watermarks/metadata_record.py b/src/remove_ai_watermarks/metadata_record.py index 8c414e2..ade323b 100644 --- a/src/remove_ai_watermarks/metadata_record.py +++ b/src/remove_ai_watermarks/metadata_record.py @@ -45,7 +45,8 @@ from typing import TYPE_CHECKING, Any if TYPE_CHECKING: from pathlib import Path -from remove_ai_watermarks._internal.constants import PNG_SIGNATURE +from remove_ai_watermarks._internal.constants import PNG_SIGNATURE, RIFF_CODED_IMAGE_CHUNKS +from remove_ai_watermarks._internal.schema import require_schema_version from remove_ai_watermarks.metadata import ( QUICK_SCAN_BYTES, SAMSUNG_EDITOR_MARKER, @@ -73,8 +74,11 @@ UNKNOWN_TRAILER_WINDOW = 64 * 1024 # PNG text keys the file path reads for a generator tag, in ITS order. NovelAI stamps # Software/Source/Title rather than EXIF, and the first match wins, so order matters. _GENERATOR_TEXT_KEYS = ("Software", "Source", "Title", "Description") -# RIFF chunks holding coded pixels rather than metadata. -_RIFF_IMAGE_CHUNKS = frozenset({b"VP8 ", b"VP8L", b"ALPH"}) +# Stable transport contract for records produced by this module. The version is +# deliberately separate from the verdict version: collection and interpretation can +# evolve independently as long as old records remain readable. +METADATA_RECORD_SCHEMA_VERSION = 1 +METADATA_RECORD_TYPE = "provenance_metadata" def _jpeg_regions(data: bytes) -> bytes: @@ -145,14 +149,15 @@ def _riff_regions(data: bytes) -> bytes: input models. """ out = bytearray(data[:12]) # 'RIFF' + size + 'WEBP' - size = len(data) + declared_end = 8 + struct.unpack("= 12 else len(data) + size = min(len(data), declared_end) position = 12 while position + 8 <= size: chunk_type = data[position : position + 4] (length,) = struct.unpack(" tuple[str, bytes]: appended bytes as chunks, inflating the record and creating false signals. """ from remove_ai_watermarks._internal.isobmff import is_isobmff - from remove_ai_watermarks.metadata import png_late_metadata + from remove_ai_watermarks.metadata import png_late_metadata, riff_late_metadata if head.startswith(b"\xff\xd8"): return "jpeg", _jpeg_regions(head) @@ -198,7 +203,7 @@ def _container_regions(image_path: Path, head: bytes) -> tuple[str, bytes]: # past the window; the same seek-past-IDAT reader the file path uses gets them. return "png", _png_regions(head) + png_late_metadata(image_path, HEAD_WINDOW) if head.startswith(b"RIFF") and head[8:12] == b"WEBP": - return "webp", _riff_regions(head) + return "webp", _riff_regions(head) + riff_late_metadata(image_path, HEAD_WINDOW) if is_isobmff(head): return "isobmff", _isobmff_regions(image_path, head) return "unknown", head @@ -223,6 +228,31 @@ def _trailer(image_path: Path, container: str) -> bytes: or one whose end lies before the window) the window is kept as-is, bounded -- that is what a byte scan of the same file would have seen anyway. """ + if container == "webp": + # RIFF declares its structural end in bytes 4..8. A fixed tail window is + # normally the last animation/frame payload, not a trailer, so preserve + # only bytes appended after the declared RIFF container. + try: + with open(image_path, "rb") as handle: + header = handle.read(12) + if len(header) < 12 or not header.startswith(b"RIFF"): + return b"" + declared_end = 8 + struct.unpack("= file_size: + return b"" + handle.seek(declared_end) + return handle.read(min(file_size - declared_end, UNKNOWN_TRAILER_WINDOW)) + except OSError as exc: + logger.debug("RIFF trailer read failed for %s: %s", image_path, exc) + return b"" + if container == "isobmff": + # ISOBMFF has no out-of-container trailer convention. Its bounded box + # walkers already collect late provenance while skipping ``mdat``; keeping + # a blind tail here would carry coded media bytes. + return b"" + tail = read_file_tail(image_path, TAIL_WINDOW) if SAMSUNG_EDITOR_MARKER in tail: # Galaxy AI splits its evidence: the marker sits in the post-EOI trailer, but @@ -311,27 +341,52 @@ def _pil_info(info: dict[str, Any]) -> dict[str, str]: return out -def collect_metadata_record(image_path: Path) -> dict[str, Any]: +def collect_metadata_record( + image_path: Path, + *, + schema_version: int = METADATA_RECORD_SCHEMA_VERSION, +) -> dict[str, Any]: """Collect everything the provenance verdict reads, as a JSON-safe record. The record is the transport format for :func:`identify.evidence_from_metadata_record`: it carries the metadata regions - (base64), the C2PA manifest store, and PIL's info mapping, and it never carries - pixel data. + (base64), the C2PA manifest store, and PIL's info mapping without carrying the + primary coded-pixel stream. Schema and collection status are explicit so a + consumer cannot mistake a failed read for an unknown provenance verdict. Args: image_path: Path to the image. + schema_version: Output schema implemented by the consumer. Returns: - A JSON-serializable dict. ``metadata_base64`` holds the concatenated - container regions, ``tail_base64`` the file trailer. + A versioned JSON-serializable dict. ``metadata_base64`` holds the + concatenated container regions, ``tail_base64`` the file trailer. """ + schema_version = require_schema_version( + schema_version, + contract="provenance metadata", + supported=(1,), + ) + from remove_ai_watermarks._internal.c2pa import read_manifest_store_json + try: + image_path.stat() + status = "complete" + issues: list[dict[str, str]] = [] + except OSError as exc: + logger.debug("metadata source unavailable for %s: %s", image_path, exc) + status = "error" + issues = [{"stage": "source", "code": "unavailable"}] + container, regions = _container_regions(image_path, _raw_head(image_path)) info = _decoder_info(image_path) record: dict[str, Any] = { + "schema_version": schema_version, + "record_type": METADATA_RECORD_TYPE, + "status": status, + "issues": issues, "container": container, "name": image_path.name, "metadata_base64": base64.b64encode(regions).decode("ascii"), diff --git a/src/remove_ai_watermarks/pixel_evidence.py b/src/remove_ai_watermarks/pixel_evidence.py index 077dfdc..a78cccc 100644 --- a/src/remove_ai_watermarks/pixel_evidence.py +++ b/src/remove_ai_watermarks/pixel_evidence.py @@ -1,13 +1,11 @@ # pyright: reportUnknownMemberType=false, reportUnknownArgumentType=false, reportUnknownVariableType=false, reportMissingTypeStubs=false -"""Experimental: the complete pixel-forensics layer for one image. +"""The complete pixel-forensics layer for one image. STATUS -Experimental and unused. Nothing in this package reads it -- not the provenance -verdict, not removal, not the CLI. It is here because the research scanner that -produced these measurements is gone, and the capability was worth keeping: whatever -asks for pixel forensics next starts from a tested implementation instead of -rebuilding one. Treat the shape as unstable until a caller exists. +Independent from provenance verdicts, removal, and the CLI. Consumers use the +versioned :meth:`PixelEvidence.to_dict` boundary; feature extraction failures are +reported per family without discarding successful measurements. WHAT IS MEASURED @@ -16,7 +14,7 @@ One decode, then six families of scale-robust statistics over it: * ``dct`` -- AC coefficient histograms over the 8x8 block DCT, plus the deviation of leading digits from Benford's law. * ``fft`` -- radial band energies of the log-magnitude spectrum, plus the - colour-filter-array periodicity peaks a demosaiced camera capture leaves. + color-filter-array periodicity peaks a demosaiced camera capture leaves. * ``noise`` -- standard deviation and kurtosis of a high-pass residual. * ``ela`` -- error level after a quality-90 JPEG re-save. * ``gradient`` -- gradient-magnitude histogram and Laplacian variance. @@ -46,9 +44,12 @@ from __future__ import annotations import base64 import io import logging +import time from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any +from remove_ai_watermarks._internal.schema import require_schema_version + if TYPE_CHECKING: from pathlib import Path @@ -63,6 +64,7 @@ FFT_BANDS = 8 # A Bayer CFA shows as symmetric peaks at half the Nyquist on the diagonals. BAYER_OFFSETS = ((1, 1), (1, -1)) INSTALL_HINT = "install the pixel extra: uv add 'remove-ai-watermarks[pixels]'" +PIXEL_EVIDENCE_SCHEMA_VERSION = 1 @dataclass(frozen=True) @@ -86,12 +88,48 @@ class PixelEvidence: color: dict[str, Any] = field(default_factory=dict[str, Any]) # Identifies the source image; see the module note. Empty unless asked for. artifacts: dict[str, Any] = field(default_factory=dict[str, Any]) + # Opt-in timings for callers measuring pipeline latency. Empty by default so + # repeated evidence collection remains value-deterministic. + timing_ms: dict[str, float] = field(default_factory=dict[str, float]) @property def decoded(self) -> bool: """False when the source could not be decoded at all.""" return "error" not in self.decode + @property + def status(self) -> str: + """``complete``, ``partial`` for a failed family, or ``error`` on decode.""" + if not self.decoded: + return "error" + sections = (self.dct, self.fft, self.noise, self.ela, self.gradient, self.color, self.artifacts) + return "partial" if any("error" in section for section in sections) else "complete" + + def to_dict( + self, + *, + schema_version: int = PIXEL_EVIDENCE_SCHEMA_VERSION, + ) -> dict[str, Any]: + """Return the selected JSON-safe transport schema without a local path.""" + schema_version = require_schema_version( + schema_version, + contract="pixel evidence", + supported=(1,), + ) + return { + "schema_version": schema_version, + "status": self.status, + "decode": dict(self.decode), + "dct": dict(self.dct), + "fft": dict(self.fft), + "noise": dict(self.noise), + "ela": dict(self.ela), + "gradient": dict(self.gradient), + "color": dict(self.color), + "artifacts": dict(self.artifacts), + "timing_ms": dict(self.timing_ms), + } + def is_available() -> bool: """True when the optional pixel dependencies are installed.""" @@ -120,14 +158,17 @@ def _dct_matrix(np: Any, n: int = 8) -> Any: def read_gray(image_path: Path) -> tuple[Any, Any, dict[str, Any]]: - """Decode to float32 grayscale (and RGB for colour stats), downscaled. + """Decode to float32 grayscale (and RGB for color stats), downscaled. Pillow, not cv2, and the source dimensions are recorded BEFORE the downscale. """ np = _numpy() from PIL import Image + from remove_ai_watermarks import image_io + try: + image_io._register_heif() # pyright: ignore[reportPrivateUsage] with Image.open(image_path) as img: info: dict[str, Any] = {"width": img.width, "height": img.height} if max(img.size) > MAX_SIDE: @@ -136,7 +177,9 @@ def read_gray(image_path: Path) -> tuple[Any, Any, dict[str, Any]]: gray = np.asarray(img.convert("L"), dtype=np.float32) except Exception as exc: logger.debug("pixel decode failed for %s: %s", image_path, exc) - return None, None, {"error": f"{type(exc).__name__}: {exc}"} + # Exception text from Pillow commonly embeds the absolute source path. + # Keep that detail in the log, not in the pathless transport contract. + return None, None, {"error": type(exc).__name__} return gray, rgb, info @@ -150,11 +193,13 @@ def dct_features(gray: Any) -> dict[str, Any]: basis = _dct_matrix(np) bins = np.linspace(-20.5, 20.5, 22) blocks = gray[:h8, :w8].reshape(h8 // 8, 8, w8 // 8, 8).swapaxes(1, 2) - coeff = np.einsum("ij,abjk,lk->abil", basis, blocks, basis) + rows = basis[[row for row, _ in AC_POSITIONS]] + columns = basis[[column for _, column in AC_POSITIONS]] + coeff = np.einsum("ki,abij,kj->abk", rows, blocks, columns) hists = [] lead_vals: list[Any] = [] - for dy, dx in AC_POSITIONS: - values = coeff[:, :, dy, dx].ravel() + for index in range(len(AC_POSITIONS)): + values = coeff[:, :, index].ravel() hists.append(np.histogram(values, bins=bins)[0].tolist()) lead_vals.append(np.abs(values)) out: dict[str, Any] = {"dct_ac_hist": hists} @@ -290,8 +335,8 @@ def perceptual_hash(gray: Any) -> str: small = np.asarray(Image.fromarray(gray.astype(np.float32), mode="F").resize((32, 32), Image.Resampling.LANCZOS)) basis = _dct_matrix(np, 32) - coeff = basis @ small @ basis.T - low = coeff[:8, :8].ravel()[1:] # drop DC + low_basis = basis[:8] + low = (low_basis @ small @ low_basis.T).ravel()[1:] # drop DC bits = low > np.median(low) return f"{int(''.join('1' if bit else '0' for bit in bits), 2):016x}" @@ -343,7 +388,7 @@ def spatial_artifacts(gray: Any, rgb: Any, *, ela: Any, residual: Any, phase: An return out -def extract_pixel_evidence(image_path: Path, *, artifacts: bool = False) -> PixelEvidence: +def extract_pixel_evidence(image_path: Path, *, artifacts: bool = False, timings: bool = False) -> PixelEvidence: """Measure every pixel-statistic family for one image in a single decode. The image is decoded ONCE and the intermediate maps (high-pass residual, ELA @@ -361,40 +406,79 @@ def extract_pixel_evidence(image_path: Path, *, artifacts: bool = False) -> Pixe artifacts: Also return the spatial layer -- perceptual hash, thumbnail and coarse maps. Off by default: those identify the source image, so asking for them is a decision the caller makes explicitly. + timings: Measure each stage and include rounded milliseconds in + :attr:`PixelEvidence.timing_ms`. Returns: A :class:`PixelEvidence`. """ - gray, rgb, info = read_gray(image_path) - if gray is None or rgb is None: - return PixelEvidence(path=image_path, decode=info) + started = time.perf_counter() + stage_started = started + measured: dict[str, float] = {} - residual = noise_residual_map(gray) - spectrum = fft_decompose(gray) - error = ela_map(rgb) + gray, rgb, info = read_gray(image_path) + measured["decode"] = time.perf_counter() - stage_started + if gray is None or rgb is None: + measured["total"] = time.perf_counter() - started + timing_ms = {name: round(seconds * 1000, 1) for name, seconds in measured.items()} if timings else {} + return PixelEvidence(path=image_path, decode=info, timing_ms=timing_ms) families: dict[str, dict[str, Any]] = {} + + residual = None + stage_started = time.perf_counter() + try: + residual = noise_residual_map(gray) + families["noise"] = noise_features(residual) if residual is not None else {} + except Exception as exc: + logger.debug("pixel family noise failed for %s: %s", image_path, exc) + families["noise"] = {"error": type(exc).__name__} + measured["noise"] = time.perf_counter() - stage_started + + spectrum = None + stage_started = time.perf_counter() + try: + spectrum = fft_decompose(gray) + families["fft"] = fft_features(spectrum[0]) if spectrum is not None else {} + except Exception as exc: + logger.debug("pixel family fft failed for %s: %s", image_path, exc) + families["fft"] = {"error": type(exc).__name__} + measured["fft"] = time.perf_counter() - stage_started + + error = None + stage_started = time.perf_counter() + try: + error = ela_map(rgb) + families["ela"] = ela_features(error) if error is not None else {} + except Exception as exc: + logger.debug("pixel family ela failed for %s: %s", image_path, exc) + families["ela"] = {"error": type(exc).__name__} + measured["ela"] = time.perf_counter() - stage_started + for name, compute in ( - ("noise", lambda: noise_features(residual) if residual is not None else {}), - ("fft", lambda: fft_features(spectrum[0]) if spectrum is not None else {}), - ("ela", lambda: ela_features(error) if error is not None else {}), ("dct", lambda: dct_features(gray)), ("gradient", lambda: gradient_features(gray)), ("color", lambda: color_features(rgb)), ): + stage_started = time.perf_counter() try: families[name] = compute() except Exception as exc: # one bad family must not lose the other five logger.debug("pixel family %s failed for %s: %s", name, image_path, exc) - families[name] = {"error": f"{type(exc).__name__}: {exc}"} + families[name] = {"error": type(exc).__name__} + measured[name] = time.perf_counter() - stage_started if artifacts: + stage_started = time.perf_counter() try: families["artifacts"] = spatial_artifacts( gray, rgb, ela=error, residual=residual, phase=spectrum[1] if spectrum is not None else None ) except Exception as exc: logger.debug("pixel artifacts failed for %s: %s", image_path, exc) - families["artifacts"] = {"error": f"{type(exc).__name__}: {exc}"} + families["artifacts"] = {"error": type(exc).__name__} + measured["full_artifacts"] = time.perf_counter() - stage_started - return PixelEvidence(path=image_path, decode=info, **families) + measured["total"] = time.perf_counter() - started + timing_ms = {name: round(seconds * 1000, 1) for name, seconds in measured.items()} if timings else {} + return PixelEvidence(path=image_path, decode=info, timing_ms=timing_ms, **families) diff --git a/tests/test_forensic_metadata.py b/tests/test_forensic_metadata.py new file mode 100644 index 0000000..cf6335e --- /dev/null +++ b/tests/test_forensic_metadata.py @@ -0,0 +1,236 @@ +"""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'{"Label":"1"}' + 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"metadata" + 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) == [] diff --git a/tests/test_identify.py b/tests/test_identify.py index a277b29..330a9e5 100644 --- a/tests/test_identify.py +++ b/tests/test_identify.py @@ -112,6 +112,36 @@ class TestProvenanceEvidence: assert evidence.exif_generator == "NovelAI" + @pytest.mark.parametrize( + "record", + [ + {"name": "trainedAlgorithmicMedia.jpg"}, + {"sha256": "jumb-c2pa-OpenAI-trainedAlgorithmicMedia"}, + {"pil": {"trainedAlgorithmicMedia": "plain"}}, + {"signals": {"provenance": {"is_ai_generated": True}}}, + {"pixel": {"error": "trainedAlgorithmicMedia"}}, + ], + ) + def test_external_diagnostics_and_arbitrary_keys_are_not_evidence(self, tmp_path: Path, record: dict): + report = identify_from_evidence(evidence_from_metadata_record(record, path=tmp_path / "plain.jpg")) + + assert report.is_ai_generated is None + assert report.signals == [] + + def test_external_metadata_value_is_evidence(self, tmp_path: Path): + record = { + "exif": { + "0th": { + "ImageDescription": "digitalSourceType=trainedAlgorithmicMedia", + } + } + } + + report = identify_from_evidence(evidence_from_metadata_record(record, path=tmp_path / "generated.jpg")) + + assert report.is_ai_generated is True + assert report.ai_source_kind == "generated" + @pytest.mark.parametrize( "filename", [ @@ -443,6 +473,18 @@ class TestIdentifyRealSamples: r = identify(p, check_visible=False, check_invisible=False) assert r.is_ai_generated is True assert r.platform == "Apple Photos (Clean Up AI edit)" + assert r.ai_source_kind == "enhanced" + + def test_standalone_iptc_composite_synthetic_is_enhanced(self, tmp_path: Path): + p = tmp_path / "composite.jpg" + p.write_bytes( + b'\xff\xd8\xff\xe1\xff\xd9' + ) + + r = identify(p, check_visible=False, check_invisible=False) + + assert r.is_ai_generated is True + assert r.ai_source_kind == "enhanced" def test_flux_bfl_c2pa_png(self): # flux-1.png: real Black Forest Labs FLUX.2 Playground output (signed C2PA). diff --git a/tests/test_metadata_record.py b/tests/test_metadata_record.py index ee6c0a7..f4b7315 100644 --- a/tests/test_metadata_record.py +++ b/tests/test_metadata_record.py @@ -24,7 +24,12 @@ from remove_ai_watermarks.identify import ( identify_from_evidence, identify_metadata_record, ) -from remove_ai_watermarks.metadata_record import HEAD_WINDOW, collect_metadata_record +from remove_ai_watermarks.metadata_record import ( + HEAD_WINDOW, + METADATA_RECORD_SCHEMA_VERSION, + METADATA_RECORD_TYPE, + collect_metadata_record, +) FIXTURES = Path(__file__).resolve().parent.parent / "data" / "fixtures" / "provenance" COMPARED = ("is_ai_generated", "platform", "confidence", "ai_source_kind", "ai_from_metadata") @@ -138,7 +143,25 @@ 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" + assert json.loads(json.dumps(record, allow_nan=False))["container"] == "png" + assert record["schema_version"] == METADATA_RECORD_SCHEMA_VERSION == 1 + assert record["record_type"] == METADATA_RECORD_TYPE == "provenance_metadata" + + @pytest.mark.parametrize("keep_version", [True, False]) + def test_transport_filename_is_not_evidence(self, tmp_path: Path, keep_version: bool): + """The path labels a record; detector tokens in it do not describe pixels.""" + path = tmp_path / "jumb-c2pa-OpenAI-trainedAlgorithmicMedia.jpg" + Image.fromarray(np.zeros((64, 64, 3), dtype=np.uint8)).save(path, "JPEG") + record = collect_metadata_record(path) + if not keep_version: + record.pop("schema_version") + record.pop("record_type") + + report = identify_metadata_record(record, path=path) + + assert report.is_ai_generated is None + assert report.platform is None + assert report.confidence == "none" def test_pixels_are_not_carried(self, tmp_path: Path): """The reason the record walks regions instead of shipping the head: a record @@ -168,6 +191,60 @@ class TestRecordShape: assert record["container"] == "unknown" assert record["metadata_base64"] == "" + assert record["status"] == "error" + assert record["issues"] == [{"stage": "source", "code": "unavailable"}] + + def test_unknown_record_schema_is_rejected(self, tmp_path: Path): + path = _noise_png(tmp_path / "plain.png") + record = collect_metadata_record(path) + record["schema_version"] = 2 + + with pytest.raises(ValueError, match="Unsupported provenance metadata schema"): + identify_metadata_record(record, path=path) + + @pytest.mark.parametrize("schema_version", [True, 1.0, "1", None]) + def test_native_record_schema_requires_the_integer_one(self, tmp_path: Path, schema_version: object): + path = _noise_png(tmp_path / "plain.png") + record = collect_metadata_record(path) + record["schema_version"] = schema_version + + with pytest.raises(ValueError, match="Unsupported provenance metadata schema"): + identify_metadata_record(record, path=path) + + @pytest.mark.parametrize("status", [None, "partial", "unknown", True]) + def test_native_record_requires_complete_collection_status(self, tmp_path: Path, status: object): + path = _noise_png(tmp_path / "plain.png") + record = collect_metadata_record(path) + if status is None: + record.pop("status") + else: + record["status"] = status + + with pytest.raises(ValueError, match="collection status"): + identify_metadata_record(record, path=path) + + @pytest.mark.parametrize("schema_version", [2, True, 1.0]) + def test_collection_rejects_unsupported_output_schema_before_reading(self, tmp_path: Path, schema_version: object): + with pytest.raises(ValueError, match="Unsupported provenance metadata schema"): + collect_metadata_record( + tmp_path / "missing.png", + schema_version=schema_version, # type: ignore[arg-type] + ) + + def test_broad_forensic_record_is_not_detector_input(self, tmp_path: Path): + path = _noise_png(tmp_path / "plain.png") + + with pytest.raises(ValueError, match="Unsupported metadata record type"): + identify_metadata_record( + {"record_type": "forensic_metadata", "schema_version": 1}, + path=path, + ) + + def test_failed_collection_cannot_be_judged_as_an_unknown_image(self, tmp_path: Path): + path = tmp_path / "missing.png" + + with pytest.raises(ValueError, match="collection failed"): + identify_metadata_record(collect_metadata_record(path), path=path) class TestReportTransport: @@ -178,7 +255,15 @@ class TestReportTransport: assert payload["schema_version"] == PROVENANCE_REPORT_SCHEMA_VERSION == 1 assert "path" not in payload - assert json.loads(json.dumps(payload)) == payload + assert json.loads(json.dumps(payload, allow_nan=False)) == payload + + @pytest.mark.parametrize("schema_version", [2, True, 1.0]) + def test_report_rejects_unsupported_output_schema(self, tmp_path: Path, schema_version: object): + path = _noise_png(tmp_path / "plain.png") + report = identify_metadata_record(collect_metadata_record(path), path=path) + + with pytest.raises(ValueError, match="Unsupported provenance report schema"): + report.to_dict(schema_version=schema_version) # type: ignore[arg-type] def test_convenience_entry_point_matches_explicit_sequence(self, tmp_path: Path): path = _noise_png(tmp_path / "plain.png") @@ -199,6 +284,38 @@ def test_a_webp_record_matches(tmp_path: Path): _assert_same_verdict(path) +def test_a_webp_record_collects_metadata_after_a_large_frame(tmp_path: Path): + """The RIFF walker seeks past coded pixels instead of stopping at the head window.""" + path = tmp_path / "late.webp" + rng = np.random.default_rng(7) + pixels = rng.integers(0, 255, (900, 900, 3), dtype=np.uint8) + xmp = b"trainedAlgorithmicMedia" + Image.fromarray(pixels).save(path, "WEBP", lossless=True, xmp=xmp) + + assert path.stat().st_size > HEAD_WINDOW + _assert_same_verdict(path) + + +def test_a_webp_record_does_not_carry_animation_frame_pixels(tmp_path: Path): + """ANMF is a coded-frame container, not a metadata chunk.""" + from remove_ai_watermarks.metadata_record import _riff_regions + + frame = b"jumb c2pa OpenAI trainedAlgorithmicMedia" * 20 + riff = b"RIFF" + (4 + 8 + len(frame)).to_bytes(4, "little") + b"WEBP" + riff += b"ANMF" + len(frame).to_bytes(4, "little") + frame + + assert frame not in _riff_regions(riff) + + +def test_an_invalid_short_riff_size_does_not_turn_the_container_into_a_trailer(tmp_path: Path): + from remove_ai_watermarks.metadata_record import _trailer + + path = tmp_path / "invalid.webp" + path.write_bytes(b"RIFF\x00\x00\x00\x00WEBPjumb c2pa trainedAlgorithmicMedia") + + assert _trailer(path, "webp") == b"" + + 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") diff --git a/tests/test_pixel_evidence.py b/tests/test_pixel_evidence.py index eb387f1..5ff96c0 100644 --- a/tests/test_pixel_evidence.py +++ b/tests/test_pixel_evidence.py @@ -1,10 +1,9 @@ -"""Tests for the experimental pixel-forensics collector. +"""Tests for the pixel-forensics collector. -It has no consumer, so there is no downstream behavior to pin. What these guard is -the part a future consumer would rely on and could not discover from the code: that -a family is empty rather than wrong when the image is too small for it, that one -failing family does not lose the other five, and that the artifacts which identify -the source image stay behind their opt-in. +These pin the service contract and the edge cases a consumer cannot infer safely: +a family is empty rather than wrong when the image is too small for it, one failing +family does not lose the other five, and artifacts that identify the source image +stay behind their opt-in. """ from __future__ import annotations @@ -16,7 +15,16 @@ import numpy as np import pytest from PIL import Image -from remove_ai_watermarks.pixel_evidence import PixelEvidence, extract_pixel_evidence, is_available +from remove_ai_watermarks.pixel_evidence import ( + AC_POSITIONS, + PIXEL_EVIDENCE_SCHEMA_VERSION, + PixelEvidence, + _dct_matrix, + dct_features, + extract_pixel_evidence, + is_available, + perceptual_hash, +) if TYPE_CHECKING: from pathlib import Path @@ -25,7 +33,7 @@ FAMILIES = ("dct", "fft", "noise", "ela", "gradient", "color") def _textured(path: Path, size: tuple[int, int] = (192, 160), *, seed: int = 0) -> Path: - """A textured image. Flat colour would make several families degenerate (zero + """A textured image. Flat color would make several families degenerate (zero residual, empty gradient histogram) and hide a real break.""" rng = np.random.default_rng(seed) base = rng.integers(0, 255, (size[1], size[0], 3), dtype=np.uint8) @@ -61,6 +69,28 @@ class TestFamilies: assert extract_pixel_evidence(path).decode == {"width": 3000, "height": 80} + def test_selected_dct_coefficients_match_the_full_transform(self): + gray = np.random.default_rng(7).uniform(0, 255, (24, 32)).astype(np.float32) + basis = _dct_matrix(np) + blocks = gray.reshape(3, 8, 4, 8).swapaxes(1, 2) + full = np.einsum("ij,abjk,lk->abil", basis, blocks, basis) + bins = np.linspace(-20.5, 20.5, 22) + expected = [ + np.histogram(full[:, :, row, column].ravel(), bins=bins)[0].tolist() for row, column in AC_POSITIONS + ] + + assert dct_features(gray)["dct_ac_hist"] == expected + + def test_perceptual_hash_matches_the_full_transform(self): + gray = np.random.default_rng(8).uniform(0, 255, (32, 32)).astype(np.float32) + basis = _dct_matrix(np, 32) + coefficients = basis @ gray @ basis.T + low = coefficients[:8, :8].ravel()[1:] + bits = low > np.median(low) + expected = f"{int(''.join('1' if bit else '0' for bit in bits), 2):016x}" + + assert perceptual_hash(gray) == expected + class TestDegenerateInputs: def test_undecodable_file_reports_the_error_and_stays_empty(self, tmp_path: Path): @@ -90,12 +120,14 @@ class TestDegenerateInputs: path = _textured(tmp_path / "textured.png") def boom(*args, **kwargs): - raise ValueError("family exploded") + raise ValueError(f"family failed for {path}") monkeypatch.setattr("remove_ai_watermarks.pixel_evidence.color_features", boom) evidence = extract_pixel_evidence(path) - assert "error" in evidence.color + assert evidence.color == {"error": "ValueError"} + assert evidence.status == "partial" + assert evidence.to_dict()["status"] == "partial" assert evidence.dct != {} assert evidence.gradient != {} @@ -154,3 +186,68 @@ def test_evidence_is_frozen(tmp_path: Path): assert isinstance(evidence, PixelEvidence) with pytest.raises(dataclasses.FrozenInstanceError): evidence.decode = {} # type: ignore[misc] + + +def test_transport_contract_is_versioned_json_and_omits_path(tmp_path: Path): + import json + + evidence = extract_pixel_evidence(_textured(tmp_path / "textured.png"), timings=True) + payload = evidence.to_dict() + + assert payload["schema_version"] == PIXEL_EVIDENCE_SCHEMA_VERSION == 1 + assert payload["status"] == "complete" + assert "path" not in payload + assert payload["timing_ms"]["total"] >= 0 + assert json.loads(json.dumps(payload, allow_nan=False)) == payload + + +@pytest.mark.parametrize("schema_version", [2, True, 1.0]) +def test_transport_rejects_unsupported_output_schema(tmp_path: Path, schema_version: object): + evidence = extract_pixel_evidence(_textured(tmp_path / "textured.png")) + + with pytest.raises(ValueError, match="Unsupported pixel evidence schema"): + evidence.to_dict(schema_version=schema_version) # type: ignore[arg-type] + + +def test_decode_error_transport_does_not_leak_the_local_path(tmp_path: Path): + import json + + path = tmp_path / "broken.png" + path.write_bytes(b"not an image") + + payload = extract_pixel_evidence(path).to_dict() + + assert payload["status"] == "error" + assert payload["decode"]["error"] == "UnidentifiedImageError" + assert str(path) not in json.dumps(payload) + + +def test_timings_are_opt_in(tmp_path: Path): + path = _textured(tmp_path / "textured.png") + + assert extract_pixel_evidence(path).timing_ms == {} + assert set(extract_pixel_evidence(path, artifacts=True, timings=True).timing_ms) == { + "decode", + "noise", + "fft", + "ela", + "dct", + "gradient", + "color", + "full_artifacts", + "total", + } + + +def test_pixel_decode_registers_optional_heif_opener(tmp_path: Path, monkeypatch): + registered = False + + def mark_registered(): + nonlocal registered + registered = True + + monkeypatch.setattr("remove_ai_watermarks.image_io._register_heif", mark_registered) + + extract_pixel_evidence(_textured(tmp_path / "textured.png")) + + assert registered is True