# Source code for pyzipper.zipfile_aes

    import struct

    from Cryptodome.Protocol.KDF import PBKDF2
    from Cryptodome.Cipher import AES
    from Cryptodome.Hash import HMAC
    from Cryptodome.Hash.SHA1 import SHA1Hash
    from Cryptodome.Util import Counter
    from Cryptodome import Random

    from .zipfile import (
        ZIP_BZIP2,
        ZIP_LZMA,
        BadZipFile,
        BaseZipDecrypter,
        ZipFile,
        ZipInfo,
        ZipExtFile,
    )

    WZ_AES = 'WZ_AES'
    WZ_AES_COMPRESS_TYPE = 99
    WZ_AES_V1 = 0x0001
    WZ_AES_V2 = 0x0002
    WZ_AES_VENDOR_ID = b'AE'

    EXTRA_WZ_AES = 0x9901

    WZ_SALT_LENGTHS = {
        1: 8,   # 128 bit
        2: 12,  # 192 bit
        3: 16,  # 256 bit
    }
    WZ_KEY_LENGTHS = {
        1: 16,  # 128 bit
        2: 24,  # 192 bit
        3: 32,  # 256 bit
    }

    class AESZipDecrypter(BaseZipDecrypter):

        hmac_size = 10

        def __init__(self, zinfo, pwd, encryption_header):
            self.filename = zinfo.filename

            key_length = WZ_KEY_LENGTHS[zinfo.wz_aes_strength]
            salt_length = WZ_SALT_LENGTHS[zinfo.wz_aes_strength]

            salt = struct.unpack(
                "<{}s".format(salt_length),
                encryption_header[:salt_length]
            )[0]
            pwd_verify_length = 2
            pwd_verify = encryption_header[salt_length:]
            dkLen = 2*key_length + pwd_verify_length
            keymaterial = PBKDF2(pwd, salt, count=1000, dkLen=dkLen)

            encpwdverify = keymaterial[2*key_length:]
            if encpwdverify != pwd_verify:
                raise RuntimeError("Bad password for file %r" % zinfo.filename)

            enckey = keymaterial[:key_length]
            self.decypter = AES.new(
                enckey,
                AES.MODE_CTR,
                counter=Counter.new(nbits=128, little_endian=True)
            )
            encmac_key = keymaterial[key_length:2*key_length]
            self.hmac = HMAC.new(encmac_key, digestmod=SHA1Hash())

        @staticmethod
        def encryption_header_length(zinfo):
            # salt_length + pwd_verify_length
            salt_length = WZ_SALT_LENGTHS[zinfo.wz_aes_strength]
            return salt_length + 2

        def decrypt(self, data):
            self.hmac.update(data)
            return self.decypter.decrypt(data)

        def check_hmac(self, hmac_check):
            if self.hmac.digest()[:10] != hmac_check:
                raise BadZipFile("Bad HMAC check for file %r" % self.filename)

    class BaseZipEncrypter:

        def update_zipinfo(self, zipinfo):
            raise NotImplementedError(
                'BaseZipEncrypter implementations must implement `update_zipinfo`.'
            )

        def encrypt(self, data):
            raise NotImplementedError(
                'BaseZipEncrypter implementations must implement `encrypt`.'
            )

        def encryption_header(self):
            raise NotImplementedError(
                'BaseZipEncrypter implementations must implement '
                '`encryption_header`.'
            )

        def flush(self):
            return b''

    class AESZipEncrypter(BaseZipEncrypter):

        hmac_size = 10

        def __init__(self, pwd, nbits=256, force_wz_aes_version=None):
            if not pwd:
                raise RuntimeError(
                    '%s encryption requires a password.' % WZ_AES
                )

            if nbits not in (128, 192, 256):
                raise RuntimeError(
                    "`nbits` must be one of 128, 192, 256. Got '%s'" % nbits
                )

            self.force_wz_aes_version = force_wz_aes_version
            salt_lengths = {
                128: 8,
                192: 12,
                256: 16,
            }
            self.salt_length = salt_lengths[nbits]
            key_lengths = {
                128: 16,
                192: 24,
                256: 32,
            }
            key_length = key_lengths[nbits]
            aes_strengths = {
                128: 1,
                192: 2,
                256: 3,
            }
            self.aes_strength = aes_strengths[nbits]

            self.salt = Random.new().read(self.salt_length)
            pwd_verify_length = 2
            dkLen = 2 * key_length + pwd_verify_length
            keymaterial = PBKDF2(pwd, self.salt, count=1000, dkLen=dkLen)

            self.encpwdverify = keymaterial[2*key_length:]

            enckey = keymaterial[:key_length]
            self.encrypter = AES.new(
                enckey,
                AES.MODE_CTR,
                counter=Counter.new(nbits=128, little_endian=True)
            )
            encmac_key = keymaterial[key_length:2*key_length]
            self.hmac = HMAC.new(encmac_key, digestmod=SHA1Hash())

        def update_zipinfo(self, zipinfo):
            zipinfo.wz_aes_vendor_id = WZ_AES_VENDOR_ID
            zipinfo.wz_aes_strength = self.aes_strength
            if self.force_wz_aes_version is not None:
                zipinfo.wz_aes_version = self.force_wz_aes_version

        def encryption_header(self):
            return self.salt + self.encpwdverify

        def encrypt(self, data):
            data = self.encrypter.encrypt(data)
            self.hmac.update(data)
            return data

        def flush(self):
            return struct.pack('<%ds' % self.hmac_size, self.hmac.digest()[:10])

    class AESZipInfo(ZipInfo):
        """Class with attributes describing each file in the ZIP archive."""

        # __slots__ on subclasses only need to contain the additional slots.
        __slots__ = (
            'wz_aes_version',
            'wz_aes_vendor_id',
            'wz_aes_strength',
            # 'wz_aes_actual_compression_type',
        )

        def __init__(self, *args, **kwargs):
            super().__init__(*args, **kwargs)
            self.wz_aes_version = None
            self.wz_aes_vendor_id = None
            self.wz_aes_strength = None

        def decode_extra_wz_aes(self, ln, extra):
            if ln == 7:
                counts = struct.unpack("<H2sBH", extra[4: ln+4])
            else:
                raise BadZipFile(
                    "Corrupt extra field %04x (size=%d)" % (EXTRA_WZ_AES, ln))

            self.wz_aes_version = counts[0]
            self.wz_aes_vendor_id = counts[1]
            # 0x01  128-bit encryption key
            # 0x02  192-bit encryption key
            # 0x03  256-bit encryption key
            self.wz_aes_strength = counts[2]

            # the compression method is the one that would otherwise have been
            # stored in the local and central headers for the file. For example, if
            # the file is imploded, this field will contain the compression code 6.
            # This is needed because a compression method of 99 is used to indicate
            # the presence of an AES-encrypted file
            self.compress_type = counts[3]

            # self.wz_aes_actual_compression_type = counts[3]

        def get_extra_decoders(self):
            extra_decoders = super().get_extra_decoders()
            extra_decoders[EXTRA_WZ_AES] = self.decode_extra_wz_aes
            return extra_decoders

        def encode_extra(self, crc, compress_type):
            wz_aes_extra = b''
            if self.wz_aes_vendor_id is not None:
                compress_type = WZ_AES_COMPRESS_TYPE
                aes_version = self.wz_aes_version
                if aes_version is None:
                    if self.file_size < 20 | self.compress_type == ZIP_BZIP2:
                        # The only difference between version 1 and 2 is the
                        # handling of the CRC values. For version 2 the CRC value
                        # is not used and must be set to 0.
                        # For small files, the CRC files can leak the contents of
                        # the encrypted data.
                        # For bzip2, the compression already has integrity checks
                        # so CRC is not required.
                        aes_version = WZ_AES_V2
                    else:
                        aes_version = WZ_AES_V1

                if aes_version == WZ_AES_V2:
                    crc = 0

                wz_aes_extra = struct.pack(
                    "<3H2sBH",
                    EXTRA_WZ_AES,
                    7,  # extra block body length: H2sBH
                    aes_version,
                    self.wz_aes_vendor_id,
                    self.wz_aes_strength,
                    self.compress_type,
                )
            return wz_aes_extra, crc, compress_type

        def encode_local_header(self, *, crc, compress_type, extra, **kwargs):
            wz_aes_extra, crc, compress_type = self.encode_extra(
                crc, compress_type)
            return super().encode_local_header(
                crc=crc,
                compress_type=compress_type,
                extra=extra+wz_aes_extra,
                **kwargs
            )

        def encode_central_directory(self, *, crc, compress_type, extra_data,
                                     **kwargs):
            wz_aes_extra, crc, compress_type = self.encode_extra(
                crc, compress_type)
            return super().encode_central_directory(
                crc=crc,
                compress_type=compress_type,
                extra_data=extra_data+wz_aes_extra,
                **kwargs)

    class AESZipExtFile(ZipExtFile):

        def setup_aeszipdecrypter(self):
            if not self._pwd:
                raise RuntimeError(
                    'File %r is encrypted with %s encryption and requires a '
                    'password.' % (self.name, WZ_AES)
                )
            encryption_header_length = AESZipDecrypter.encryption_header_length(
                self._zinfo)
            self.encryption_header = self._fileobj.read(encryption_header_length)
            # Adjust read size for encrypted files since the start of the file
            # may be used for the encryption/password information.
            self._orig_compress_left -= encryption_header_length
            # Also remove the hmac length from the end of the file.
            self._orig_compress_left -= AESZipDecrypter.hmac_size

            return AESZipDecrypter

        def setup_decrypter(self):
            if self._zinfo.wz_aes_version is not None:
                return self.setup_aeszipdecrypter()
            return super().setup_decrypter()

        def check_wz_aes(self):
            if self._zinfo.compress_type == ZIP_LZMA:
                # LZMA may have an end of stream marker or padding. Make sure we
                # read that to get the proper HMAC of the compressed byte stream.
                while self._compress_left > 0:
                    data = self._read2(self.MIN_READ_SIZE)
                    # but we don't want to find any more data here.
                    data = self._decompressor.decompress(data)
                    if data:
                        raise BadZipFile(
                            "More data found than indicated by uncompressed size for "
                            "'{}'".format(self.filename)
                        )

            hmac_check = self._fileobj.read(self._decrypter.hmac_size)
            self._decrypter.check_hmac(hmac_check)

        def check_integrity(self):
            if self._zinfo.wz_aes_version is not None:
                self.check_wz_aes()
                if self._expected_crc is not None and self._expected_crc != 0:
                    # Not part of the spec but still check the CRC if it is
                    # supplied when WZ_AES_V2 is specified (no CRC check and CRC
                    # should be 0).
                    self.check_crc()
                elif self._zinfo.wz_aes_version != WZ_AES_V2:
                    # CRC value should be 0 for AES vendor version 2.
                    self.check_crc()
            else:
                super().check_integrity()

    class AESZipFile(ZipFile):
        zipinfo_cls = AESZipInfo
        zipextfile_cls = AESZipExtFile

        def __init__(self, *args, **kwargs):
            encryption = kwargs.pop('encryption', None)
            encryption_kwargs = kwargs.pop('encryption_kwargs', None)
            super().__init__(*args, **kwargs)
            self.encryption = encryption
            self.encryption_kwargs = encryption_kwargs

        def get_encrypter(self):
            if self.encryption == WZ_AES:
                if self.encryption_kwargs is None:
                    encryption_kwargs = {}
                else:
                    encryption_kwargs = self.encryption_kwargs

                return AESZipEncrypter(pwd=self.pwd, **encryption_kwargs)