#!/usr/bin/env python3
"""
MCRO Certificate Deduplication Registry Builder
================================================
Takes the output of mcro_sig_extract.py and:

  1. Deduplicates all exported certificates by SHA-256 fingerprint
  2. Builds a canonical 59-cert (or however many unique) reference library
  3. Creates a cross-reference index: cert → [PDFs, sig fields, timestamps]
  4. Categorizes certs by PKI layer, role, and CA chain
  5. Generates a SHA-256 manifest of the entire registry for OTS stamping
  6. Produces human-readable registry report + machine-readable JSON

Usage:
  python3 mcro_cert_registry.py /path/to/sig_extraction_output [/path/to/registry_output]

Expects the master_signature_index.json from mcro_sig_extract.py.
"""

import sys
import os
import json
import hashlib
import shutil
from pathlib import Path
from datetime import datetime, timezone
from collections import defaultdict, OrderedDict


def sha256_file(path: str) -> str:
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(65536), b""):
            h.update(chunk)
    return h.hexdigest()


def sha256_bytes(data: bytes) -> str:
    return hashlib.sha256(data).hexdigest()


def classify_ca_chain(cert: dict) -> str:
    """Classify which PKI layer a certificate belongs to."""
    issuer_cn = cert.get("issuer", {}).get("commonName",
                cert.get("issuer", {}).get("Common Name", ""))
    subject_cn = cert.get("subject", {}).get("commonName",
                 cert.get("subject", {}).get("Common Name", ""))
    
    if "MJB - Issuing CA" in str(issuer_cn):
        return "Sectigo Enterprise (MJB - Issuing CA)"
    elif "MJB - Root CA" in str(issuer_cn) or "MJB - Root CA" in str(subject_cn):
        return "Sectigo Enterprise (MJB Root)"
    elif "courts-J00PCERTCA-CA" in str(issuer_cn) or "courts-J00PCERTCA-CA" in str(subject_cn):
        return "Internal AD CS (courts-J00PCERTCA-CA)"
    elif "MN Courts Signing CA" in str(issuer_cn) or "MN Courts Signing CA" in str(subject_cn):
        return "Legacy E-Filing (MN Courts Signing CA)"
    elif "ESolutions" in str(subject_cn):
        return "Vendor (ESolutions)"
    else:
        return "Other / External"


def classify_role(cert: dict) -> str:
    """Classify the functional role of a certificate."""
    subject_cn = cert.get("subject", {}).get("commonName",
                 cert.get("subject", {}).get("Common Name", ""))
    subject_cn = str(subject_cn)
    
    if cert.get("is_self_signed"):
        if "Root CA" in subject_cn or "Signing CA" in subject_cn or "CERTCA" in subject_cn:
            return "Root CA"
        elif "ESolutions" in subject_cn:
            return "Vendor Root CA"
        else:
            return "Self-Signed (other)"
    
    if "Issuing CA" in subject_cn:
        return "Intermediate CA"
    elif "MCRO" in subject_cn and "Watermark" in subject_cn:
        return "MCRO Watermark (automated)"
    elif "Integration Service Account" in subject_cn:
        return "Integration Service (automated)"
    elif "(Judge)" in subject_cn or "(Anoka Judge)" in subject_cn:
        return "Judicial Officer (Judge)"
    else:
        # Check if it looks like a person name (Last, First pattern)
        if "," in subject_cn and len(subject_cn.split(",")) == 2:
            return "Court Staff / Officer"
        return "Other"


def determine_verifiability(cert: dict) -> dict:
    """Assess external verifiability of a certificate."""
    exts = cert.get("extensions", {})
    aia = exts.get("authority_information_access", {})
    crls = exts.get("crl_distribution_points", [])
    
    has_public_ocsp = False
    has_public_crl = False
    has_ldap_crl = False
    ocsp_url = None
    crl_url = None
    
    for method, url in aia.items():
        if "OCSP" in method or "ocsp" in str(url):
            url_str = str(url)
            if url_str.startswith("http"):
                has_public_ocsp = True
                ocsp_url = url_str
    
    for crl in crls:
        crl_str = str(crl)
        if crl_str.startswith("http"):
            has_public_crl = True
            crl_url = crl_str
        elif crl_str.startswith("ldap"):
            has_ldap_crl = True
    
    return {
        "externally_verifiable": has_public_ocsp or has_public_crl,
        "has_public_ocsp": has_public_ocsp,
        "has_public_crl": has_public_crl,
        "has_ldap_only_crl": has_ldap_crl and not has_public_crl,
        "ocsp_url": ocsp_url,
        "crl_url": crl_url,
    }


def main():
    if len(sys.argv) < 2:
        print(f"Usage: {sys.argv[0]} <sig_extraction_output_dir> [registry_output_dir]")
        sys.exit(1)

    src_dir = Path(sys.argv[1])
    master_json = src_dir / "master_signature_index.json"
    
    if not master_json.exists():
        print(f"Error: {master_json} not found. Run mcro_sig_extract.py first.")
        sys.exit(1)
    
    if len(sys.argv) >= 3:
        out_dir = Path(sys.argv[2])
    else:
        out_dir = src_dir / "cert_registry"
    
    out_dir.mkdir(parents=True, exist_ok=True)
    canonical_dir = out_dir / "canonical_certs"
    canonical_dir.mkdir(exist_ok=True)

    print(f"Loading master index from {master_json} ...")
    with open(master_json) as f:
        master = json.load(f)
    
    print(f"Processing {master['total_pdfs_processed']} PDFs with {master['total_signatures_found']} signatures ...\n")

    # -----------------------------------------------------------------------
    # Pass 1: Deduplicate certs and build cross-reference
    # -----------------------------------------------------------------------
    
    # cert_sha256 -> { cert_data, appearances: [ {pdf, sig_field, sig_timestamp, sig_subfilter} ] }
    registry = OrderedDict()
    
    total_cert_instances = 0
    
    for pdf_rec in master.get("pdfs", []):
        pdf_file = pdf_rec.get("file", "unknown")
        pdf_name = Path(pdf_file).name
        pdf_sha256 = pdf_rec.get("file_sha256", "")
        
        for sig in pdf_rec.get("signatures", []):
            sig_field = sig.get("field_name", "unknown")
            sig_meta = sig.get("metadata", {})
            sig_timestamp = sig_meta.get("M", "")
            sig_subfilter = sig_meta.get("SubFilter", "unknown")
            sig_reason = sig_meta.get("Reason", "")
            sig_location = sig_meta.get("Location", "")
            
            for cert in sig.get("certificates", []):
                total_cert_instances += 1
                fp = cert.get("fingerprints", {}).get("sha256", "")
                if not fp:
                    continue
                
                if fp not in registry:
                    registry[fp] = {
                        "cert": cert,
                        "appearances": [],
                        "pdfs_set": set(),
                        "sig_subfilters": set(),
                    }
                
                registry[fp]["appearances"].append({
                    "pdf_filename": pdf_name,
                    "pdf_sha256": pdf_sha256,
                    "sig_field": sig_field,
                    "sig_timestamp": sig_timestamp,
                    "sig_subfilter": sig_subfilter,
                    "sig_reason": sig_reason,
                    "sig_location": sig_location,
                })
                registry[fp]["pdfs_set"].add(pdf_name)
                registry[fp]["sig_subfilters"].add(sig_subfilter)
    
    print(f"Total certificate instances across all signatures: {total_cert_instances:,}")
    print(f"Unique certificates by SHA-256: {len(registry)}")

    # -----------------------------------------------------------------------
    # Pass 2: Enrich each unique cert with classification + verifiability
    # -----------------------------------------------------------------------
    
    registry_entries = []
    
    for i, (fp, data) in enumerate(sorted(registry.items(),
            key=lambda x: len(x[1]["pdfs_set"]), reverse=True), 1):
        
        cert = data["cert"]
        subject_cn = cert.get("subject", {}).get("commonName",
                     cert.get("subject", {}).get("Common Name", "unknown"))
        issuer_cn = cert.get("issuer", {}).get("commonName",
                    cert.get("issuer", {}).get("Common Name", "unknown"))
        
        entry = {
            "registry_id": i,
            "sha256_fingerprint": fp,
            "sha1_fingerprint": cert.get("fingerprints", {}).get("sha1", ""),
            "subject_cn": subject_cn,
            "subject_full": cert.get("subject", {}),
            "issuer_cn": issuer_cn,
            "issuer_full": cert.get("issuer", {}),
            "serial_hex": cert.get("serial_number_hex", ""),
            "serial_int": cert.get("serial_number_int", ""),
            "not_before_utc": cert.get("not_before_utc", ""),
            "not_after_utc": cert.get("not_after_utc", ""),
            "validity_days": None,
            "is_currently_valid": False,
            "is_expired": False,
            "signature_algorithm": cert.get("signature_algorithm", ""),
            "public_key_algorithm": cert.get("public_key_algorithm", ""),
            "public_key_size_bits": cert.get("public_key_size_bits", 0),
            "is_self_signed": cert.get("is_self_signed", False),
            "pki_layer": classify_ca_chain(cert),
            "functional_role": classify_role(cert),
            "verifiability": determine_verifiability(cert),
            "extensions": cert.get("extensions", {}),
            "sig_subfilters_seen": sorted(data["sig_subfilters"]),
            "total_pdf_appearances": len(data["pdfs_set"]),
            "total_signature_appearances": len(data["appearances"]),
            "canonical_cer_file": None,
            "canonical_pem_file": None,
            "appearances": data["appearances"],
        }
        
        # Validity calculations
        try:
            nb = datetime.fromisoformat(entry["not_before_utc"])
            na = datetime.fromisoformat(entry["not_after_utc"])
            now = datetime.now(timezone.utc)
            entry["validity_days"] = (na - nb).days
            entry["is_currently_valid"] = nb <= now <= na
            entry["is_expired"] = now > na
        except Exception:
            pass
        
        # Copy canonical cert files
        # Find one of the exported .cer files for this cert
        exported_cer = cert.get("exported_cer", "")
        exported_pem = cert.get("exported_pem", "")
        
        cn_safe = "".join(c if c.isalnum() or c in "-_ " else "_" for c in str(subject_cn))[:50]
        canon_name = f"REG{i:03d}__{cn_safe}"
        
        if exported_cer and os.path.exists(exported_cer):
            dest_cer = canonical_dir / f"{canon_name}.cer"
            shutil.copy2(exported_cer, dest_cer)
            entry["canonical_cer_file"] = str(dest_cer)
        
        if exported_pem and os.path.exists(exported_pem):
            dest_pem = canonical_dir / f"{canon_name}.pem"
            shutil.copy2(exported_pem, dest_pem)
            entry["canonical_pem_file"] = str(dest_pem)
        
        registry_entries.append(entry)
    
    # -----------------------------------------------------------------------
    # Pass 3: Build chain linkage map
    # -----------------------------------------------------------------------
    
    # Map issuer CN + AKI -> cert entries that could be the issuer
    cert_by_ski = {}
    for entry in registry_entries:
        ski = entry.get("extensions", {}).get("subject_key_identifier", "")
        if ski:
            cert_by_ski[ski] = entry["registry_id"]
    
    for entry in registry_entries:
        aki = entry.get("extensions", {}).get("authority_key_identifier", "")
        if aki and aki in cert_by_ski:
            entry["issuer_registry_id"] = cert_by_ski[aki]
        else:
            entry["issuer_registry_id"] = None
    
    # -----------------------------------------------------------------------
    # Build output structures
    # -----------------------------------------------------------------------
    
    registry_output = {
        "generated_utc": datetime.now(timezone.utc).isoformat(),
        "source_extraction_dir": str(src_dir.resolve()),
        "total_pdfs_in_corpus": master["total_pdfs_processed"],
        "total_signatures_in_corpus": master["total_signatures_found"],
        "total_cert_instances": total_cert_instances,
        "unique_certificates": len(registry_entries),
        "summary_by_pki_layer": {},
        "summary_by_role": {},
        "summary_by_verifiability": {
            "externally_verifiable": 0,
            "internal_only": 0,
        },
        "summary_by_status": {
            "currently_valid": 0,
            "expired": 0,
        },
        "certificates": registry_entries,
    }
    
    # Summaries
    for entry in registry_entries:
        layer = entry["pki_layer"]
        role = entry["functional_role"]
        registry_output["summary_by_pki_layer"][layer] = \
            registry_output["summary_by_pki_layer"].get(layer, 0) + 1
        registry_output["summary_by_role"][role] = \
            registry_output["summary_by_role"].get(role, 0) + 1
        if entry["verifiability"]["externally_verifiable"]:
            registry_output["summary_by_verifiability"]["externally_verifiable"] += 1
        else:
            registry_output["summary_by_verifiability"]["internal_only"] += 1
        if entry["is_currently_valid"]:
            registry_output["summary_by_status"]["currently_valid"] += 1
        if entry["is_expired"]:
            registry_output["summary_by_status"]["expired"] += 1

    # Write registry JSON
    registry_json_path = out_dir / "cert_registry.json"
    with open(registry_json_path, "w") as f:
        json.dump(registry_output, f, indent=2, default=str)
    print(f"\nRegistry JSON: {registry_json_path}")
    
    # -----------------------------------------------------------------------
    # Cross-reference CSV (cert -> PDFs)
    # -----------------------------------------------------------------------
    
    xref_path = out_dir / "cert_pdf_crossref.csv"
    with open(xref_path, "w") as f:
        f.write("registry_id,cert_sha256,subject_cn,issuer_cn,pki_layer,role,"
                "not_before,not_after,is_expired,externally_verifiable,"
                "pdf_filename,sig_field,sig_timestamp,sig_subfilter\n")
        for entry in registry_entries:
            for app in entry["appearances"]:
                cn = str(entry["subject_cn"]).replace('"', '""')
                icn = str(entry["issuer_cn"]).replace('"', '""')
                f.write(f'{entry["registry_id"]},'
                        f'{entry["sha256_fingerprint"]},'
                        f'"{cn}",'
                        f'"{icn}",'
                        f'"{entry["pki_layer"]}",'
                        f'"{entry["functional_role"]}",'
                        f'{entry["not_before_utc"]},'
                        f'{entry["not_after_utc"]},'
                        f'{entry["is_expired"]},'
                        f'{entry["verifiability"]["externally_verifiable"]},'
                        f'"{app["pdf_filename"]}",'
                        f'"{app["sig_field"]}",'
                        f'"{app["sig_timestamp"]}",'
                        f'"{app["sig_subfilter"]}"\n')
    print(f"Cross-reference CSV: {xref_path}")
    
    # -----------------------------------------------------------------------
    # Compact cross-reference (cert -> PDF list only, no per-sig detail)
    # -----------------------------------------------------------------------
    
    compact_xref_path = out_dir / "cert_pdf_compact.csv"
    with open(compact_xref_path, "w") as f:
        f.write("registry_id,cert_sha256,subject_cn,pki_layer,role,"
                "total_pdfs,total_sigs,is_expired,verifiable\n")
        for entry in registry_entries:
            cn = str(entry["subject_cn"]).replace('"', '""')
            f.write(f'{entry["registry_id"]},'
                    f'{entry["sha256_fingerprint"][:16]}...,'
                    f'"{cn}",'
                    f'"{entry["pki_layer"]}",'
                    f'"{entry["functional_role"]}",'
                    f'{entry["total_pdf_appearances"]},'
                    f'{entry["total_signature_appearances"]},'
                    f'{entry["is_expired"]},'
                    f'{entry["verifiability"]["externally_verifiable"]}\n')
    print(f"Compact cross-ref: {compact_xref_path}")
    
    # -----------------------------------------------------------------------
    # Human-readable report
    # -----------------------------------------------------------------------
    
    lines = []
    lines.append("=" * 95)
    lines.append("MCRO CERTIFICATE DEDUPLICATION REGISTRY")
    lines.append(f"Generated: {registry_output['generated_utc']}")
    lines.append(f"Source: {registry_output['source_extraction_dir']}")
    lines.append("=" * 95)
    lines.append(f"\nCorpus: {registry_output['total_pdfs_in_corpus']:,} PDFs, "
                 f"{registry_output['total_signatures_in_corpus']:,} signatures")
    lines.append(f"Certificate instances: {total_cert_instances:,}")
    lines.append(f"Unique certificates:  {len(registry_entries)}")
    
    lines.append(f"\nBY PKI LAYER:")
    for layer, count in sorted(registry_output["summary_by_pki_layer"].items()):
        lines.append(f"  {layer}: {count}")
    
    lines.append(f"\nBY FUNCTIONAL ROLE:")
    for role, count in sorted(registry_output["summary_by_role"].items()):
        lines.append(f"  {role}: {count}")
    
    lines.append(f"\nVERIFIABILITY:")
    lines.append(f"  Externally verifiable (public OCSP/CRL): "
                 f"{registry_output['summary_by_verifiability']['externally_verifiable']}")
    lines.append(f"  Internal only (LDAP/self-signed):        "
                 f"{registry_output['summary_by_verifiability']['internal_only']}")
    
    lines.append(f"\nSTATUS:")
    lines.append(f"  Currently valid: {registry_output['summary_by_status']['currently_valid']}")
    lines.append(f"  Expired:         {registry_output['summary_by_status']['expired']}")
    
    lines.append(f"\n{'=' * 95}")
    lines.append("CANONICAL CERTIFICATE REGISTRY")
    lines.append(f"{'=' * 95}")
    
    for entry in registry_entries:
        lines.append(f"\n  REG#{entry['registry_id']:03d}  {entry['subject_cn']}")
        lines.append(f"    Issuer:       {entry['issuer_cn']}"
                     f"{' → REG#'+str(entry['issuer_registry_id']).zfill(3) if entry.get('issuer_registry_id') else ''}")
        lines.append(f"    PKI Layer:    {entry['pki_layer']}")
        lines.append(f"    Role:         {entry['functional_role']}")
        lines.append(f"    Serial:       {entry['serial_hex']}")
        lines.append(f"    Valid:        {entry['not_before_utc'][:10]}  →  {entry['not_after_utc'][:10]}"
                     f"  ({entry['validity_days']} days)"
                     f"  {'✓ ACTIVE' if entry['is_currently_valid'] else '✗ EXPIRED'}")
        lines.append(f"    Algorithm:    {entry['signature_algorithm']}  /  "
                     f"{entry['public_key_algorithm']} {entry['public_key_size_bits']}-bit")
        lines.append(f"    SHA-256:      {entry['sha256_fingerprint']}")
        lines.append(f"    SHA-1:        {entry['sha1_fingerprint']}")
        lines.append(f"    Self-signed:  {entry['is_self_signed']}")
        lines.append(f"    Verifiable:   {'YES (public OCSP/CRL)' if entry['verifiability']['externally_verifiable'] else 'NO (internal only)'}")
        if entry['verifiability']['ocsp_url']:
            lines.append(f"    OCSP:         {entry['verifiability']['ocsp_url']}")
        if entry['verifiability']['crl_url']:
            lines.append(f"    CRL:          {entry['verifiability']['crl_url']}")
        lines.append(f"    SubFilters:   {', '.join(entry['sig_subfilters_seen'])}")
        lines.append(f"    Seen in:      {entry['total_pdf_appearances']} PDFs, "
                     f"{entry['total_signature_appearances']} signatures")
    
    lines.append(f"\n{'=' * 95}")
    lines.append("END OF REGISTRY")
    lines.append(f"{'=' * 95}")
    
    report_path = out_dir / "cert_registry_report.txt"
    report_path.write_text("\n".join(lines))
    print(f"Registry report: {report_path}")
    
    # -----------------------------------------------------------------------
    # SHA-256 manifest of the entire registry
    # -----------------------------------------------------------------------
    
    manifest_lines = []
    for fpath in sorted(out_dir.rglob("*")):
        if fpath.is_file() and fpath.name != "MANIFEST.sha256":
            rel = fpath.relative_to(out_dir)
            h = sha256_file(str(fpath))
            manifest_lines.append(f"{h}  {rel}")
    
    manifest_path = out_dir / "MANIFEST.sha256"
    manifest_path.write_text("\n".join(manifest_lines) + "\n")
    print(f"SHA-256 manifest: {manifest_path}")
    print(f"  ({len(manifest_lines)} files in manifest — OTS-stamp this file)")
    
    # Stats
    print(f"\n{'=' * 60}")
    print(f"Unique certs in registry:     {len(registry_entries)}")
    print(f"Canonical .cer files:         {len(list(canonical_dir.glob('*.cer')))}")
    print(f"Canonical .pem files:         {len(list(canonical_dir.glob('*.pem')))}")
    print(f"Cross-reference CSV rows:     {total_cert_instances:,}")
    print(f"{'=' * 60}")


if __name__ == "__main__":
    main()
