#!/usr/bin/env python3
"""
MCRO OCSP Batch Query & CRL Verification Tool
===============================================
Takes the cert_registry.json from mcro_cert_registry.py and:

  1. Identifies all certificates with public OCSP/CRL endpoints
  2. Fetches the issuer certificate needed for OCSP verification
  3. Queries OCSP for every verifiable leaf certificate
  4. Captures and saves the raw signed OCSP responses (.ocsp.der)
  5. Fetches and parses the current CRL
  6. Cross-checks every cert serial against the CRL revoked list
  7. Generates a comprehensive status report

The raw OCSP responses are cryptographically signed by the responder,
making them independently verifiable evidence of cert status at query time.
OTS-stamp the output directory after running.

Usage:
  python3 mcro_ocsp_batch.py /path/to/cert_registry [/path/to/ocsp_output]

Dependencies:
  pip install cryptography requests
  openssl (CLI) must be available in PATH

Note: Run from a machine with unrestricted network access (not a sandbox).
"""

import sys
import os
import json
import subprocess
import hashlib
import time
from pathlib import Path
from datetime import datetime, timezone
from collections import defaultdict

try:
    import requests
    HAS_REQUESTS = True
except ImportError:
    HAS_REQUESTS = False
    print("[WARN] requests library not available — using curl fallback")


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


def run_cmd(cmd: list, timeout: int = 30) -> dict:
    """Run a command and return stdout, stderr, returncode."""
    try:
        result = subprocess.run(cmd, capture_output=True, timeout=timeout)
        return {
            "stdout": result.stdout.decode("utf-8", errors="replace"),
            "stderr": result.stderr.decode("utf-8", errors="replace"),
            "returncode": result.returncode,
        }
    except subprocess.TimeoutExpired:
        return {"stdout": "", "stderr": "TIMEOUT", "returncode": -1}
    except Exception as e:
        return {"stdout": "", "stderr": str(e), "returncode": -1}


def fetch_url_to_file(url: str, dest_path: str, timeout: int = 30) -> bool:
    """Download a URL to a file. Returns True on success."""
    if HAS_REQUESTS:
        try:
            r = requests.get(url, timeout=timeout)
            if r.status_code == 200:
                with open(dest_path, "wb") as f:
                    f.write(r.content)
                return True
            else:
                print(f"      HTTP {r.status_code} for {url}")
                return False
        except Exception as e:
            print(f"      Request error: {e}")
            return False
    else:
        result = run_cmd(["curl", "-s", "-o", dest_path, "-w", "%{http_code}", url], timeout)
        return result["returncode"] == 0 and "200" in result["stdout"]


def query_ocsp_openssl(cert_pem: str, issuer_pem: str, ocsp_url: str,
                       response_der_path: str, timeout: int = 30) -> dict:
    """
    Query OCSP using openssl and capture the response.
    Returns parsed status dict.
    """
    cmd = [
        "openssl", "ocsp",
        "-issuer", issuer_pem,
        "-cert", cert_pem,
        "-url", ocsp_url,
        "-resp_text",
        "-respout", response_der_path,
    ]
    
    result = run_cmd(cmd, timeout)
    
    status = {
        "ocsp_url": ocsp_url,
        "query_time_utc": datetime.now(timezone.utc).isoformat(),
        "raw_stdout": result["stdout"],
        "raw_stderr": result["stderr"],
        "returncode": result["returncode"],
        "response_saved": os.path.exists(response_der_path) and os.path.getsize(response_der_path) > 0,
        "response_der_sha256": None,
        "cert_status": "error",
        "this_update": None,
        "next_update": None,
        "produced_at": None,
        "revocation_time": None,
        "revocation_reason": None,
        "response_verify": None,
    }
    
    if status["response_saved"]:
        with open(response_der_path, "rb") as f:
            status["response_der_sha256"] = sha256hex(f.read())
    
    # Parse the text output
    stdout = result["stdout"]
    
    for line in stdout.split("\n"):
        line = line.strip()
        if line.startswith("Cert Status:"):
            status["cert_status"] = line.split(":", 1)[1].strip()
        elif line.startswith("This Update:"):
            status["this_update"] = line.split(":", 1)[1].strip()
        elif line.startswith("Next Update:"):
            status["next_update"] = line.split(":", 1)[1].strip()
        elif line.startswith("Produced At:"):
            status["produced_at"] = line.split(":", 1)[1].strip()
        elif line.startswith("Revocation Time:"):
            status["revocation_time"] = line.split(":", 1)[1].strip()
        elif line.startswith("Revocation Reason:"):
            status["revocation_reason"] = line.split(":", 1)[1].strip()
        elif "Response verify OK" in line:
            status["response_verify"] = "OK"
        elif "Response Verify Failure" in line:
            status["response_verify"] = "FAILURE"
    
    # If we didn't find a cert_status in the structured output, check for 
    # the summary line at the end (format: "filename: status")
    if status["cert_status"] == "error":
        for line in stdout.split("\n"):
            line = line.strip()
            if ".pem:" in line or ".cer:" in line:
                parts = line.split(":")
                if len(parts) >= 2:
                    s = parts[-1].strip()
                    if s in ["good", "revoked", "unknown"]:
                        status["cert_status"] = s
    
    return status


def parse_crl_for_serials(crl_der_path: str) -> dict:
    """Parse a DER-encoded CRL and extract revoked serial numbers."""
    result = run_cmd(["openssl", "crl", "-in", crl_der_path, "-inform", "DER",
                      "-text", "-noout"], timeout=60)
    
    crl_info = {
        "issuer": None,
        "last_update": None,
        "next_update": None,
        "total_revoked": 0,
        "revoked_serials": set(),
        "parse_error": None,
    }
    
    if result["returncode"] != 0:
        crl_info["parse_error"] = result["stderr"]
        return crl_info
    
    current_serial = None
    for line in result["stdout"].split("\n"):
        line = line.strip()
        if line.startswith("Issuer:"):
            crl_info["issuer"] = line.split(":", 1)[1].strip()
        elif line.startswith("Last Update:"):
            crl_info["last_update"] = line.split(":", 1)[1].strip()
        elif line.startswith("Next Update:"):
            crl_info["next_update"] = line.split(":", 1)[1].strip()
        elif line.startswith("Serial Number:"):
            serial = line.split(":", 1)[1].strip().lower().replace(":", "")
            crl_info["revoked_serials"].add(serial)
    
    crl_info["total_revoked"] = len(crl_info["revoked_serials"])
    return crl_info


def main():
    if len(sys.argv) < 2:
        print(f"Usage: {sys.argv[0]} <cert_registry_dir> [ocsp_output_dir]")
        sys.exit(1)
    
    reg_dir = Path(sys.argv[1])
    reg_json = reg_dir / "cert_registry.json"
    
    if not reg_json.exists():
        # Try parent
        reg_json = reg_dir / "cert_registry" / "cert_registry.json"
        if not reg_json.exists():
            print(f"Error: cert_registry.json not found in {reg_dir}")
            sys.exit(1)
        reg_dir = reg_dir / "cert_registry"
    
    if len(sys.argv) >= 3:
        out_dir = Path(sys.argv[2])
    else:
        out_dir = reg_dir / "ocsp_verification"
    
    out_dir.mkdir(parents=True, exist_ok=True)
    (out_dir / "ocsp_responses").mkdir(exist_ok=True)
    (out_dir / "crl_snapshots").mkdir(exist_ok=True)
    (out_dir / "issuer_certs").mkdir(exist_ok=True)

    # Check openssl
    r = run_cmd(["openssl", "version"])
    if r["returncode"] != 0:
        print("Error: openssl not found in PATH")
        sys.exit(1)
    print(f"Using: {r['stdout'].strip()}\n")

    print(f"Loading registry from {reg_json} ...")
    with open(reg_json) as f:
        registry = json.load(f)
    
    certs = registry.get("certificates", [])
    print(f"Registry contains {len(certs)} unique certificates\n")
    
    # -----------------------------------------------------------------------
    # Identify verifiable certs and their issuers
    # -----------------------------------------------------------------------
    
    verifiable = []
    issuer_map = {}  # issuer_registry_id -> entry
    
    for entry in certs:
        if entry["verifiability"]["externally_verifiable"]:
            verifiable.append(entry)
        # Build issuer lookup
        issuer_map[entry["registry_id"]] = entry
    
    print(f"Certificates with public OCSP/CRL: {len(verifiable)}")
    
    # Group by OCSP URL
    by_ocsp = defaultdict(list)
    for entry in verifiable:
        url = entry["verifiability"].get("ocsp_url", "none")
        by_ocsp[url].append(entry)
    
    print(f"Unique OCSP endpoints: {len(by_ocsp)}")
    for url, entries in by_ocsp.items():
        print(f"  {url}: {len(entries)} certs")
    
    # -----------------------------------------------------------------------
    # Fetch issuer certs needed for OCSP queries
    # -----------------------------------------------------------------------
    
    print(f"\n{'=' * 70}")
    print("STEP 1: Fetching issuer certificates for OCSP verification")
    print(f"{'=' * 70}\n")
    
    # For Sectigo-chained certs, we need the MJB Issuing CA cert
    issuer_certs_needed = set()
    for entry in verifiable:
        aia = entry.get("extensions", {}).get("authority_information_access", {})
        for method, url in aia.items():
            if "caIssuers" in method or "caIssuers" in str(method):
                url_str = str(url)
                if url_str.startswith("http"):
                    issuer_certs_needed.add(url_str)
    
    issuer_pem_paths = {}  # CA Issuer URL -> local PEM path
    
    for url in sorted(issuer_certs_needed):
        fname = Path(url).name
        der_path = str(out_dir / "issuer_certs" / fname)
        pem_path = str(out_dir / "issuer_certs" / f"{Path(fname).stem}.pem")
        
        print(f"  Fetching: {url}")
        if fetch_url_to_file(url, der_path):
            # Convert DER to PEM
            r = run_cmd(["openssl", "x509", "-in", der_path, "-inform", "DER",
                         "-out", pem_path, "-outform", "PEM"])
            if r["returncode"] == 0:
                print(f"    ✓ Saved: {fname} → PEM")
                issuer_pem_paths[url] = pem_path
            else:
                # Maybe already PEM?
                r2 = run_cmd(["openssl", "x509", "-in", der_path, "-inform", "PEM",
                              "-out", pem_path, "-outform", "PEM"])
                if r2["returncode"] == 0:
                    print(f"    ✓ Saved (was already PEM): {fname}")
                    issuer_pem_paths[url] = pem_path
                else:
                    print(f"    ✗ Failed to convert: {r['stderr'][:100]}")
        else:
            print(f"    ✗ Download failed")
    
    # Also check if we have issuer certs in the canonical registry
    # (for cases where the issuer is also in our corpus)
    for entry in verifiable:
        issuer_id = entry.get("issuer_registry_id")
        if issuer_id and issuer_id in issuer_map:
            issuer_entry = issuer_map[issuer_id]
            pem = issuer_entry.get("canonical_pem_file", "")
            if pem and os.path.exists(pem):
                # Use this as fallback
                aia = entry.get("extensions", {}).get("authority_information_access", {})
                for method, url in aia.items():
                    if "caIssuers" in str(method):
                        url_str = str(url)
                        if url_str not in issuer_pem_paths:
                            issuer_pem_paths[url_str] = pem
                            print(f"  Using canonical cert for issuer REG#{issuer_id:03d}")
    
    # -----------------------------------------------------------------------
    # OCSP Queries
    # -----------------------------------------------------------------------
    
    print(f"\n{'=' * 70}")
    print("STEP 2: Querying OCSP for all verifiable certificates")
    print(f"{'=' * 70}\n")
    
    ocsp_results = []
    
    for i, entry in enumerate(verifiable, 1):
        reg_id = entry["registry_id"]
        cn = entry["subject_cn"]
        ocsp_url = entry["verifiability"].get("ocsp_url")
        serial = entry["serial_hex"]
        
        print(f"  [{i}/{len(verifiable)}] REG#{reg_id:03d} {cn}")
        print(f"    Serial: {serial}")
        print(f"    Status: {'EXPIRED' if entry['is_expired'] else 'VALID'} "
              f"(until {entry['not_after_utc'][:10]})")
        
        if not ocsp_url:
            print(f"    ⊘ No OCSP URL — skipping")
            ocsp_results.append({
                "registry_id": reg_id,
                "subject_cn": cn,
                "serial_hex": serial,
                "ocsp_status": "no_ocsp_url",
            })
            continue
        
        # Find cert PEM
        cert_pem = entry.get("canonical_pem_file", "")
        if not cert_pem or not os.path.exists(cert_pem):
            print(f"    ⊘ No PEM file available — skipping")
            ocsp_results.append({
                "registry_id": reg_id,
                "subject_cn": cn,
                "serial_hex": serial,
                "ocsp_status": "no_cert_file",
            })
            continue
        
        # Find issuer PEM
        aia = entry.get("extensions", {}).get("authority_information_access", {})
        issuer_pem = None
        for method, url in aia.items():
            if "caIssuers" in str(method):
                url_str = str(url)
                if url_str in issuer_pem_paths:
                    issuer_pem = issuer_pem_paths[url_str]
                    break
        
        if not issuer_pem:
            print(f"    ⊘ No issuer cert available — skipping")
            ocsp_results.append({
                "registry_id": reg_id,
                "subject_cn": cn,
                "serial_hex": serial,
                "ocsp_status": "no_issuer_cert",
            })
            continue
        
        # Query OCSP
        cn_safe = "".join(c if c.isalnum() or c in "-_" else "_" for c in cn)[:40]
        resp_path = str(out_dir / "ocsp_responses" / f"REG{reg_id:03d}__{cn_safe}.ocsp.der")
        
        print(f"    Querying: {ocsp_url}")
        status = query_ocsp_openssl(cert_pem, issuer_pem, ocsp_url, resp_path)
        
        result_entry = {
            "registry_id": reg_id,
            "subject_cn": cn,
            "serial_hex": serial,
            "is_expired": entry["is_expired"],
            "not_after_utc": entry["not_after_utc"],
            "pki_layer": entry["pki_layer"],
            "functional_role": entry["functional_role"],
            "ocsp_status": status["cert_status"],
            "ocsp_verify": status["response_verify"],
            "ocsp_produced_at": status["produced_at"],
            "ocsp_this_update": status["this_update"],
            "ocsp_next_update": status["next_update"],
            "ocsp_revocation_time": status["revocation_time"],
            "ocsp_revocation_reason": status["revocation_reason"],
            "response_der_saved": status["response_saved"],
            "response_der_sha256": status["response_der_sha256"],
            "response_der_path": resp_path if status["response_saved"] else None,
            "query_time_utc": status["query_time_utc"],
            "raw_stdout": status["raw_stdout"],
        }
        
        ocsp_results.append(result_entry)
        
        status_icon = {"good": "✓", "revoked": "✗", "unknown": "?"}.get(
            status["cert_status"], "⊘")
        print(f"    {status_icon} OCSP Status: {status['cert_status'].upper()}"
              f"  (verify: {status['response_verify'] or 'N/A'})")
        if status["revocation_time"]:
            print(f"    ✗ Revoked at: {status['revocation_time']}")
            if status["revocation_reason"]:
                print(f"    ✗ Reason: {status['revocation_reason']}")
        
        # Be polite to the OCSP responder
        time.sleep(0.5)
    
    # -----------------------------------------------------------------------
    # CRL Fetch and Cross-Check
    # -----------------------------------------------------------------------
    
    print(f"\n{'=' * 70}")
    print("STEP 3: Fetching CRL and cross-checking serials")
    print(f"{'=' * 70}\n")
    
    crl_urls = set()
    for entry in verifiable:
        url = entry["verifiability"].get("crl_url")
        if url:
            crl_urls.add(url)
    
    crl_data = {}
    
    for url in sorted(crl_urls):
        fname = Path(url).name
        crl_path = str(out_dir / "crl_snapshots" / fname)
        
        print(f"  Fetching CRL: {url}")
        if fetch_url_to_file(url, crl_path):
            fsize = os.path.getsize(crl_path)
            print(f"    ✓ Downloaded: {fsize:,} bytes")
            
            crl_info = parse_crl_for_serials(crl_path)
            crl_data[url] = crl_info
            
            print(f"    Issuer: {crl_info['issuer']}")
            print(f"    Last Update: {crl_info['last_update']}")
            print(f"    Next Update: {crl_info['next_update']}")
            print(f"    Revoked certs: {crl_info['total_revoked']}")
        else:
            print(f"    ✗ Download failed")
    
    # Cross-check each verifiable cert against CRL
    crl_results = []
    
    print(f"\n  Cross-checking {len(verifiable)} cert serials against CRL(s)...")
    
    for entry in verifiable:
        crl_url = entry["verifiability"].get("crl_url")
        serial = entry["serial_hex"].lower().lstrip("0") or "0"
        
        crl_status = "no_crl"
        if crl_url and crl_url in crl_data:
            crl_info = crl_data[crl_url]
            # Normalize serial comparison (strip leading zeros)
            revoked_normalized = {s.lstrip("0") or "0" for s in crl_info["revoked_serials"]}
            if serial in revoked_normalized:
                crl_status = "REVOKED"
            else:
                crl_status = "not_revoked"
        
        crl_results.append({
            "registry_id": entry["registry_id"],
            "subject_cn": entry["subject_cn"],
            "serial_hex": entry["serial_hex"],
            "crl_url": crl_url,
            "crl_status": crl_status,
        })
        
        if crl_status == "REVOKED":
            print(f"    ✗ REG#{entry['registry_id']:03d} {entry['subject_cn']}: REVOKED in CRL!")
    
    not_revoked = sum(1 for r in crl_results if r["crl_status"] == "not_revoked")
    revoked = sum(1 for r in crl_results if r["crl_status"] == "REVOKED")
    print(f"\n  CRL results: {not_revoked} not revoked, {revoked} revoked, "
          f"{len(crl_results) - not_revoked - revoked} no CRL available")
    
    # -----------------------------------------------------------------------
    # Output
    # -----------------------------------------------------------------------
    
    print(f"\n{'=' * 70}")
    print("STEP 4: Writing results")
    print(f"{'=' * 70}\n")
    
    output = {
        "generated_utc": datetime.now(timezone.utc).isoformat(),
        "registry_source": str(reg_json),
        "total_certs_queried": len(verifiable),
        "ocsp_summary": {
            "good": sum(1 for r in ocsp_results if r.get("ocsp_status") == "good"),
            "revoked": sum(1 for r in ocsp_results if r.get("ocsp_status") == "revoked"),
            "unknown": sum(1 for r in ocsp_results if r.get("ocsp_status") == "unknown"),
            "error": sum(1 for r in ocsp_results if r.get("ocsp_status") not in 
                        ["good", "revoked", "unknown", "no_ocsp_url", "no_cert_file", "no_issuer_cert"]),
            "skipped": sum(1 for r in ocsp_results if r.get("ocsp_status") in 
                         ["no_ocsp_url", "no_cert_file", "no_issuer_cert"]),
        },
        "crl_summary": {
            "not_revoked": not_revoked,
            "revoked": revoked,
            "no_crl": len(crl_results) - not_revoked - revoked,
        },
        "crl_snapshots": {url: {
            "issuer": info["issuer"],
            "last_update": info["last_update"],
            "next_update": info["next_update"],
            "total_revoked": info["total_revoked"],
        } for url, info in crl_data.items()},
        "ocsp_results": ocsp_results,
        "crl_results": crl_results,
    }
    
    # Remove raw_stdout from saved JSON (too bulky, saved in .der files)
    for r in output["ocsp_results"]:
        r.pop("raw_stdout", None)
    
    results_json = out_dir / "ocsp_crl_results.json"
    with open(results_json, "w") as f:
        json.dump(output, f, indent=2, default=str)
    print(f"  Results JSON: {results_json}")
    
    # Human-readable report
    lines = []
    lines.append("=" * 90)
    lines.append("MCRO OCSP/CRL VERIFICATION REPORT")
    lines.append(f"Generated: {output['generated_utc']}")
    lines.append("=" * 90)
    
    lines.append(f"\nCertificates queried: {output['total_certs_queried']}")
    lines.append(f"\nOCSP RESULTS:")
    for k, v in output["ocsp_summary"].items():
        lines.append(f"  {k}: {v}")
    lines.append(f"\nCRL RESULTS:")
    for k, v in output["crl_summary"].items():
        lines.append(f"  {k}: {v}")
    
    for url, info in output["crl_snapshots"].items():
        lines.append(f"\nCRL SNAPSHOT: {url}")
        lines.append(f"  Last Update: {info['last_update']}")
        lines.append(f"  Next Update: {info['next_update']}")
        lines.append(f"  Revoked certs in CRL: {info['total_revoked']}")
    
    lines.append(f"\n{'=' * 90}")
    lines.append("PER-CERTIFICATE DETAILS")
    lines.append(f"{'=' * 90}")
    
    # Merge OCSP and CRL results
    crl_by_id = {r["registry_id"]: r for r in crl_results}
    
    for r in ocsp_results:
        rid = r["registry_id"]
        cn = r["subject_cn"]
        lines.append(f"\n  REG#{rid:03d}  {cn}")
        lines.append(f"    Serial:     {r['serial_hex']}")
        lines.append(f"    Expired:    {r.get('is_expired', 'N/A')}")
        lines.append(f"    OCSP:       {r.get('ocsp_status', 'N/A').upper()}"
                     f"  (verify: {r.get('ocsp_verify', 'N/A')})")
        if r.get("ocsp_produced_at"):
            lines.append(f"    Produced:   {r['ocsp_produced_at']}")
        if r.get("ocsp_revocation_time"):
            lines.append(f"    Revoked at: {r['ocsp_revocation_time']}")
            if r.get("ocsp_revocation_reason"):
                lines.append(f"    Reason:     {r['ocsp_revocation_reason']}")
        if r.get("response_der_sha256"):
            lines.append(f"    Resp SHA256:{r['response_der_sha256']}")
        
        crl_r = crl_by_id.get(rid, {})
        lines.append(f"    CRL:        {crl_r.get('crl_status', 'N/A').upper()}")
    
    lines.append(f"\n{'=' * 90}")
    lines.append("END OF REPORT")
    lines.append(f"{'=' * 90}")
    
    report_path = out_dir / "ocsp_crl_report.txt"
    report_path.write_text("\n".join(lines))
    print(f"  Report: {report_path}")
    
    # Save raw OCSP text outputs
    raw_dir = out_dir / "ocsp_raw_text"
    raw_dir.mkdir(exist_ok=True)
    for r in ocsp_results:
        if "raw_stdout" in r and r.get("ocsp_status") not in ["no_ocsp_url", "no_cert_file", "no_issuer_cert"]:
            rid = r["registry_id"]
            cn_safe = "".join(c if c.isalnum() or c in "-_" else "_" for c in r["subject_cn"])[:40]
            txt_path = raw_dir / f"REG{rid:03d}__{cn_safe}__ocsp.txt"
            txt_path.write_text(r["raw_stdout"])
    
    # Manifest
    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 = hashlib.sha256(fpath.read_bytes()).hexdigest()
            manifest_lines.append(f"{h}  {rel}")
    
    manifest_path = out_dir / "MANIFEST.sha256"
    manifest_path.write_text("\n".join(manifest_lines) + "\n")
    print(f"  Manifest: {manifest_path} ({len(manifest_lines)} files)")
    print(f"\n  ➤ OTS-stamp {manifest_path} to anchor this verification snapshot")
    
    # Final summary
    print(f"\n{'=' * 70}")
    print(f"OCSP responses saved: {sum(1 for r in ocsp_results if r.get('response_der_saved'))}")
    print(f"  good:    {output['ocsp_summary']['good']}")
    print(f"  revoked: {output['ocsp_summary']['revoked']}")
    print(f"  unknown: {output['ocsp_summary']['unknown']}")
    print(f"  error:   {output['ocsp_summary']['error']}")
    print(f"  skipped: {output['ocsp_summary']['skipped']}")
    print(f"CRL check: {not_revoked} clean, {revoked} revoked")
    print(f"{'=' * 70}")


if __name__ == "__main__":
    main()