provenance_gate.py

"""Self-contained provenance (C2PA/metadata) gate — works on raw upload bytes.

Ported from the detector repo's scripts/provenance_check.py. No external service
(:8061) and no special c2pa library: pure PIL + byte-signature scan. When a
generator/provenance signal is present it is HIGH precision -> force AI verdict.
Absence proves nothing (screenshots strip it) -> defer to the pixel model.
"""
from __future__ import annotations

import io
import json

from PIL import ExifTags, Image

SIGNATURES = {
    "C2PA/JUMBF": [b"c2pa", b"jumbf", b"urn:c2pa", b"claim_generator", b"caBX", b"contentauth", b"c2pa.assertions"],
    "OpenAI": [b"openai", b"dall-e", b"dall\xc2\xb7e", b"dalle", b"gpt-image", b"chatgpt"],
    "Adobe/Firefly": [b"firefly", b"cai:"],
    "StableDiffusion/Comfy": [b"stable diffusion", b"stable-diffusion", b"sd-webui", b"automatic1111",
                               b"comfyui", b"negative prompt"],
    "Midjourney": [b"midjourney", b"niji"],
    "Google/SynthID": [b"synthid", b"made with google", b"gemini", b"google ai", b"imagen 3", b"imagen 4"],
    "Other-gen": [b"novelai", b"leonardo.ai", b"ideogram", b"playground", b"seedream", b"flux.1", b"grok-image"],
    "AIGC-generic": [b"trainedalgorithmicmedia", b"compositesynthetic", b"ai-generated", b"genai"],
}

_EXIF_TAGS = {v: k for k, v in ExifTags.TAGS.items()}


def _scan(raw: bytes) -> set[str]:
    low = raw.lower()
    return {label for label, sigs in SIGNATURES.items() if any(s in low for s in sigs)}


def check_bytes(raw: bytes) -> dict:
    """Returns {'signals': [...], 'details': {...}} ; signals non-empty => AI."""
    details: dict = {}
    sig = _scan(raw)
    try:
        img = Image.open(io.BytesIO(raw))
        details["format"] = img.format
        info = {k: v for k, v in (img.info or {}).items()
                if isinstance(v, (str, bytes)) and k.lower() != "icc_profile"}
        if info:
            details["text_chunks"] = {
                k: (v if isinstance(v, str) else v.decode("latin-1", "ignore"))[:300]
                for k, v in info.items()
            }
        if "xmp" in img.info:
            xmp = img.info["xmp"]
            sig |= _scan(xmp if isinstance(xmp, bytes) else xmp.encode("latin-1", "ignore"))
        exif = img.getexif()
        if exif:
            wanted = ("Make", "Model", "Software", "Artist", "XPComment", "ImageDescription")
            ex = {name: str(exif[_EXIF_TAGS[name]])[:200]
                  for name in wanted if name in _EXIF_TAGS and _EXIF_TAGS[name] in exif}
            if ex:
                details["exif"] = ex
    except Exception as exc:  # noqa: BLE001
        details["pil_error"] = str(exc)

    blob = json.dumps(details, ensure_ascii=False).lower()
    for label, sigs in SIGNATURES.items():
        if any(s.decode("latin-1", "ignore") in blob for s in sigs):
            sig.add(label)

    return {"signals": sorted(sig), "details": details}


def check_filename(name: str | None) -> list[str]:
    """Weak mobile fallback for generator-exported filenames.

    Some phone gallery providers hand Expo a transcoded/cache copy that strips
    C2PA/JUMBF chunks. The filename is not cryptographic provenance, but app
    exports like "ChatGPT Image ..." are still useful as a high-confidence hint
    for the demo path when raw provenance was removed before upload.
    """
    low = (name or "").lower()
    found: list[str] = []
    if "chatgpt image" in low or "gpt-image" in low or "dall-e" in low or "dalle" in low:
        found.extend(["OpenAI", "Filename:AI-export"])
    if "midjourney" in low or "niji" in low:
        found.extend(["Midjourney", "Filename:AI-export"])
    if "stable diffusion" in low or "comfyui" in low or "sd-webui" in low:
        found.extend(["StableDiffusion/Comfy", "Filename:AI-export"])
    # Gemini app saves as "Gemini_Generated_Image_xxxxx.png"; Imagen exports similarly.
    if "gemini" in low or "imagen" in low or "google_ai" in low or "made with google" in low:
        found.extend(["Google/SynthID", "Filename:AI-export"])
    return sorted(set(found))