#!/usr/bin/env python3
"""Varsel AI-transparency reference verifier (standalone, no network).

Checks whether a Varsel document carries the machine-readable AI-content marking
described in the published specification (varsel-ai-transparency/1.0), prints the
manifest it carries, and verifies the signature against the Varsel trust anchor
you pass in.

    python verify_varsel_ai_transparency.py <report.pdf> <anchor.pem>
    python verify_varsel_ai_transparency.py <report.html> <anchor.pem>
    python verify_varsel_ai_transparency.py --json <preview.json> <anchor.pem>

Exit codes:
    0  marked, every check passed, and the signature verified against the anchor
    4  marked and internally consistent, but NO anchor was given, so the signer is
       unverified (a forged file can reach this state; never treat it as verified)
    3  unmarked (no carrier, or a carrier without a signature)
    1  marked, but at least one check failed (the failed checks are named)
    2  usage or environment error (bad arguments, unreadable anchor file, a
       required library missing, or an internal error while reading the file)

Dependencies (all from PyPI, no Varsel code):
    cryptography                      JSON and HTML modes
    pikepdf, pyHanko (pyhanko)        PDF mode in addition

What exit 0 means: the file carries a manifest stating that the document contains
AI-generated content and listing the AI providers recorded for the run, and the
signature over it was made with the private key matching the anchor certificate
you supplied. It says nothing about the accuracy of the screening; the decision
and the accountability for it stay with the reader of the report.

Published by Conshultocracy AB (org. nr 559419-4697). Specification and trust anchor:
https://www.conshultocracy.com/tools/varsel/ai-transparency
"""

import hashlib
import io
import json
import re
import sys
import xml.etree.ElementTree as ET
from datetime import datetime, timezone
from pathlib import Path

# --- constants, identical to scripts/ai_transparency.py in the Varsel code base ---
SPEC = "varsel-ai-transparency/1.0"
DIGITAL_SOURCE_TYPE = "http://cv.iptc.org/newscodes/digitalsourcetype/trainedAlgorithmicMedia"
XMP_NAMESPACE = "https://varsel.nu/ns/ai-transparency/1.0/"
IPTC_NAMESPACE = "http://iptc.org/std/Iptc4xmpExt/2008-02-29/"
RDF_NAMESPACE = "http://www.w3.org/1999/02/22-rdf-syntax-ns#"
XMP_MANIFEST_PROPERTY = "aiManifest"
XMP_SHA256_PROPERTY = "aiManifestSha256"
SIGNATURE_FIELD_NAME = "VarselAITransparency"
CARRIER_KEYS = ("manifest", "manifest_sha256", "signed", "reason", "signature")
SIGNATURE_KEYS = ("algorithm", "value", "certificate_sha256", "signed_over")
CARRIER_SIGNED_OVER = "manifest_canonical_json_utf8"
ALGORITHM_EC_P256 = "ecdsa-p256-sha256"
ALGORITHM_RSA = "rsa-pkcs1v15-sha256"
GENERATED_AT_FORMAT = "%Y-%m-%dT%H:%M:%SZ"
HTML_META_NAMES = (
    "varsel-ai-transparency",
    "varsel-ai-transparency-sha256",
    "varsel-ai-transparency-signature",
    "varsel-ai-transparency-cert-sha256",
)
MANIFEST_KEYS = (
    "spec", "statement", "system", "provider", "generated_at", "ai_generated",
    "digital_source_type", "models", "models_scope",
)

EXIT_VERIFIED = 0
EXIT_FAILED = 1
EXIT_USAGE = 2
EXIT_UNMARKED = 3
EXIT_UNVERIFIED = 4


class EnvironmentProblem(Exception):
    """A missing library, an unreadable anchor, or an internal reading error: exit 2."""


def canonical_json(manifest) -> str:
    """Canonical serialisation: sorted keys at every level, no whitespace, UTF-8 text."""
    return json.dumps(manifest, sort_keys=True, separators=(",", ":"), ensure_ascii=False)


class _Checks:
    def __init__(self):
        self.failures = []

    def __call__(self, label, ok, detail=""):
        print(f"  [{'ok' if ok else 'FAIL'}] {label}" + (f": {detail}" if detail else ""))
        if not ok:
            self.failures.append(label)


class _Anchor:
    """The trust anchor: the cryptography certificate object and its DER SHA-256."""

    def __init__(self, path: Path):
        try:
            from cryptography import x509
            from cryptography.hazmat.primitives import serialization
        except ImportError as exc:
            raise EnvironmentProblem(f"missing library: {exc.name}") from exc
        try:
            self.cert = x509.load_pem_x509_certificate(path.read_bytes())
        except Exception as exc:
            raise EnvironmentProblem(f"anchor is not a PEM certificate: {path} ({type(exc).__name__})") from exc
        self.name = path.name
        self.sha256 = hashlib.sha256(self.cert.public_bytes(serialization.Encoding.DER)).hexdigest()

    def algorithm(self) -> str:
        from cryptography.hazmat.primitives.asymmetric import ec, rsa

        key = self.cert.public_key()
        if isinstance(key, ec.EllipticCurvePublicKey):
            return ALGORITHM_EC_P256 if key.curve.name == "secp256r1" else f"ecdsa-{key.curve.name}-sha256"
        if isinstance(key, rsa.RSAPublicKey):
            return ALGORITHM_RSA
        return "unsupported"

    def signature_valid(self, algorithm, value_b64, data: bytes) -> bool:
        """Detached signature check with the certificate's public key; any fault is False."""
        import base64
        from cryptography.hazmat.primitives import hashes
        from cryptography.hazmat.primitives.asymmetric import ec, padding

        try:
            sig = base64.b64decode(value_b64, validate=True)
            if algorithm != self.algorithm():
                return False
            if algorithm.startswith("ecdsa-"):
                self.cert.public_key().verify(sig, data, ec.ECDSA(hashes.SHA256()))
            elif algorithm == ALGORITHM_RSA:
                self.cert.public_key().verify(sig, data, padding.PKCS1v15(), hashes.SHA256())
            else:
                return False
            return True
        except Exception:
            return False


def _manifest_checks(check, manifest_text, digest):
    """Shape, canonical form, digest and spec of a manifest carried as text."""
    try:
        manifest = json.loads(manifest_text)
    except Exception:
        manifest = None
    check("manifest is JSON", isinstance(manifest, dict))
    if not isinstance(manifest, dict):
        return None
    check("manifest key set", sorted(manifest) == sorted(MANIFEST_KEYS), str(sorted(manifest)))
    check("manifest is canonical", canonical_json(manifest) == manifest_text)
    check("manifest_sha256 = sha256(canonical manifest)",
          digest == hashlib.sha256(manifest_text.encode("utf-8")).hexdigest(), str(digest))
    check("spec", manifest.get("spec") == SPEC, str(manifest.get("spec")))
    check("digital_source_type", manifest.get("digital_source_type") == DIGITAL_SOURCE_TYPE)
    check("generated_at format", bool(re.match(r"^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}Z$",
                                               str(manifest.get("generated_at")))),
          str(manifest.get("generated_at")))
    return manifest


def _finish(check, manifest, anchor) -> int:
    if isinstance(manifest, dict):
        print("manifest:")
        print(json.dumps(manifest, indent=2, ensure_ascii=False))
    if check.failures:
        print(f"RESULT: MARKED WITH FAILURES: {check.failures}")
        return EXIT_FAILED
    if anchor is None:
        print("RESULT: MARKED, NOT VERIFIED (no anchor given: the signer is unverified; "
              "pass the Varsel certificate as the second argument)")
        return EXIT_UNVERIFIED
    print(f"RESULT: MARKED OK (verified against anchor {anchor.name}, sha256 {anchor.sha256})")
    return EXIT_VERIFIED


def _print_file_line(path: Path, data: bytes, kind: str) -> None:
    print(hashlib.sha256(data).hexdigest())
    print(f"file: {path.name} ({len(data)} bytes, {kind})")


# --------------------------------------------------------------------------- JSON
def verify_json(path: Path, anchor) -> int:
    """The `ai_transparency` key of a POST /screen/preview response."""
    data = path.read_bytes()
    _print_file_line(path, data, "preview JSON")
    check = _Checks()
    try:
        response = json.loads(data.decode("utf-8"))
    except Exception as exc:
        print(f"RESULT: UNMARKED (file is not JSON: {type(exc).__name__})")
        return EXIT_UNMARKED
    carrier = response.get("ai_transparency") if isinstance(response, dict) else None
    if not isinstance(carrier, dict):
        print("RESULT: UNMARKED (no ai_transparency object in the response)")
        return EXIT_UNMARKED
    if carrier.get("signed") is not True:
        print(f"RESULT: UNMARKED (carrier present, unsigned, reason={carrier.get('reason')})")
        return EXIT_UNMARKED
    check("carrier key set", sorted(carrier) == sorted(CARRIER_KEYS), str(sorted(carrier)))
    manifest = carrier.get("manifest")
    text = canonical_json(manifest) if isinstance(manifest, dict) else ""
    manifest = _manifest_checks(check, text, carrier.get("manifest_sha256"))
    signature = carrier.get("signature")
    check("signature object", isinstance(signature, dict) and sorted(signature) == sorted(SIGNATURE_KEYS),
          str(sorted(signature)) if isinstance(signature, dict) else str(signature))
    if isinstance(signature, dict):
        check("signed_over", signature.get("signed_over") == CARRIER_SIGNED_OVER, str(signature.get("signed_over")))
        print(f"  algorithm {signature.get('algorithm')} certificate_sha256 {signature.get('certificate_sha256')}")
        if anchor is not None:
            print(f"  anchor {anchor.name} sha256 {anchor.sha256}")
            check("certificate_sha256 = anchor sha256", signature.get("certificate_sha256") == anchor.sha256)
            check("signature valid over the canonical manifest",
                  anchor.signature_valid(signature.get("algorithm"), signature.get("value"), text.encode("utf-8")))
        else:
            print("  (no anchor given: the signature value cannot be checked)")
    return _finish(check, manifest, anchor)


# --------------------------------------------------------------------------- HTML
def _html_metas(text: str) -> dict:
    """name -> [content, ...] of every <meta> inside <head>."""
    from html.parser import HTMLParser

    class _Head(HTMLParser):
        def __init__(self):
            super().__init__(convert_charrefs=True)
            self.in_head = False
            self.metas = {}

        def handle_starttag(self, tag, attrs):
            if tag == "head":
                self.in_head = True
            elif tag == "meta" and self.in_head:
                d = dict(attrs)
                if d.get("name"):
                    self.metas.setdefault(d["name"], []).append(d.get("content") or "")

        def handle_startendtag(self, tag, attrs):
            self.handle_starttag(tag, attrs)

        def handle_endtag(self, tag):
            if tag == "head":
                self.in_head = False

    parser = _Head()
    parser.feed(text)
    parser.close()
    return parser.metas


def verify_html(path: Path, anchor) -> int:
    """The <meta> carriers in the head of a Varsel report HTML."""
    data = path.read_bytes()
    _print_file_line(path, data, "report HTML")
    check = _Checks()
    metas = _html_metas(data.decode("utf-8", "replace"))
    carried = {n: metas.get(n, []) for n in HTML_META_NAMES}
    if not carried[HTML_META_NAMES[0]]:
        print("RESULT: UNMARKED (no varsel-ai-transparency meta in the head)")
        return EXIT_UNMARKED
    if not carried[HTML_META_NAMES[2]]:
        print("RESULT: UNMARKED (manifest carried unsigned: no signature meta)")
        return EXIT_UNMARKED
    for name in HTML_META_NAMES:
        check(f"exactly one {name} meta", len(carried[name]) == 1, str(len(carried[name])))
    text = carried[HTML_META_NAMES[0]][0]
    digest = (carried[HTML_META_NAMES[1]] or [None])[0]
    manifest = _manifest_checks(check, text, digest)
    value = carried[HTML_META_NAMES[2]][0]
    cert_sha = (carried[HTML_META_NAMES[3]] or [None])[0]
    print(f"  certificate_sha256 {cert_sha}")
    if anchor is not None:
        print(f"  anchor {anchor.name} sha256 {anchor.sha256}")
        check("cert-sha256 meta = anchor sha256", cert_sha == anchor.sha256)
        check("signature valid over the manifest meta",
              anchor.signature_valid(anchor.algorithm(), value, text.encode("utf-8")))
    else:
        print("  (no anchor given: the signature value cannot be checked)")
    return _finish(check, manifest, anchor)


# --------------------------------------------------------------------------- PDF
def _read_xmp_raw(data: bytes) -> dict:
    """Second, library-free reader of the one XMP packet (stdlib only). Never raises."""
    raw = {"manifest": None, "sha": None, "dst_resource": None, "dst_text": None, "dst_count": 0,
           "packets": 0, "error": None}
    starts = [m.start() for m in re.finditer(rb"<x:xmpmeta\b", data)]
    raw["packets"] = len(starts)
    if not starts:
        return raw
    try:
        end = data.index(b"</x:xmpmeta>", starts[0]) + len(b"</x:xmpmeta>")
        root = ET.fromstring(data[starts[0]:end].decode("utf-8"))
    except Exception as exc:
        raw["error"] = type(exc).__name__
        return raw
    for desc in root.iter(f"{{{RDF_NAMESPACE}}}Description"):
        for child in desc:
            if child.tag == f"{{{XMP_NAMESPACE}}}{XMP_MANIFEST_PROPERTY}":
                raw["manifest"] = child.text or ""
            elif child.tag == f"{{{XMP_NAMESPACE}}}{XMP_SHA256_PROPERTY}":
                raw["sha"] = child.text
            elif child.tag == f"{{{IPTC_NAMESPACE}}}DigitalSourceType":
                raw["dst_count"] += 1
                raw["dst_resource"] = child.get(f"{{{RDF_NAMESPACE}}}resource")
                raw["dst_text"] = child.text
    return raw


def verify_pdf(pdf_path: Path, anchor) -> int:
    import logging

    logging.disable(logging.CRITICAL)  # pyHanko logs expected validation failures as errors
    try:
        import pikepdf
        from pyhanko.pdf_utils.reader import PdfFileReader
        from pyhanko.sign.diff_analysis import ModificationLevel
        from pyhanko.sign.validation import SignatureCoverageLevel, validate_pdf_signature
        from pyhanko_certvalidator import ValidationContext
    except ImportError as exc:
        raise EnvironmentProblem(f"missing library for PDF mode: {exc.name} (pip install pikepdf pyHanko)") from exc

    data = pdf_path.read_bytes()
    _print_file_line(pdf_path, data, "PDF")
    check = _Checks()

    sig_error = None
    try:
        sigs = list(PdfFileReader(io.BytesIO(data)).embedded_signatures)
    except Exception as exc:
        sigs = []
        sig_error = type(exc).__name__
        print(f"  signature reader error: {sig_error}")

    pike = {"manifest": None, "sha": None}
    try:
        with pikepdf.open(io.BytesIO(data)) as pdf:
            meta = pdf.open_metadata()
            pike = {"manifest": meta.get(f"{{{XMP_NAMESPACE}}}{XMP_MANIFEST_PROPERTY}"),
                    "sha": meta.get(f"{{{XMP_NAMESPACE}}}{XMP_SHA256_PROPERTY}")}
    except Exception as exc:
        print(f"  pikepdf reader error: {type(exc).__name__}")

    raw = _read_xmp_raw(data)
    if raw["error"]:
        print(f"  raw reader error: {raw['error']}")
    has_manifest = bool(pike["manifest"] or raw["manifest"])
    if sig_error and not has_manifest:
        print(f"RESULT: UNMARKED (file could not be read as a PDF: {sig_error}; no XMP manifest found)")
        return EXIT_UNMARKED
    if not has_manifest and not sigs:
        print("RESULT: UNMARKED (no XMP manifest, no signature)")
        return EXIT_UNMARKED
    if has_manifest and not sigs and not sig_error:
        print("RESULT: UNMARKED (manifest carried unsigned: no signature in the file)")
        return EXIT_UNMARKED
    # From here the file carries a manifest, or a signature, or could not be read although it
    # carries a manifest (a damaged signed file): every one of these is reported as a marked file
    # with named failures (exit 1), never as "nothing to see".
    if sig_error:
        check("file readable by the signature reader (damaged or truncated signed file?)", False, sig_error)
    if not has_manifest:
        check("a signature implies a manifest", False, f"sigs={len(sigs)}, no XMP manifest")

    check("exactly one XMP packet", raw["packets"] == 1, str(raw["packets"]))
    manifest_text = pike["manifest"] or raw["manifest"] or ""
    check("two readers agree on the manifest", bool(manifest_text) and pike["manifest"] == raw["manifest"])
    check("two readers agree on the stored hash", pike["sha"] == raw["sha"] and bool(pike["sha"]))
    manifest = _manifest_checks(check, manifest_text, pike["sha"] or raw["sha"])
    check("DigitalSourceType rdf:resource is trainedAlgorithmicMedia",
          raw["dst_resource"] == DIGITAL_SOURCE_TYPE and raw["dst_count"] == 1,
          f"resource={raw['dst_resource']} count={raw['dst_count']}")

    check("exactly one signature", len(sigs) == 1, str(len(sigs)))
    if not sigs:
        return _finish(check, manifest, anchor)  # failures are already recorded: exit 1
    sig = sigs[0]
    check(f"signature field {SIGNATURE_FIELD_NAME}", sig.field_name == SIGNATURE_FIELD_NAME, str(sig.field_name))
    try:
        st0 = validate_pdf_signature(sig, ValidationContext(trust_roots=[], allow_fetching=False))
        print(f"  without anchor: intact={st0.intact} valid={st0.valid} trusted={st0.trusted}")
        check("intact (file unchanged since signing)", st0.intact)
        check("valid (signature mathematics hold)", st0.valid)
    except Exception as exc:
        check("signature can be validated", False, type(exc).__name__)
    if anchor is not None:
        from asn1crypto import x509 as asn1_x509
        from cryptography.hazmat.primitives import serialization

        root = asn1_x509.Certificate.load(anchor.cert.public_bytes(serialization.Encoding.DER))
        print(f"  anchor {anchor.name} sha256 {anchor.sha256}")
        try:
            st = validate_pdf_signature(sig, ValidationContext(trust_roots=[root], allow_fetching=False,
                                                               revocation_mode="soft-fail"))
            print(f"  with anchor: intact={st.intact} valid={st.valid} trusted={st.trusted} "
                  f"coverage={st.coverage} modification={st.modification_level}")
            check("trusted with the anchor", st.trusted)
            check("coverage ENTIRE_FILE", st.coverage == SignatureCoverageLevel.ENTIRE_FILE)
            check("modification NONE", st.modification_level == ModificationLevel.NONE)
        except Exception as exc:
            check("trusted with the anchor", False, type(exc).__name__)
    else:
        print("  (no anchor given: a self-signed certificate is never trusted without it)")
    try:
        reported = sig.self_reported_timestamp
    except Exception:
        reported = None
    gen = (manifest or {}).get("generated_at")
    try:
        expected = datetime.strptime(gen, GENERATED_AT_FORMAT).replace(tzinfo=timezone.utc)
    except Exception:
        expected = None
    rep_utc = reported.astimezone(timezone.utc) if reported else None
    check("signing time equals generated_at (self-asserted, signed by Varsel)",
          expected is not None and rep_utc == expected, f"signed={rep_utc} generated_at={gen}")
    return _finish(check, manifest, anchor)


# --------------------------------------------------------------------------- main
def _is_html(path: Path) -> bool:
    if path.suffix.lower() in (".html", ".htm"):
        return True
    try:
        head = path.read_bytes()[:64].lstrip().lower()
    except OSError:
        return False
    return head.startswith(b"<!doctype") or head.startswith(b"<html")


def main(argv) -> int:
    args = [a for a in argv[1:] if a != "--json"]
    if not args or len(args) > 2:
        print(__doc__)
        return EXIT_USAGE
    target = Path(args[0])
    if not target.is_file():
        print(f"not a file: {target}")
        return EXIT_USAGE
    anchor = None
    if len(args) > 1:
        anchor_path = Path(args[1])
        if not anchor_path.is_file():
            print(f"not a file: {anchor_path}")
            return EXIT_USAGE
        anchor = _Anchor(anchor_path)
    if "--json" in argv or target.suffix.lower() == ".json":
        return verify_json(target, anchor)
    if _is_html(target):
        return verify_html(target, anchor)
    return verify_pdf(target, anchor)


def run(argv) -> int:
    """main() with the exit-code contract enforced: no traceback ever reaches the caller."""
    try:
        return main(argv)
    except EnvironmentProblem as exc:
        print(f"ERROR: {exc}")
        return EXIT_USAGE
    except Exception as exc:  # an unexpected reading error is an environment problem, not a verdict
        print(f"ERROR: internal error while reading the input ({type(exc).__name__})")
        return EXIT_USAGE


if __name__ == "__main__":
    sys.exit(run(sys.argv))
