#!/usr/bin/env python3
"""
Generate a small local CA plus leaf certificates for Phase 2 HTTPS testing.
"""

from __future__ import annotations

import argparse
import ipaddress
from datetime import datetime, timedelta, timezone
from pathlib import Path

from cryptography import x509
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.x509.oid import ExtendedKeyUsageOID, NameOID


def build_name(common_name: str) -> x509.Name:
    return x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, common_name)])


def new_private_key() -> rsa.RSAPrivateKey:
    return rsa.generate_private_key(public_exponent=65537, key_size=2048)


def write_private_key(path: Path, key: rsa.RSAPrivateKey) -> None:
    path.write_bytes(
        key.private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.TraditionalOpenSSL,
            encryption_algorithm=serialization.NoEncryption(),
        )
    )


def write_cert(path: Path, cert: x509.Certificate) -> None:
    path.write_bytes(cert.public_bytes(serialization.Encoding.PEM))


def parse_sans(names: list[str]) -> list[x509.GeneralName]:
    sans: list[x509.GeneralName] = []
    seen = set()
    for value in names:
        if value in seen:
            continue
        seen.add(value)
        try:
            sans.append(x509.IPAddress(ipaddress.ip_address(value)))
        except ValueError:
            sans.append(x509.DNSName(value))
    return sans


def issue_ca(common_name: str, valid_days: int) -> tuple[rsa.RSAPrivateKey, x509.Certificate]:
    now = datetime.now(timezone.utc)
    key = new_private_key()
    subject = issuer = build_name(common_name)
    cert = (
        x509.CertificateBuilder()
        .subject_name(subject)
        .issuer_name(issuer)
        .public_key(key.public_key())
        .serial_number(x509.random_serial_number())
        .not_valid_before(now - timedelta(minutes=5))
        .not_valid_after(now + timedelta(days=valid_days))
        .add_extension(x509.BasicConstraints(ca=True, path_length=None), critical=True)
        .add_extension(x509.SubjectKeyIdentifier.from_public_key(key.public_key()), critical=False)
        .add_extension(x509.AuthorityKeyIdentifier.from_issuer_public_key(key.public_key()), critical=False)
        .add_extension(x509.KeyUsage(digital_signature=True, key_encipherment=False, key_cert_sign=True, crl_sign=True, content_commitment=False, data_encipherment=False, key_agreement=False, encipher_only=False, decipher_only=False), critical=True)
        .sign(key, hashes.SHA256())
    )
    return key, cert


def issue_leaf(
    ca_key: rsa.RSAPrivateKey,
    ca_cert: x509.Certificate,
    common_name: str,
    san_values: list[str],
    valid_days: int,
) -> tuple[rsa.RSAPrivateKey, x509.Certificate]:
    now = datetime.now(timezone.utc)
    key = new_private_key()
    cert = (
        x509.CertificateBuilder()
        .subject_name(build_name(common_name))
        .issuer_name(ca_cert.subject)
        .public_key(key.public_key())
        .serial_number(x509.random_serial_number())
        .not_valid_before(now - timedelta(minutes=5))
        .not_valid_after(now + timedelta(days=valid_days))
        .add_extension(x509.BasicConstraints(ca=False, path_length=None), critical=True)
        .add_extension(x509.SubjectAlternativeName(parse_sans(san_values)), critical=False)
        .add_extension(x509.SubjectKeyIdentifier.from_public_key(key.public_key()), critical=False)
        .add_extension(x509.AuthorityKeyIdentifier.from_issuer_public_key(ca_key.public_key()), critical=False)
        .add_extension(x509.ExtendedKeyUsage([ExtendedKeyUsageOID.SERVER_AUTH]), critical=False)
        .add_extension(x509.KeyUsage(digital_signature=True, key_encipherment=True, key_cert_sign=False, crl_sign=False, content_commitment=False, data_encipherment=False, key_agreement=False, encipher_only=False, decipher_only=False), critical=True)
        .sign(ca_key, hashes.SHA256())
    )
    return key, cert


def emit_leaf_bundle(
    out_dir: Path,
    leaf_name: str,
    ca_key: rsa.RSAPrivateKey,
    ca_cert: x509.Certificate,
    san_values: list[str],
    valid_days: int,
) -> None:
    key, cert = issue_leaf(ca_key, ca_cert, leaf_name, san_values, valid_days)
    write_private_key(out_dir / f"{leaf_name}.key", key)
    write_cert(out_dir / f"{leaf_name}.crt", cert)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Generate local CA and Phase 2 service certificates")
    parser.add_argument("--out-dir", default="tls/phase2")
    parser.add_argument("--valid-days", type=int, default=30)
    parser.add_argument("--ca-common-name", default="ChromeCard Phase2 Local CA")
    parser.add_argument(
        "--proxy-san",
        action="append",
        default=[],
        help="Extra SAN for k_proxy certificate; may be repeated",
    )
    parser.add_argument(
        "--server-san",
        action="append",
        default=[],
        help="Extra SAN for k_server certificate; may be repeated",
    )
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    ca_key, ca_cert = issue_ca(args.ca_common_name, args.valid_days)
    write_private_key(out_dir / "ca.key", ca_key)
    write_cert(out_dir / "ca.crt", ca_cert)

    proxy_sans = ["localhost", "127.0.0.1", "k_proxy", *args.proxy_san]
    server_sans = ["localhost", "127.0.0.1", "k_server", *args.server_san]

    emit_leaf_bundle(out_dir, "k_proxy", ca_key, ca_cert, proxy_sans, args.valid_days)
    emit_leaf_bundle(out_dir, "k_server", ca_key, ca_cert, server_sans, args.valid_days)

    print(f"Generated CA and leaf certificates in {out_dir}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())