import datetime
import os
import typing
import cryptography

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

COUNTRY_NAME = "CH"
STATE_OR_PROVINCE_NAME = "Zurich"
LOCALITY_NAME = "Zurich"
ORGANIZATION_NAME = "MaybeBug"

CA_CERTIFICATE_PATH = "cacert.pem"


def create_ca() -> typing.Tuple[rsa.RSAPrivateKey, x509.Certificate]:

    private_key: rsa.RSAPrivateKey = rsa.generate_private_key(
        public_exponent=65537,
        key_size=4096,
    )

    ca_cert = (
        x509.CertificateBuilder()
        .subject_name(
            x509.Name(
                [
                    # Provide various details about who we are.
                    x509.NameAttribute(NameOID.COUNTRY_NAME, COUNTRY_NAME),
                    x509.NameAttribute(
                        NameOID.STATE_OR_PROVINCE_NAME, STATE_OR_PROVINCE_NAME
                    ),
                    x509.NameAttribute(NameOID.LOCALITY_NAME, LOCALITY_NAME),
                    x509.NameAttribute(NameOID.ORGANIZATION_NAME, ORGANIZATION_NAME),
                    x509.NameAttribute(NameOID.COMMON_NAME, "CA"),
                ]
            )
        )
        .issuer_name(
            x509.Name(
                [
                    # Provide various details about who we are.
                    x509.NameAttribute(NameOID.COUNTRY_NAME, COUNTRY_NAME),
                    x509.NameAttribute(
                        NameOID.STATE_OR_PROVINCE_NAME, STATE_OR_PROVINCE_NAME
                    ),
                    x509.NameAttribute(NameOID.LOCALITY_NAME, LOCALITY_NAME),
                    x509.NameAttribute(NameOID.ORGANIZATION_NAME, ORGANIZATION_NAME),
                    x509.NameAttribute(NameOID.COMMON_NAME, "CA"),
                ]
            )
        )
        .public_key(private_key.public_key())
        .serial_number(1)
        .not_valid_before(datetime.datetime.utcnow())
        .not_valid_after(
            # Our certificate will be valid for 10 days
            datetime.datetime.utcnow()
            + datetime.timedelta(days=3650)
        )
        .add_extension(x509.BasicConstraints(ca=True, path_length=None), critical=True)
        .add_extension(
            x509.KeyUsage(
                digital_signature=False,
                content_commitment=False,
                key_encipherment=False,
                data_encipherment=False,
                key_agreement=False,
                key_cert_sign=True,
                crl_sign=True,
                encipher_only=False,
                decipher_only=False,
            ),
            critical=False,
        )
        .add_extension(
            x509.AuthorityKeyIdentifier.from_issuer_public_key(
                private_key.public_key()
            ),
            critical=False,
        )
        .sign(private_key, hashes.SHA256())
    )

    with open(CA_CERTIFICATE_PATH, "wb") as f:
        f.write(ca_cert.public_bytes(serialization.Encoding.PEM))

    return private_key, ca_cert


def create_truststore_keytool(truststore_path: str, truststore_password: str):
    command = f'keytool -keystore {truststore_path} -alias CARoot -import -file {CA_CERTIFICATE_PATH} -storepass "{truststore_password}" -noprompt -storetype PKCS12'
    os.system(command)


def create_truststore_cryptography(
    ca_certificate: x509.Certificate, truststore_path: str, truststore_password: str
):

    pcks12_ca_certificate = pkcs12.PKCS12Certificate(ca_certificate, "CAroot".encode())

    p12 = pkcs12.serialize_key_and_certificates(
        name="CARoot".encode(),
        key=None,
        cert=None,
        cas=[pcks12_ca_certificate],
        encryption_algorithm=serialization.BestAvailableEncryption(
            truststore_password.encode()
        ),
    )
    with open(truststore_path, "wb") as f:
        f.write(p12)


def verify_truststore(truststore_path, truststore_password: str):
    command = (
        f"keytool -list -keystore {truststore_path} -storepass '{truststore_password}'"
    )
    os.system(command)

    with open(truststore_path, "rb") as f:
        data = f.read()
        (
            truststore_private_key,
            truststore_certificate,
            truststore_additional_certificates,
        ) = pkcs12.load_key_and_certificates(
            data, password=truststore_password.encode()
        )

        print("truststore_private_key", truststore_private_key)
        print("truststore_certificate", truststore_certificate)
        print("truststore_additional_certificates", truststore_additional_certificates)


password = "a_very_secure_password"
ca_key, ca_certificate = create_ca()

print("Cryptography version:", cryptography.__version__)
print("Create truststore using keytool:")
create_truststore_keytool("keytool.truststore.p12", password)
verify_truststore("keytool.truststore.p12", password)
print()
print("Create truststore using cryptography:")
create_truststore_cryptography(ca_certificate, "cryptography.truststore.p12", password)
verify_truststore("cryptography.truststore.p12", password)

Cryptography version: 37.0.0.dev1
Create truststore using keytool:
Certificate was added to keystore
Keystore type: PKCS12
Keystore provider: SUN

Your keystore contains 1 entry

caroot, Apr 13, 2022, trustedCertEntry, 
Certificate fingerprint (SHA-256): BF:E5:79:AB:98:AF:DA:BF:5B:A1:75:A7:27:2E:BE:BC:9E:16:8E:A7:DD:D3:10:3F:D2:41:07:AE:27:5A:2B:2A
truststore_private_key None
truststore_certificate None
truststore_additional_certificates [<Certificate(subject=<Name(C=CH,ST=Zurich,L=Zurich,O=MaybeBug,CN=CA)>, ...)>]

Create truststore using cryptography:
Keystore type: PKCS12
Keystore provider: SUN

Your keystore contains 0 entries

truststore_private_key None
truststore_certificate None
truststore_additional_certificates [<Certificate(subject=<Name(C=CH,ST=Zurich,L=Zurich,O=MaybeBug,CN=CA)>, ...)>]