remove-ai-watermarks/src/remove_ai_watermarks/_internal/c2pa.py

"""C2PA inspection through the official reader with a bounded PNG fallback."""

from __future__ import annotations

import contextlib
import functools
import json
import logging
import re
import struct
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING, Any, cast

from remove_ai_watermarks._internal.constants import (
    C2PA_ACTIONS,
    C2PA_AI_TOOLS,
    C2PA_CHUNK_TYPE,
    C2PA_ISSUERS,
    C2PA_SIGNATURES,
    C2PA_SOFT_BINDINGS,
    PNG_SIGNATURE,
    SYNTHID_C2PA_ISSUERS,
)

logger = logging.getLogger(__name__)

if TYPE_CHECKING:
    from typing import BinaryIO

_C2paReader: Any = None
with contextlib.suppress(Exception):
    from c2pa import Reader as _C2paReader  # pyright: ignore[reportMissingTypeStubs]

_C2PA_READER_AVAILABLE = _C2paReader is not None
_PNG_HEADER = struct.Struct(">I4s")


@dataclass(frozen=True)
class _PngChunk:
    payload: bytes
    serialized: bytes


def reader_available() -> bool:
    """Return whether the official C2PA reader loaded successfully."""
    return _C2PA_READER_AVAILABLE


def _manifest_json_uncached(path: str) -> str | None:
    try:
        reader = _C2paReader.try_create(path)
    except Exception as error:
        logger.debug("C2PA reader rejected %s: %s", path, error)
        return None
    if reader is None:
        return None
    try:
        with reader:
            return cast("str", reader.json())
    except Exception as error:
        logger.debug("C2PA reader could not serialize %s: %s", path, error)
        return None


@functools.lru_cache(maxsize=8)
def _manifest_json_cached(path: str, _mtime_ns: int) -> str | None:
    return _manifest_json_uncached(path)


def read_manifest_store_json(image_path: Path) -> str | None:
    """Read the complete manifest-store JSON, caching it until the file changes."""
    if not reader_available():
        return None
    path = str(image_path)
    try:
        return _manifest_json_cached(path, image_path.stat().st_mtime_ns)
    except OSError:
        return _manifest_json_uncached(path)


def _find_c2pa_chunk(path: Path) -> _PngChunk | None:
    """Return the first recognizable C2PA chunk without loading the whole PNG."""
    try:
        stream = path.open("rb")
    except OSError:
        return None
    with stream:
        if stream.read(len(PNG_SIGNATURE)) != PNG_SIGNATURE:
            return None
        file_size = stream.seek(0, 2)
        stream.seek(len(PNG_SIGNATURE))
        while True:
            header = stream.read(_PNG_HEADER.size)
            if len(header) != _PNG_HEADER.size:
                return None
            length, kind = _PNG_HEADER.unpack(header)
            if length + 4 > file_size - stream.tell():
                return None
            if kind == C2PA_CHUNK_TYPE:
                payload = stream.read(length)
                crc = stream.read(4)
                if _looks_like_c2pa(payload):
                    return _PngChunk(payload, header + payload + crc)
            else:
                stream.seek(length + 4, 1)
            if kind == b"IEND":
                return None


def _is_well_formed_png(path: Path) -> bool:
    """Validate PNG chunk bounds with seeks rather than payload allocations."""
    try:
        stream = path.open("rb")
    except OSError:
        return False
    with stream:
        if stream.read(len(PNG_SIGNATURE)) != PNG_SIGNATURE:
            return False
        file_size = stream.seek(0, 2)
        stream.seek(len(PNG_SIGNATURE))
        while True:
            header = stream.read(_PNG_HEADER.size)
            if len(header) != _PNG_HEADER.size:
                return False
            length, kind = _PNG_HEADER.unpack(header)
            if length + 4 > file_size - stream.tell():
                return False
            stream.seek(length + 4, 1)
            if kind == b"IEND":
                return True


def _copy_bytes(source: BinaryIO, target: BinaryIO, byte_count: int) -> None:
    """Copy exactly one bounded chunk without allocating its complete payload."""
    remaining = byte_count
    while remaining:
        block = source.read(min(remaining, 1024 * 1024))
        if not block:
            raise OSError("PNG changed while it was being copied")
        target.write(block)
        remaining -= len(block)


def _looks_like_c2pa(payload: bytes) -> bool:
    lowered = payload.lower()
    return any(signature in payload for signature in C2PA_SIGNATURES) or b"c2pa" in lowered or b"jumb" in lowered


def extract_c2pa_chunk(image_path: Path) -> bytes | None:
    """Return the first complete C2PA PNG chunk, including header and CRC."""
    if image_path.suffix.casefold() != ".png":
        return None
    chunk = _find_c2pa_chunk(image_path)
    return None if chunk is None else chunk.serialized


def has_c2pa_metadata(image_path: Path) -> bool:
    """Return whether a validly bounded PNG contains a recognizable C2PA chunk."""
    return extract_c2pa_chunk(Path(image_path)) is not None


def _active_manifest(store: dict[str, Any]) -> dict[str, Any]:
    manifests = store.get("manifests")
    if not isinstance(manifests, dict):
        return {}
    typed_manifests = cast("dict[object, object]", manifests)
    active = typed_manifests.get(store.get("active_manifest"))
    return cast("dict[str, Any]", active) if isinstance(active, dict) else {}


def _claim_generator_from_store(store: dict[str, Any]) -> str | None:
    active = _active_manifest(store)
    direct = active.get("claim_generator")
    if isinstance(direct, str) and direct.isprintable() and direct:
        return direct
    candidates = active.get("claim_generator_info")
    if isinstance(candidates, list) and candidates and isinstance(candidates[0], dict):
        candidate = cast("dict[object, object]", candidates[0])
        name = candidate.get("name")
        if isinstance(name, str) and name.isprintable() and name:
            return name
    return None


def synthid_verdict(vendors: str) -> str:
    """Describe why metadata implies a likely pixel-level SynthID watermark."""
    return f"likely present ({vendors} embeds SynthID with C2PA)"


def _names_present(buffer: bytes, registry: dict[bytes, str]) -> list[str]:
    return sorted({label for token, label in registry.items() if token in buffer})


def synthid_vendors_in(buffer: bytes) -> list[str]:
    """List matching C2PA issuers known to pair their manifests with SynthID."""
    registry = {token: label for token, label in C2PA_ISSUERS.items() if token in SYNTHID_C2PA_ISSUERS}
    return _names_present(buffer, registry)


def soft_binding_vendors_in(buffer: bytes) -> list[str]:
    """List the soft-binding algorithms named in manifest bytes."""
    return _names_present(buffer, C2PA_SOFT_BINDINGS)


def _ordered_matches(buffer: bytes, registry: dict[bytes, str]) -> list[str]:
    return list(dict.fromkeys(label for token, label in registry.items() if token in buffer))


def _populate_registry_fields(buffer: bytes, info: dict[str, Any]) -> bool:
    issuers = _ordered_matches(buffer, C2PA_ISSUERS)
    tools = _ordered_matches(buffer, C2PA_AI_TOOLS)
    actions = _ordered_matches(buffer, C2PA_ACTIONS)
    if issuers:
        info["issuer"] = ", ".join(issuers)
    if tools:
        info["ai_tool"] = ", ".join(tools)
    if actions:
        info["actions"] = ", ".join(actions)

    ai_source = False
    if b"trainedAlgorithmicMedia" in buffer:
        info.update(source_type="trainedAlgorithmicMedia (AI-generated)", ai_source_kind="generated")
        ai_source = True
    elif b"compositeWithTrainedAlgorithmicMedia" in buffer:
        info.update(source_type="compositeWithTrainedAlgorithmicMedia (AI-enhanced)", ai_source_kind="enhanced")
        ai_source = True
    elif b"algorithmicMedia" in buffer:
        info["source_type"] = "algorithmicMedia"

    synthid = synthid_vendors_in(buffer)
    if ai_source and synthid:
        info["synthid_vendors"] = synthid
        info["synthid_watermark"] = synthid_verdict(", ".join(synthid))

    soft_bindings = soft_binding_vendors_in(buffer)
    if soft_bindings:
        info["soft_binding_vendors"] = soft_bindings
        info["soft_binding"] = ", ".join(soft_bindings)
    return ai_source


def _base_info(byte_count: int, *, fallback: bool = False) -> dict[str, Any]:
    container = "C2PA manifest" if fallback else "C2PA manifest store"
    return {
        "has_c2pa": True,
        "type": "C2PA (Coalition for Content Provenance and Authenticity)",
        "c2pa_manifest": f"{container} ({byte_count} bytes)",
    }


def _info_from_store(store: dict[str, Any], encoded: bytes) -> dict[str, Any]:
    info = _base_info(len(encoded))
    _populate_registry_fields(encoded, info)
    generator = _claim_generator_from_store(store)
    if generator is not None:
        info["claim_generator"] = generator
    signature_value = _active_manifest(store).get("signature_info")
    if isinstance(signature_value, dict):
        signature = cast("dict[object, object]", signature_value)
        timestamp = signature.get("time")
        if timestamp:
            info["timestamp"] = str(timestamp)
    return info


def c2pa_info_from_manifest_store(store: str | dict[str, Any]) -> dict[str, Any]:
    """Normalize a manifest store supplied as JSON text or a decoded object."""
    try:
        raw_decoded: object = store if isinstance(store, dict) else json.loads(store)
        decoded = cast("dict[str, Any]", raw_decoded) if isinstance(raw_decoded, dict) else None
        if not isinstance(decoded, dict) or not decoded or decoded.get("error"):
            return {}
        encoded = json.dumps(decoded, ensure_ascii=False).encode() if isinstance(store, dict) else store.encode()
    except (TypeError, ValueError):
        return {}
    return _info_from_store(decoded, encoded)


def cbor_text_after(payload: bytes, key: bytes) -> str | None:
    """Decode a definite-length CBOR text value immediately following ``key``."""
    key_end = payload.find(key)
    if key_end < 0:
        return None
    cursor = key_end + len(key)
    if cursor >= len(payload):
        return None
    initial = payload[cursor]
    if 0x60 <= initial <= 0x77:
        length, cursor = initial & 0x1F, cursor + 1
    elif initial == 0x78 and cursor + 1 < len(payload):
        length, cursor = payload[cursor + 1], cursor + 2
    elif initial == 0x79 and cursor + 2 < len(payload):
        length = int.from_bytes(payload[cursor + 1 : cursor + 3], "big")
        cursor += 3
    else:
        return None
    raw = payload[cursor : cursor + length]
    if len(raw) != length:
        return None
    try:
        return raw.decode()
    except UnicodeDecodeError:
        return raw.decode("latin1", errors="replace")


def _parse_c2pa_chunk(payload: bytes, info: dict[str, Any]) -> None:
    info.update(_base_info(len(payload), fallback=True))
    _populate_registry_fields(payload, info)
    for key, output_key in ((b"name", "claim_generator"), (b"specVersion", "c2pa_spec")):
        value = cbor_text_after(payload, key)
        if value and value.isprintable():
            info[output_key] = value
    timestamps = [item.decode() for item in re.findall(rb"\d{14}Z", payload)]
    if timestamps:
        info["timestamp"] = timestamps[0]
        if len(timestamps) > 1:
            info["timestamps"] = timestamps[:3]


def _extract_c2pa_info_png(image_path: Path) -> dict[str, Any]:
    if image_path.suffix.casefold() != ".png":
        return {}
    chunk = _find_c2pa_chunk(image_path)
    if chunk is None:
        return {}
    info: dict[str, Any] = {}
    _parse_c2pa_chunk(chunk.payload, info)
    return info


def _extract_c2pa_info_impl(image_path: Path) -> dict[str, Any]:
    store = read_manifest_store_json(Path(image_path))
    if store is not None:
        return c2pa_info_from_manifest_store(store)
    return _extract_c2pa_info_png(Path(image_path))


@functools.lru_cache(maxsize=4)
def _extract_c2pa_info_cached(path_str: str, _mtime_ns: int, _size: int, _reader: bool) -> dict[str, Any]:
    """Cache shim: every argument after the path is key-only.

    ``_reader`` is in the key because the answer genuinely depends on it -- with the
    official reader the manifest comes back as a store, without it from the hand-rolled
    PNG chunk parser, and the two produce different ``c2pa_manifest`` labels. Keying on
    the file alone handed a reader-path result to a caller that had disabled the reader.
    """
    return _extract_c2pa_info_impl(Path(path_str))


def extract_c2pa_info(image_path: Path) -> dict[str, Any]:
    """Return normalized C2PA evidence from the official reader or PNG fallback.

    Memoized on ``(path, mtime_ns, size)``: one ``identify`` reaches this twice (once
    directly, once inside ``get_ai_metadata``) and each call re-runs the Rust manifest
    reader and re-parses its JSON. Size joins mtime in the key because this package
    rewrites files in place, and an in-place rewrite can land inside one mtime tick.
    """
    try:
        stat = image_path.stat()
    except OSError:
        # No stat (a pipe, or a race): read uncached rather than fail.
        return _extract_c2pa_info_impl(image_path)
    cached = _extract_c2pa_info_cached(str(image_path), stat.st_mtime_ns, stat.st_size, _C2PA_READER_AVAILABLE)
    # Deep-ish copy: values are scalars plus a few lists, and a caller mutating one of
    # those lists would otherwise poison every later reader of the same file.
    return {key: list(cast("list[Any]", value)) if isinstance(value, list) else value for key, value in cached.items()}


def inject_c2pa_chunk(target_path: Path, output_path: Path, c2pa_chunk: bytes) -> None:
    """Replace any C2PA chunks in a PNG and insert ``c2pa_chunk`` before IDAT."""
    if target_path.suffix.casefold() != ".png" or output_path.suffix.casefold() != ".png":
        raise ValueError("C2PA chunk injection is only supported for PNG files")
    if not _is_well_formed_png(target_path):
        raise ValueError("Target is not a well-formed PNG file")

    output_path.parent.mkdir(parents=True, exist_ok=True)
    with target_path.open("rb") as source, output_path.open("wb") as target:
        target.write(source.read(len(PNG_SIGNATURE)))
        inserted = False
        while True:
            header = source.read(_PNG_HEADER.size)
            length, kind = _PNG_HEADER.unpack(header)
            if kind == b"IDAT" and not inserted:
                target.write(c2pa_chunk)
                inserted = True
            if kind == C2PA_CHUNK_TYPE:
                source.seek(length + 4, 1)
            else:
                target.write(header)
                _copy_bytes(source, target, length + 4)
            if kind == b"IEND":
                break