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


class AutoCertGen:
    def __init__(self, name="website", org="organization", locale="place", province="province", country="ZZ", dns="localhost"):
        self.name = name
        self.org = org
        self.locale = locale
        self.province = province
        self.country = country
        self.dns = dns

    def generate_self_signed_cert(self, different_issuer=False):
        private_key = rsa.generate_private_key(
            public_exponent=65537,
            key_size=2048,
        )

        if different_issuer is True:
            subject = x509.Name(
                [
                    x509.NameAttribute(NameOID.COUNTRY_NAME, self.country),
                    x509.NameAttribute(NameOID.STATE_OR_PROVINCE_NAME, self.province),
                    x509.NameAttribute(NameOID.LOCALITY_NAME, self.locale),
                    x509.NameAttribute(NameOID.ORGANIZATION_NAME, self.org),
                    x509.NameAttribute(NameOID.COMMON_NAME, self.name),
                ]
            )
            issuer = x509.Name(
                [
                    x509.NameAttribute(NameOID.COUNTRY_NAME, "RU"),
                    x509.NameAttribute(NameOID.STATE_OR_PROVINCE_NAME, "Красноярск Область"),
                    x509.NameAttribute(NameOID.LOCALITY_NAME, "Красноярск"),
                    x509.NameAttribute(NameOID.ORGANIZATION_NAME, "Amethyst Group"),
                    x509.NameAttribute(NameOID.COMMON_NAME, "Amethyst Untrusted Signing Certificate"),
                ]
            )
        else:
            subject = issuer = x509.Name(
                [
                    x509.NameAttribute(NameOID.COUNTRY_NAME, self.country),
                    x509.NameAttribute(NameOID.STATE_OR_PROVINCE_NAME, self.province),
                    x509.NameAttribute(NameOID.LOCALITY_NAME, self.locale),
                    x509.NameAttribute(NameOID.ORGANIZATION_NAME, self.org),
                    x509.NameAttribute(NameOID.COMMON_NAME, self.name),
                ]
            )

        certificate = (
            x509.CertificateBuilder()
            .subject_name(subject)
            .issuer_name(issuer)
            .public_key(private_key.public_key())
            .serial_number(x509.random_serial_number())
            .not_valid_before(datetime.datetime.utcnow())
            .not_valid_after(
                datetime.datetime.utcnow() + datetime.timedelta(days=365)
            )  # 1 year validity
            .add_extension(
                x509.SubjectAlternativeName([x509.DNSName(self.dns)]), critical=False
            )
            .sign(private_key, hashes.SHA256())
        )

        with open("key.pem", "wb") as f:
            f.write(
                private_key.private_bytes(
                    encoding=serialization.Encoding.PEM,
                    format=serialization.PrivateFormat.TraditionalOpenSSL,
                    encryption_algorithm=serialization.NoEncryption(),
                )
            )

        with open("cert.pem", "wb") as f:
            f.write(certificate.public_bytes(serialization.Encoding.PEM))