watermarks-remover/skills/remove-ai-marks/scripts/image_meta.py

"""Detect and strip C2PA / AI-related metadata from PNG and JPEG (stdlib)."""

from __future__ import annotations

import json
import os
import re
import struct
import subprocess
import sys
import zlib
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

from common import classify_finding_confidence, safe_arg, safe_write_bytes, subprocess_preexec_fn, which

SCRIPTS_DIR = Path(__file__).resolve().parent

PNG_SIG = b"\x89PNG\r\n\x1a\n"
JPEG_SOI = b"\xff\xd8"

# Chunk types / content patterns associated with C2PA / AI provenance
C2PA_MARKERS = (
    b"c2pa",
    b"C2PA",
    b"jumb",
    b"JUMB",
    b"c2ma",
    b"contentcredentials",
    b"contentauth",
    b"cai:",
    b"http://ns.adobe.com/xmp/InstanceID/",  # not solely C2PA
)

AI_META_HINTS = (
    b"c2pa",
    b"C2PA",
    b"contentcredentials",
    b"ContentCredentials",
    b"digitalSourceType",
    b"trainedAlgorithmicMedia",
    b"compositeWithTrainedAlgorithmicMedia",
    b"algorithmicMedia",
    b"AIGC",
    b"aigc",
    b"AI generated",
    b"Generated by",
    b"Claude",
    b"Anthropic",
    b"OpenAI",
    b"SynthID",
    b"dcterms:provenance",
)


@dataclass
class ImageInspectReport:
    path: str
    format: str  # png | jpeg | unknown
    has_c2pa: bool
    has_ai_metadata: bool
    findings: list[str] = field(default_factory=list)
    tools: dict[str, Any] = field(default_factory=dict)
    synthid: dict[str, Any] | None = None
    notes: list[str] = field(default_factory=list)

    def to_dict(self) -> dict:
        return {
            "path": self.path,
            "format": self.format,
            "has_c2pa": self.has_c2pa,
            "has_ai_metadata": self.has_ai_metadata,
            "findings": self.findings,
            "findings_confidence": [
                classify_finding_confidence(f) for f in self.findings
            ],
            "tools": self.tools,
            "synthid": self.synthid,
            "notes": self.notes,
        }


def detect_format(data: bytes) -> str:
    if data.startswith(PNG_SIG):
        return "png"
    if data.startswith(JPEG_SOI):
        return "jpeg"
    return "unknown"


def _contains_any(blob: bytes, needles: tuple[bytes, ...]) -> list[str]:
    found = []
    lower = blob.lower()
    for n in needles:
        if n.lower() in lower:
            try:
                found.append(n.decode("ascii", errors="replace"))
            except Exception:
                found.append(repr(n))
    return found


def inspect_png(data: bytes) -> tuple[bool, bool, list[str]]:
    findings: list[str] = []
    has_c2pa = False
    has_ai = False
    if not data.startswith(PNG_SIG):
        return False, False, ["not a PNG"]
    pos = 8
    while pos + 8 <= len(data):
        length = struct.unpack(">I", data[pos : pos + 4])[0]
        ctype = data[pos + 4 : pos + 8]
        chunk_start = pos + 8
        chunk_end = chunk_start + length
        if chunk_end + 4 > len(data):
            findings.append(f"truncated chunk {ctype!r}")
            break
        payload = data[chunk_start:chunk_end]
        name = ctype.decode("latin-1", errors="replace")
        # Private/ancillary chunks sometimes used for JUMBF/C2PA
        if ctype in (b"caBX", b"juMB", b"jumb") or ctype.startswith(b"c2"):
            has_c2pa = True
            findings.append(f"PNG chunk {name} (possible C2PA container)")
        if ctype in (b"tEXt", b"zTXt", b"iTXt", b"eXIf"):
            hits = _contains_any(payload, AI_META_HINTS + C2PA_MARKERS)
            if hits:
                has_ai = True
                if any(h.lower() in ("c2pa", "contentcredentials", "jumb") for h in hits):
                    has_c2pa = True
                findings.append(f"PNG {name}: {', '.join(hits[:8])}")
        if ctype == b"IEND":
            break
        pos = chunk_end + 4  # skip CRC
    # Whole-file scan fallback
    whole = _contains_any(data, C2PA_MARKERS)
    if whole and not has_c2pa:
        has_c2pa = True
        findings.append(f"byte-scan C2PA markers: {', '.join(whole[:6])}")
    return has_c2pa, has_ai or has_c2pa, findings


def inspect_jpeg(data: bytes) -> tuple[bool, bool, list[str]]:
    findings: list[str] = []
    has_c2pa = False
    has_ai = False
    if not data.startswith(JPEG_SOI):
        return False, False, ["not a JPEG"]
    i = 2
    n = len(data)
    while i + 4 <= n:
        if data[i] != 0xFF:
            i += 1
            continue
        # skip fill
        while i < n and data[i] == 0xFF:
            i += 1
        if i >= n:
            break
        marker = data[i]
        i += 1
        if marker in (0xD8, 0xD9):  # SOI/EOI
            continue
        if marker == 0xDA:  # SOS — image data follows
            break
        if marker >= 0xD0 and marker <= 0xD7:  # RSTn
            continue
        if i + 2 > n:
            break
        seglen = struct.unpack(">H", data[i : i + 2])[0]
        if seglen < 2 or i + seglen > n:
            findings.append(f"bad segment length at marker 0x{marker:02X}")
            break
        payload = data[i + 2 : i + seglen]
        i += seglen

        # APP11 (0xEB) often holds JUMBF/C2PA
        if marker == 0xEB:
            has_c2pa = True
            findings.append("JPEG APP11 segment (JUMBF/C2PA common)")
        if marker in (0xE1, 0xE2, 0xED, 0xEE, 0xEB):  # APP1,2,13,14,11
            hits = _contains_any(payload, AI_META_HINTS + C2PA_MARKERS)
            if hits:
                has_ai = True
                if any(
                    h.lower() in ("c2pa", "contentcredentials", "jumb", "contentauth")
                    for h in hits
                ):
                    has_c2pa = True
                findings.append(f"JPEG APP{marker - 0xE0}: {', '.join(hits[:8])}")

    whole = _contains_any(data, C2PA_MARKERS)
    if whole and not has_c2pa:
        has_c2pa = True
        findings.append(f"byte-scan C2PA markers: {', '.join(whole[:6])}")
    return has_c2pa, has_ai or has_c2pa, findings


def run_optional_tools(path: Path) -> dict[str, Any]:
    tools: dict[str, Any] = {}
    c2patool = which("c2patool")
    if c2patool:
        try:
            r = subprocess.run(
                [c2patool, safe_arg(str(path))],
                capture_output=True,
                text=True,
                timeout=30,
                preexec_fn=subprocess_preexec_fn,
            )
            out = (r.stdout or "") + (r.stderr or "")
            low = out.lower()
            # Negative markers must veto every positive branch, so the
            # positive alternatives are parenthesised: c2patool reports a
            # missing manifest as "Error: No claim found", which contains
            # the substring "claim" and would otherwise read as a hit.
            no_manifest = "no claim" in low or "no jumbf" in low
            tools["c2patool"] = {
                "available": True,
                "returncode": r.returncode,
                "snippet": out[:2000],
                "has_manifest": (
                    "claim" in low or "c2pa" in low or "manifest" in low
                )
                and not no_manifest,
            }
        except Exception as e:
            tools["c2patool"] = {"available": True, "error": str(e)}
    else:
        tools["c2patool"] = {"available": False}

    exiftool = which("exiftool")
    if exiftool:
        try:
            r = subprocess.run(
                [exiftool, "-G1", "-a", "-s", safe_arg(str(path))],
                capture_output=True,
                text=True,
                timeout=30,
                preexec_fn=subprocess_preexec_fn,
            )
            out = r.stdout or ""
            interesting = [
                line
                for line in out.splitlines()
                if re.search(
                    r"c2pa|content.?credential|AIGC|digitalSource|XMP|EXIF|IPTC|jumb",
                    line,
                    re.I,
                )
            ]
            tools["exiftool"] = {
                "available": True,
                "interesting_lines": interesting[:50],
            }
        except Exception as e:
            tools["exiftool"] = {"available": True, "error": str(e)}
    else:
        tools["exiftool"] = {"available": False}
    return tools


def run_synthid_score(
    path: Path,
    upstream_dir: str | None = None,
) -> dict[str, Any] | None:
    """Run the optional reverse-SynthID scorer in a subprocess.

    Returns None when the scorer is not configured or unavailable (exit 3),
    so callers can keep the default "no SynthID score" behavior.
    """
    if upstream_dir is None:
        upstream_dir = os.environ.get("REVERSE_SYNTHID_DIR")
    if not upstream_dir:
        return None

    script = SCRIPTS_DIR / "score_synthid.py"
    cmd = [
        sys.executable,
        str(script),
        str(path),
        "--upstream-dir",
        str(upstream_dir),
        "--json",
    ]
    try:
        r = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            timeout=180,
            preexec_fn=subprocess_preexec_fn,
        )
    except Exception as e:
        return {"available": False, "error": str(e)}

    if r.returncode == 3:
        return None
    if r.returncode != 0:
        return {"available": False, "error": (r.stderr or "").strip()[:2000]}
    try:
        return json.loads(r.stdout or "{}")
    except json.JSONDecodeError as e:
        return {"available": False, "error": f"bad scorer JSON: {e}"}


def inspect_image(
    path: Path,
    synthid_dir: str | None = None,
) -> ImageInspectReport:
    data = path.read_bytes()
    fmt = detect_format(data)
    if fmt == "png":
        has_c2pa, has_ai, findings = inspect_png(data)
    elif fmt == "jpeg":
        has_c2pa, has_ai, findings = inspect_jpeg(data)
    else:
        has_c2pa, has_ai, findings = False, False, ["unsupported format (MVP: PNG/JPEG)"]

    notes: list[str] = []
    if fmt == "unknown":
        notes.append("format not fully inspected; only PNG/JPEG are supported")

    tools = run_optional_tools(path)
    # Elevate flags from tools
    ct = tools.get("c2patool") or {}
    if ct.get("has_manifest"):
        has_c2pa = True
        findings.append("c2patool reports a C2PA-related manifest")

    return ImageInspectReport(
        path=str(path),
        format=fmt,
        has_c2pa=has_c2pa,
        has_ai_metadata=has_ai,
        findings=findings,
        tools=tools,
        synthid=run_synthid_score(path, synthid_dir),
        notes=notes,
    )


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


def strip_png(data: bytes, *, strip_all_text: bool = True) -> tuple[bytes, list[str]]:
    if not data.startswith(PNG_SIG):
        raise ValueError("not PNG")
    actions: list[str] = []
    out = bytearray(PNG_SIG)
    pos = 8
    while pos + 8 <= len(data):
        length = struct.unpack(">I", data[pos : pos + 4])[0]
        ctype = data[pos + 4 : pos + 8]
        chunk_start = pos + 8
        chunk_end = chunk_start + length
        if chunk_end + 4 > len(data):
            break
        payload = data[chunk_start:chunk_end]
        crc_bytes = data[chunk_end : chunk_end + 4]
        pos = chunk_end + 4
        name = ctype.decode("latin-1", errors="replace")

        drop = False
        if ctype in (b"eXIf", b"caBX") or ctype.startswith(b"c2"):
            drop = True
            actions.append(f"drop chunk {name}")
        elif ctype in (b"tEXt", b"zTXt", b"iTXt"):
            if strip_all_text or _contains_any(payload, AI_META_HINTS + C2PA_MARKERS):
                drop = True
                actions.append(f"drop chunk {name}")
        elif _contains_any(ctype + payload, C2PA_MARKERS) and ctype not in (
            b"IHDR",
            b"IDAT",
            b"IEND",
            b"PLTE",
            b"tRNS",
            b"gAMA",
            b"pHYs",
            b"sRGB",
            b"cHRM",
            b"iCCP",
        ):
            drop = True
            actions.append(f"drop chunk {name} (C2PA marker in payload)")

        if not drop:
            out.extend(struct.pack(">I", length) + ctype + payload + crc_bytes)
        if ctype == b"IEND":
            break
    if not actions:
        actions.append("no PNG metadata chunks removed (already clean or none matched)")
    return bytes(out), actions


def strip_jpeg(data: bytes, *, strip_all_app: bool = True) -> tuple[bytes, list[str]]:
    if not data.startswith(JPEG_SOI):
        raise ValueError("not JPEG")
    actions: list[str] = []
    out = bytearray(JPEG_SOI)
    i = 2
    n = len(data)
    while i < n:
        if data[i] != 0xFF:
            # unexpected; copy rest
            out.extend(data[i:])
            actions.append("copied remainder after non-marker byte")
            break
        while i < n and data[i] == 0xFF:
            i += 1
        if i >= n:
            break
        marker = data[i]
        i += 1

        if marker == 0xD9:  # EOI
            out.extend(b"\xff\xd9")
            break
        if marker == 0xD8:  # nested SOI?
            continue
        if 0xD0 <= marker <= 0xD7:
            out.extend(bytes([0xFF, marker]))
            continue
        if marker == 0xDA:  # SOS — copy through end
            # need length of SOS header then entropy-coded data to EOI
            if i + 2 > n:
                break
            seglen = struct.unpack(">H", data[i : i + 2])[0]
            # Find EOI from here carefully: after SOS segment header, scan for FF D9
            # not preceded by stuffed FF 00 issues — simple approach: copy from FF DA to end
            sos_start = i - 2  # points at 0xFF before marker... actually marker already consumed
            # Reconstruct: FF DA + rest of file
            out.extend(b"\xff\xda")
            out.extend(data[i:])
            actions.append("preserved entropy-coded scan (SOS→EOF)")
            break

        if i + 2 > n:
            break
        seglen = struct.unpack(">H", data[i : i + 2])[0]
        if seglen < 2 or i + seglen > n:
            out.extend(data[i - 2 :])  # best effort
            actions.append("truncated segment; copied remainder")
            break
        payload = data[i + 2 : i + seglen]
        next_i = i + seglen

        # APP0 is JFIF — keep for compatibility unless full strip of all APP
        keep = False
        drop = False
        if 0xE0 <= marker <= 0xEF:  # APPn
            if marker == 0xEB:  # APP11 JUMBF/C2PA
                drop = True
                actions.append("drop APP11 (C2PA/JUMBF)")
            elif strip_all_app and marker != 0xE0:
                # keep APP0 (JFIF) by default
                drop = True
                actions.append(f"drop APP{marker - 0xE0}")
            elif _contains_any(payload, AI_META_HINTS + C2PA_MARKERS):
                drop = True
                actions.append(f"drop APP{marker - 0xE0} (AI/C2PA markers)")
            else:
                keep = True
        elif marker in (0xFE,):  # COM
            drop = True
            actions.append("drop COM comment")
        else:
            keep = True

        if keep and not drop:
            out.extend(bytes([0xFF, marker]))
            out.extend(data[i : i + seglen])
        i = next_i

    if not actions:
        actions.append("no JPEG APP segments removed")
    return bytes(out), actions


def clean_image(
    path: Path,
    dest: Path,
    *,
    strip_all_metadata: bool = True,
    synthid_dir: str | None = None,
) -> dict[str, Any]:
    synthid_before = run_synthid_score(path, synthid_dir)
    data = path.read_bytes()
    fmt = detect_format(data)
    if fmt == "png":
        cleaned, actions = strip_png(data, strip_all_text=strip_all_metadata)
    elif fmt == "jpeg":
        cleaned, actions = strip_jpeg(data, strip_all_app=strip_all_metadata)
    else:
        raise ValueError(f"unsupported format: {fmt}")

    # Optional exiftool pass for residual tags
    safe_write_bytes(dest, cleaned)
    exiftool = which("exiftool")
    if exiftool and strip_all_metadata:
        try:
            subprocess.run(
                [
                    exiftool,
                    "-all=",
                    "-overwrite_original",
                    safe_arg(str(dest)),
                ],
                capture_output=True,
                text=True,
                timeout=60,
                check=False,
                preexec_fn=subprocess_preexec_fn,
            )
            actions.append("exiftool -all= pass")
        except Exception as e:
            actions.append(f"exiftool failed: {e}")

    after = inspect_image(dest, synthid_dir=synthid_dir)
    return {
        "input": str(path),
        "output": str(dest),
        "format": fmt,
        "actions": actions,
        "bytes_in": len(data),
        "bytes_out": dest.stat().st_size,
        "still_has_c2pa": after.has_c2pa,
        "still_has_ai_metadata": after.has_ai_metadata,
        "post_findings": after.findings,
        "synthid_before": synthid_before,
        "synthid_after": after.synthid,
    }