import datetime
from cryptography.x509 import Name, NameAttribute, CertificateBuilder, BasicConstraints, random_serial_number
from cryptography.x509.oid import NameOID
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import rsa
import pathlib

def create_self_signed_certificate(
        file_path: str,
        subject_name: str,
        organization_name: str,
        country_name: str,
        valid_days: int = 365,
        key_size: int = 2048
):
    """
    Create a self-signed certificate. It saves the certificate and private key
    in PEM format to the specified file path. Three files are created:

    - <file_path>.crt : The public certificate
    - <file_path>.key : The private key
    - <file_path>.pem : The certificate and private key combined in one file

    :param file_path: Base file path to save the certificate and key.
    :param subject_name: Common Name (CN) for the certificate.
    :param organization_name: Organization Name (O) for the certificate.
    :param country_name: Country Name (C) for the certificate.
    :param valid_days: Number of days the certificate is valid for.
    :param key_size: Size of the RSA key.

    Use the following command to replace any credentials already defined for the 
    App Registration with that certificate:
    
    az ad app credential reset --id <CLIENT_ID> --cert @<file_path>.crt

    Use --append to add the certificate without removing existing credentials.

    > Note: Do not upload the private key file (.key) nor the combined PEM file (.pem).
    """
    key = rsa.generate_private_key(public_exponent=65537, key_size=key_size)

    subject = issuer = Name([
        NameAttribute(NameOID.COUNTRY_NAME, country_name),
        NameAttribute(NameOID.ORGANIZATION_NAME, organization_name),
        NameAttribute(NameOID.COMMON_NAME, subject_name),
    ])

    cert = (
        CertificateBuilder()
        .subject_name(subject)
        .issuer_name(issuer)
        .public_key(key.public_key())
        .serial_number(random_serial_number())
        .not_valid_before(datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(days=1))
        .not_valid_after(datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(days=valid_days))
        .add_extension(BasicConstraints(ca=True, path_length=None), critical=True)
        .sign(key, hashes.SHA256())
    )

    crt_path = pathlib.Path(file_path).with_suffix('.crt')
    key_path = pathlib.Path(file_path).with_suffix('.key')
    pem_path = pathlib.Path(file_path).with_suffix('.pem')
    
    public_key = cert.public_bytes(serialization.Encoding.PEM)
    private_key = key.private_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PrivateFormat.TraditionalOpenSSL,
        encryption_algorithm=serialization.NoEncryption(),
    )
    
    # Write the certificate
    with open(crt_path, "wb") as f:
        f.write(public_key)

    # Write the private key
    with open(key_path, "wb") as f:
        f.write(private_key)
    
    with open(pem_path, "wb") as f:
        f.write(public_key)
        f.write(private_key)

    print(f"Certificate created and saved.")
    print(f" - Certificate: {crt_path}")
    print(f" - Private Key: {key_path}")
    print(f" - PEM File: {pem_path}")