#!/usr/bin/env python3
"""Creates the batch attestation key and certificate and emits them as C source.

The key and certificate are kept in --key-dir and reused, so rebuilding keeps
the same attestation identity. Delete that directory to start over.
"""

import argparse
import datetime
import pathlib
import uuid

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

AAGUID = uuid.UUID("225e7977-54f1-4d15-8e35-295b11658d5c")
OID_FIDO_AAGUID = x509.ObjectIdentifier("1.3.6.1.4.1.45724.1.1.4")


def create(key_path: pathlib.Path, cert_path: pathlib.Path) -> None:
    key = ec.generate_private_key(ec.SECP256R1())
    # Subject fields required for packed attestation certificates (WebAuthn 8.2.1).
    name = x509.Name(
        [
            x509.NameAttribute(NameOID.COUNTRY_NAME, "US"),
            x509.NameAttribute(NameOID.ORGANIZATION_NAME, "Skey DIY"),
            x509.NameAttribute(NameOID.ORGANIZATIONAL_UNIT_NAME, "Authenticator Attestation"),
            x509.NameAttribute(NameOID.COMMON_NAME, "Skey Pico Attestation"),
        ]
    )
    now = datetime.datetime.now(datetime.timezone.utc)
    cert = (
        x509.CertificateBuilder()
        .subject_name(name)
        .issuer_name(name)
        .public_key(key.public_key())
        .serial_number(x509.random_serial_number())
        .not_valid_before(now - datetime.timedelta(days=1))
        .not_valid_after(datetime.datetime(2049, 12, 31, tzinfo=datetime.timezone.utc))
        .add_extension(x509.BasicConstraints(ca=False, path_length=None), critical=True)
        .add_extension(x509.UnrecognizedExtension(OID_FIDO_AAGUID, b"\x04\x10" + AAGUID.bytes), critical=False)
        .sign(key, hashes.SHA256())
    )
    key_path.parent.mkdir(parents=True, exist_ok=True)
    key_path.write_bytes(
        key.private_bytes(
            serialization.Encoding.PEM,
            serialization.PrivateFormat.PKCS8,
            serialization.NoEncryption(),
        )
    )
    key_path.chmod(0o600)
    cert_path.write_bytes(cert.public_bytes(serialization.Encoding.DER))


def cert_aaguid(cert_der: bytes) -> bytes:
    cert = x509.load_der_x509_certificate(cert_der)
    try:
        return cert.extensions.get_extension_for_oid(OID_FIDO_AAGUID).value.value[2:]
    except x509.ExtensionNotFound:
        return b""


def c_bytes(data: bytes) -> str:
    rows = [", ".join(f"0x{b:02x}" for b in data[i : i + 12]) for i in range(0, len(data), 12)]
    return "{\n    " + ",\n    ".join(rows) + ",\n}"


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--key-dir", type=pathlib.Path, required=True)
    parser.add_argument("--out", type=pathlib.Path, required=True)
    args = parser.parse_args()

    key_path = args.key_dir / "attestation_key.pem"
    cert_path = args.key_dir / "attestation_cert.der"
    if not key_path.exists() or not cert_path.exists() or cert_aaguid(cert_path.read_bytes()) != AAGUID.bytes:
        create(key_path, cert_path)

    key = serialization.load_pem_private_key(key_path.read_bytes(), password=None)
    private_value = key.private_numbers().private_value.to_bytes(32, "big")
    cert_der = cert_path.read_bytes()

    args.out.parent.mkdir(parents=True, exist_ok=True)
    args.out.write_text(
        f"""// Generated by tools/gen_attestation.py. Do not edit.
#include "attestation.h"

const uint8_t skey_aaguid[16] = {c_bytes(AAGUID.bytes)};

const uint8_t attestation_key[32] = {c_bytes(private_value)};

const uint8_t attestation_cert[] = {c_bytes(cert_der)};

const size_t attestation_cert_len = sizeof(attestation_cert);
"""
    )


if __name__ == "__main__":
    main()