"""
PROJECT:     ReactOS tools
LICENSE:     MIT (https://spdx.org/licenses/MIT)
PURPOSE:     Script to update caroots.inf with the latest CA root certificates from Mozilla NSS
COPYRIGHT:   Copyright 2025 Mark Jansen <mark.jansen@reactos.org>
"""

from enum import Enum
from pathlib import Path
import argparse
import urllib.request
from datetime import datetime
from dataclasses import dataclass

# Additional sources:
# https://hg.mozilla.org/releases/mozilla-release/raw-file/default/security/nss/lib/ckfw/builtins/certdata.txt
# https://hg.mozilla.org/projects/nss/raw-file/tip/lib/ckfw/builtins/certdata.txt


REPO_ROOT = Path(__file__).resolve().parent.parent.parent
CAROOTS_INF = REPO_ROOT / "boot" / "bootdata" / "caroots.inf"
DEFAULT_DOWNLOAD_URL = "https://hg.mozilla.org/mozilla-central/raw-file/default/security/nss/lib/ckfw/builtins/certdata.txt"
CAROOTS_HEADER = """; Auto-generated caroots.inf
; Do not edit manually.
; Source: {source_url}
; Generated on: {date}
; Number of certificates: {cert_count}
; Generated by sdk/tools/update_caroots.py

[Version]
Signature = "$Windows NT$"

[AddReg]

"""

CERTIFICATE_HEADER = """; "{cert_name}" ({cert_size} bytes)
HKLM,"SOFTWARE\\Microsoft\\SystemCertificates\\AuthRoot\\Certificates\\{fingerprint_sha1}","Blob",0x00000001,\\
    0x20,0x00,0x00,0x00,\\
    0x01,0x00,0x00,0x00,\\
    {cert_size_hex},\\
{data}
"""


@dataclass
class Certificate:
    name: str
    fingerprint_sha1: str
    data: bytes


def format_certificate(cert: Certificate) -> str:
    # Format the data as comma-separated hex bytes, 16 bytes per line
    hex_bytes = [f"0x{b:02X}" for b in cert.data]
    lines = []
    for i in range(0, len(hex_bytes), 16):
        line = ",".join(hex_bytes[i : i + 16])
        if i + 16 < len(hex_bytes):
            line += ",\\"
        lines.append("    " + line)

    cert_size = len(cert.data)
    cert_size_hex = ",".join(f"0x{b:02X}" for b in cert_size.to_bytes(4, byteorder="little"))

    return CERTIFICATE_HEADER.format(
        cert_name=cert.name,
        fingerprint_sha1=cert.fingerprint_sha1,
        cert_size=cert_size,
        cert_size_hex=cert_size_hex,
        data="\n".join(lines),
    )


class CurrentBlock(Enum):
    Begin = 0
    Certificate = 1
    Trust = 2


def parse_certificates2(certdata: str) -> list[Certificate]:
    cert = None
    result = []
    temp_value = None
    block = CurrentBlock.Begin

    for line in certdata.splitlines():
        line = line.strip()
        if not line or line.startswith("#"):
            continue

        if line == "CKA_CLASS CK_OBJECT_CLASS CKO_NSS_TRUST":
            block = CurrentBlock.Trust
            continue
        elif line == "CKA_CLASS CK_OBJECT_CLASS CKO_CERTIFICATE":
            block = CurrentBlock.Certificate
            continue
        elif line.startswith("CKA_LABEL"):
            _, name, _ = line.split('"')
            if block == CurrentBlock.Certificate:
                cert = Certificate(name=name, fingerprint_sha1="", data=b"")
            elif block == CurrentBlock.Trust:
                assert cert is not None, line
                assert name == cert.name
            else:
                pass
        elif line == "CKA_VALUE MULTILINE_OCTAL":
            assert cert is not None
            temp_value = []
            assert cert.data == b""
            assert cert.fingerprint_sha1 == ""
        elif line == "CKA_CERT_SHA1_HASH MULTILINE_OCTAL":
            assert cert is not None
            temp_value = []
            assert cert.data != b""
            assert cert.fingerprint_sha1 == ""
        elif temp_value is not None:
            assert cert is not None
            if line == "END":
                if cert.data == b"":
                    cert.data = bytes(temp_value)
                else:
                    assert cert.fingerprint_sha1 == ""
                    cert.fingerprint_sha1 = "".join(f"{b:02X}" for b in temp_value)
                temp_value = None
            else:
                for number in line.split("\\"):
                    if not number:
                        continue
                    temp_value.append(int(number, 8))

        elif line.startswith("CKA_TRUST_SERVER_AUTH"):
            assert cert is not None
            if "CKT_NSS_TRUSTED_DELEGATOR" in line:
                print(f"Trusted cert: {cert.name} ({len(cert.data)} bytes)")
                result.append(cert)
            cert = None
    return result


def main():
    parser = argparse.ArgumentParser(
        description="Update the caroots.inf file with the latest CA root certificates."
    )
    parser.add_argument(
        "--url",
        type=str,
        default=DEFAULT_DOWNLOAD_URL,
        help="URL to download the certdata.txt file from.",
    )
    args = parser.parse_args()

    print(f"Downloading certdata.txt from {args.url}...")
    with urllib.request.urlopen(args.url) as response:
        certdata = response.read().decode("utf-8")

    print("Parsing certdata.txt...")
    certificates = parse_certificates2(certdata)
    certificates.sort(key=lambda c: c.name.lower())
    print(f"Found {len(certificates)} trusted root certificates.")
    print(f"Updating {CAROOTS_INF}...")
    with CAROOTS_INF.open("w", encoding="utf-8", newline="\r\n") as f:
        f.write(
            CAROOTS_HEADER.format(
                source_url=args.url,
                date=datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
                cert_count=len(certificates),
            )
        )
        for cert in certificates:
            f.write(format_certificate(cert))
            f.write("\n")

    print("Update complete.")


if __name__ == "__main__":
    main()