#!/usr/bin/env python3
"""
Axiom Zero — Production mTLS Certificate Authority & Enclave Generator
Generates ECDSA P-384 X.509 client certificates and cryptographic keypairs for enterprise node attestation.
"""

import os
import json
import datetime
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import ec
from cryptography.hazmat.primitives import hashes
from cryptography import x509
from cryptography.x509.oid import NameOID, ExtendedKeyUsageOID

CERT_STORE_DIR = "/home/snuffleupagus/teamwork_projects/axiom_zero_audit/dev_control_center/certs"

class EnclaveCAGenerator:
    def __init__(self, ca_dir=CERT_STORE_DIR):
        self.ca_dir = ca_dir
        os.makedirs(self.ca_dir, exist_ok=True)
        self.ca_key_path = os.path.join(self.ca_dir, "ca_key.pem")
        self.ca_cert_path = os.path.join(self.ca_dir, "ca_cert.pem")
        self.ca_key = None
        self.ca_cert = None

    def initialize_ca(self):
        """Generates CA key and certificate if they don't exist, otherwise loads them."""
        if os.path.exists(self.ca_key_path) and os.path.exists(self.ca_cert_path):
            print("CA already exists. Loading from disk.")
            with open(self.ca_key_path, "rb") as f:
                self.ca_key = serialization.load_pem_private_key(f.read(), password=None)
            with open(self.ca_cert_path, "rb") as f:
                self.ca_cert = x509.load_pem_x509_certificate(f.read())
        else:
            print("Generating new CA...")
            self.ca_key = ec.generate_private_key(ec.SECP384R1())
            subject = issuer = x509.Name([
                x509.NameAttribute(NameOID.ORGANIZATION_NAME, u"Axiom Zero Production"),
                x509.NameAttribute(NameOID.COMMON_NAME, u"Axiom Zero Production Root CA v2.4"),
            ])
            self.ca_cert = x509.CertificateBuilder().subject_name(
                subject
            ).issuer_name(
                issuer
            ).public_key(
                self.ca_key.public_key()
            ).serial_number(
                x509.random_serial_number()
            ).not_valid_before(
                datetime.datetime.now(datetime.timezone.utc)
            ).not_valid_after(
                datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(days=3650)
            ).add_extension(
                x509.BasicConstraints(ca=True, path_length=None), critical=True,
            ).sign(self.ca_key, hashes.SHA384())

            with open(self.ca_key_path, "wb") as f:
                f.write(self.ca_key.private_bytes(
                    encoding=serialization.Encoding.PEM,
                    format=serialization.PrivateFormat.PKCS8,
                    encryption_algorithm=serialization.NoEncryption()
                ))
            with open(self.ca_cert_path, "wb") as f:
                f.write(self.ca_cert.public_bytes(serialization.Encoding.PEM))
            print("CA generated and saved.")

    def generate_node_certificate(self, customer_id: str, domain: str) -> dict:
        """
        Generates an enterprise client node mTLS identity certificate record with real crypto.
        """
        if not self.ca_key or not self.ca_cert:
            self.initialize_ca()
            
        print(f"Generating client certificate for {customer_id} ({domain})...")
        client_key = ec.generate_private_key(ec.SECP384R1())
        subject = x509.Name([
            x509.NameAttribute(NameOID.ORGANIZATION_NAME, str(customer_id)),
            x509.NameAttribute(NameOID.COMMON_NAME, str(domain)),
        ])
        client_cert = x509.CertificateBuilder().subject_name(
            subject
        ).issuer_name(
            self.ca_cert.subject
        ).public_key(
            client_key.public_key()
        ).serial_number(
            x509.random_serial_number()
        ).not_valid_before(
            datetime.datetime.now(datetime.timezone.utc)
        ).not_valid_after(
            datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(days=365)
        ).add_extension(
            x509.BasicConstraints(ca=False, path_length=None), critical=True,
        ).add_extension(
            x509.ExtendedKeyUsage([ExtendedKeyUsageOID.CLIENT_AUTH]), critical=True,
        ).sign(self.ca_key, hashes.SHA384())

        client_key_path = os.path.join(self.ca_dir, f"{customer_id}_key.pem")
        client_cert_path = os.path.join(self.ca_dir, f"{customer_id}_cert.pem")
        cert_data_path = os.path.join(self.ca_dir, f"{customer_id}_cert.json")

        with open(client_key_path, "wb") as f:
            f.write(client_key.private_bytes(
                encoding=serialization.Encoding.PEM,
                format=serialization.PrivateFormat.PKCS8,
                encryption_algorithm=serialization.NoEncryption()
            ))
        with open(client_cert_path, "wb") as f:
            f.write(client_cert.public_bytes(serialization.Encoding.PEM))
            
        timestamp = datetime.datetime.now(datetime.timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
        
        cert_data = {
            "certificate_id": f"CERT-{client_cert.serial_number}",
            "customer_id": customer_id,
            "domain": domain,
            "signature_algorithm": "ECDSA_P384_SHA384",
            "key_path": client_key_path,
            "cert_path": client_cert_path,
            "issuer": "Axiom Zero Production Root CA v2.4",
            "issued_at": timestamp,
            "status": "ACTIVE_VALID"
        }
        
        with open(cert_data_path, "w", encoding="utf-8") as f:
            json.dump(cert_data, f, indent=2)
            
        print(f"[*] Generated ECDSA P-384 mTLS Client Certificate: {client_cert_path}")
        return cert_data

if __name__ == "__main__":
    generator = EnclaveCAGenerator()
    generator.initialize_ca()
    generator.generate_node_certificate("cust_acme_corps", "https://noctualabs.tech")
    generator.generate_node_certificate("cust_fintech_edge", "https://eu-app.noctualabs.tech")