#!/usr/bin/python3
"""Verify all IPA/ipacta certificates use the expected algorithm and
validate properly against the CA chain.

Usage:
    ipacta-verify-certs [--alg ML-DSA-44]

Defaults to the CA signing algorithm found in the CA certificate.
"""

import argparse
import subprocess
import sys
from pathlib import Path

from cryptography import x509
from cryptography.x509.oid import PublicKeyAlgorithmOID


# NSSDB that holds CA and subsystem certs
NSSDB_DIR = Path("/etc/pki/pki-tomcat/alias")
NSSDB_PWDFILE = NSSDB_DIR / "pwdfile.txt"

# PEM files outside the NSSDB
PEM_CERTS = {
    "RA agent cert": "/var/lib/ipa/ra-agent.pem",
    "IPA CA bundle": "/etc/ipa/ca.crt",
    "HTTPD cert": "/var/lib/ipa/certs/httpd.crt",
    "KDC cert": "/var/kerberos/krb5kdc/kdc.crt",
}

# Expected NSSDB nicknames
NSSDB_NICKNAMES = [
    "caSigningCert cert-pki-ca",
    "ocspSigningCert cert-pki-ca",
    "subsystemCert cert-pki-ca",
    "auditSigningCert cert-pki-ca",
    "Server-Cert cert-pki-ca",
    # KRA (only if --setup-kra was used)
    "transportCert cert-pki-kra",
    "storageCert cert-pki-kra",
    "auditSigningCert cert-pki-kra",
]

# OID to friendly name
ALG_NAMES = {
    "1.2.840.113549.1.1.1": "RSA",
    "1.2.840.113549.1.1.11": "SHA256withRSA",
    "1.2.840.113549.1.1.12": "SHA384withRSA",
    "1.2.840.113549.1.1.13": "SHA512withRSA",
    "1.2.840.10045.2.1": "EC",
    "1.2.840.10045.4.3.2": "SHA256withECDSA",
    "1.2.840.10045.4.3.3": "SHA384withECDSA",
    "1.2.840.10045.4.3.4": "SHA512withECDSA",
    str(PublicKeyAlgorithmOID.ML_DSA_44
        if hasattr(PublicKeyAlgorithmOID, "ML_DSA_44")
        else "2.16.840.1.101.3.4.3.17"): "ML-DSA-44",
    "2.16.840.1.101.3.4.3.17": "ML-DSA-44",
    "2.16.840.1.101.3.4.3.18": "ML-DSA-65",
    "2.16.840.1.101.3.4.3.19": "ML-DSA-87",
    "2.16.840.1.101.3.4.4.1": "ML-KEM-512",
    "2.16.840.1.101.3.4.4.2": "ML-KEM-768",
    "2.16.840.1.101.3.4.4.3": "ML-KEM-1024",
}

OK = "\033[32mOK\033[0m"
FAIL = "\033[31mFAIL\033[0m"
WARN = "\033[33mWARN\033[0m"
SKIP = "\033[90mSKIP\033[0m"


def alg_name(oid):
    oid_str = oid.dotted_string if hasattr(oid, "dotted_string") else str(oid)
    return ALG_NAMES.get(oid_str, oid_str)


def get_sig_alg(cert):
    return alg_name(cert.signature_algorithm_oid)


def get_key_alg(cert):
    oid = cert.public_key_algorithm_oid
    return alg_name(oid)


def extract_cert_from_nssdb(nickname, nssdb_dir, nssdb_pwdfile):
    """Extract a certificate from NSSDB as PEM."""
    if nssdb_pwdfile.exists():
        pwd_args = ["-f", str(nssdb_pwdfile)]
    else:
        pwd_args = ["-W", ""]

    try:
        result = subprocess.run(
            ["certutil", "-L", "-d", f"sql:{nssdb_dir}",
             "-n", nickname, "-a"] + pwd_args,
            capture_output=True, text=True, check=True,
        )
        return result.stdout
    except subprocess.CalledProcessError:
        return None


def load_pem_cert(pem_data):
    if isinstance(pem_data, str):
        pem_data = pem_data.encode()
    return x509.load_pem_x509_certificate(pem_data)


def verify_signature(cert, ca_cert):
    """Verify cert was signed by ca_cert's key."""
    try:
        ca_pub = ca_cert.public_key()
        # ML-DSA: signature verification uses verify() with no prehash
        ca_pub.verify(
            cert.signature,
            cert.tbs_certificate_bytes,
        )
        return True
    except Exception:
        pass

    # RSA/EC path: need hash algorithm
    try:
        from cryptography.hazmat.primitives.asymmetric import padding
        ca_pub = ca_cert.public_key()
        sig_alg = str(cert.signature_algorithm_oid)
        if "1.2.840.113549" in sig_alg:  # RSA
            from cryptography.hazmat.primitives import hashes
            hash_map = {
                "1.2.840.113549.1.1.11": hashes.SHA256(),
                "1.2.840.113549.1.1.12": hashes.SHA384(),
                "1.2.840.113549.1.1.13": hashes.SHA512(),
            }
            h = hash_map.get(sig_alg, hashes.SHA256())
            ca_pub.verify(
                cert.signature,
                cert.tbs_certificate_bytes,
                padding.PKCS1v15(),
                h,
            )
            return True
        elif "1.2.840.10045" in sig_alg:  # EC
            from cryptography.hazmat.primitives.asymmetric import ec
            from cryptography.hazmat.primitives import hashes
            hash_map = {
                "1.2.840.10045.4.3.2": hashes.SHA256(),
                "1.2.840.10045.4.3.3": hashes.SHA384(),
                "1.2.840.10045.4.3.4": hashes.SHA512(),
            }
            h = hash_map.get(sig_alg, hashes.SHA256())
            ca_pub.verify(
                cert.signature,
                cert.tbs_certificate_bytes,
                ec.ECDSA(h),
            )
            return True
    except Exception:
        pass

    return False


def check_cert(label, cert, expected_alg, ca_cert, results):
    """Check one certificate and record results."""
    sig_alg = get_sig_alg(cert)
    key_alg = get_key_alg(cert)
    subject = cert.subject.rfc4514_string()
    serial = format(cert.serial_number, 'x')

    is_self_signed = cert.subject == cert.issuer
    verify_against = cert if is_self_signed else ca_cert

    # Check signature algorithm matches expected
    alg_ok = (expected_alg is None or
              sig_alg.upper() == expected_alg.upper())

    # Verify signature
    if verify_against:
        sig_ok = verify_signature(cert, verify_against)
    else:
        sig_ok = None

    status = OK if (alg_ok and sig_ok) else FAIL
    if sig_ok is None and alg_ok:
        status = WARN

    results.append({
        "label": label,
        "status": "ok" if (alg_ok and sig_ok) else "fail",
        "sig_alg": sig_alg,
        "key_alg": key_alg,
        "alg_ok": alg_ok,
        "sig_ok": sig_ok,
    })

    print(f"  {status}  {label}")
    print(f"         Subject:   {subject}")
    print(f"         Serial:    {serial}")
    print(f"         Key alg:   {key_alg}")
    print(f"         Sig alg:   {sig_alg}", end="")
    if not alg_ok:
        print(f"  (expected {expected_alg})", end="")
    print()
    print("         Signature: ", end="")
    if sig_ok is True:
        print("valid")
    elif sig_ok is False:
        print("INVALID")
    else:
        print("not checked (no CA cert)")
    print()


def main():
    parser = argparse.ArgumentParser(
        description="Verify IPA certificate algorithms and signatures"
    )
    parser.add_argument(
        "--alg",
        help="Expected signing algorithm (default: auto-detect from CA cert)",
    )
    parser.add_argument(
        "--nssdb", default=str(NSSDB_DIR),
        help="NSSDB directory (default: %(default)s)",
    )
    args = parser.parse_args()

    nssdb_dir = Path(args.nssdb)
    nssdb_pwdfile = nssdb_dir / "pwdfile.txt"

    results = []
    ca_cert = None

    # Load CA certificate first
    print("=" * 60)
    print("IPA Certificate Verification")
    print("=" * 60)
    print()

    ca_pem = extract_cert_from_nssdb("caSigningCert cert-pki-ca",
                                     nssdb_dir, nssdb_pwdfile)
    if ca_pem:
        ca_cert = load_pem_cert(ca_pem)
        expected_alg = args.alg or get_sig_alg(ca_cert)
        print(f"CA signing algorithm: {get_sig_alg(ca_cert)}")
        print(f"CA key algorithm:     {get_key_alg(ca_cert)}")
        print(f"Expected algorithm:   {expected_alg}")
    else:
        expected_alg = args.alg
        print("WARNING: Could not extract CA cert from NSSDB")
        if expected_alg:
            print(f"Expected algorithm:   {expected_alg}")
        else:
            print("No --alg specified and no CA cert found, "
                  "skipping algorithm checks")

    print()
    print("-" * 60)
    print("NSSDB Certificates")
    print("-" * 60)
    print()

    for nickname in NSSDB_NICKNAMES:
        pem = extract_cert_from_nssdb(nickname, nssdb_dir, nssdb_pwdfile)
        if pem is None:
            print(f"  {SKIP}  {nickname}  (not found)")
            print()
            continue
        cert = load_pem_cert(pem)
        check_cert(nickname, cert, expected_alg, ca_cert, results)

    print("-" * 60)
    print("PEM File Certificates")
    print("-" * 60)
    print()

    for label, path in PEM_CERTS.items():
        p = Path(path)
        if not p.exists():
            print(f"  {SKIP}  {label}  ({path} not found)")
            print()
            continue
        try:
            pem_data = p.read_bytes()
            # ca.crt may contain multiple certs (bundle) — check each
            certs_in_file = []
            remaining = pem_data
            while b"-----BEGIN CERTIFICATE-----" in remaining:
                cert = x509.load_pem_x509_certificate(remaining)
                certs_in_file.append(cert)
                # Find end of this cert and continue
                end = remaining.find(b"-----END CERTIFICATE-----")
                remaining = remaining[end + len(b"-----END CERTIFICATE-----"):]

            if not certs_in_file:
                certs_in_file = [load_pem_cert(pem_data)]

            for i, cert in enumerate(certs_in_file):
                suffix = f" [{i+1}]" if len(certs_in_file) > 1 else ""
                check_cert(f"{label}{suffix} ({path})",
                           cert, expected_alg, ca_cert, results)
        except Exception as e:
            print(f"  {FAIL}  {label}  ({path}: {e})")
            print()
            results.append({"label": label, "status": "fail"})

    # Summary
    print("=" * 60)
    print("Summary")
    print("=" * 60)

    total = len(results)
    passed = sum(1 for r in results if r["status"] == "ok")
    failed = sum(1 for r in results if r["status"] == "fail")

    print(f"  Checked: {total}")
    print(f"  Passed:  {passed}")
    print(f"  Failed:  {failed}")
    print()

    if failed:
        print(f"  {FAIL}  {failed} certificate(s) failed verification")
        return 1
    else:
        print(f"  {OK}  All certificates verified successfully")
        return 0


if __name__ == "__main__":
    sys.exit(main())
