Initial code to write aes files.

Author: adiroiban

README.md

@@ -1,15 +1,16 @@
 # zipfile-aes
+
 Extend stdlib zipfile to support [WinZip AES](https://www.winzip.com/en/support/aes-encryption/) encryption.
 
 This is an alternative to [pyzipper](https://github.com/danifus/pyzipper),
 with the difference that *zipfile-aes* uses the [cryptography](https://github.com/pyca/cryptography/) library.
-The code from here is influced by pyzipper.
+The code from here is based on pyzipper.
 
 This code can't be included in Python `stdlib` since there is no `crypto` support.
 The [stdlib ssl](https://docs.python.org/3/library/ssl.html) library only exposes the TLS part.
 
-Supported operations (check the automted tests for details):
- * WinZIP AES V2 read without seek
+Supported operations (check the automated tests for details):
+ * WinZIP AES V2 and V1 read without seek.
 
 
 # Contribution
@@ -21,14 +22,14 @@ python3 -m venv .venv
 . .venv/bin/activate
 pip install uv
 uv sync
-uv pip install -e .
-python patch_zipfile.py
+uv pip install -e .[dev]
+uv run patch_zipfile.py
 ```
 
 Use dev tools
 
 ```sh
-pytest
-ruff format .
-ruff check
+uv run pytest
+uv run ruff format .
+uv run ruff check
 ```

src/zipfile_aes/__init__.py

@@ -3,6 +3,8 @@
 Code for handling ZIP aES archives.
 """
 
+import io
+import os
 import struct
 import zipfile_patched as zipfile
 
@@ -56,6 +58,30 @@ def __init__(self, *args, **kwargs):
         self.aes_version = None
         self.aes_compression = None
 
+    def setAESInfo(self, version, strength):
+        """
+        Set the attributes to prepare writing an AES encrypted file.
+        """
+        self.aes_strength = strength
+        self.aes_version = version
+        self.aes_compression = self.compress_type
+
+        self.compress_type = AES_COMPRESSION_TYPE
+
+        # We set it here... but when the file is written,
+        # the stdlib will reset all the flags.
+        self.flag_bits |= zipfile._MASK_ENCRYPTED
+
+        # FIXME
+        # Handle setting the extra multiple times.
+        self.extra += (
+            struct.pack("<HH", EXTRA_AES_HEADER_ID, 7)
+            + self.aes_version
+            + AES_VENDOR_ID
+            + self.aes_strength
+            + struct.pack("<H", self.aes_compression)
+        )
+
     def _decodeExtra(self, filename_crc):
         """
         Little endian encoding.
@@ -125,6 +151,9 @@ def __init__(self, fileobj, mode, zipinfo, pwd=None, close_fileobj=False):
             # This is reverted at the end of init.
             zipinfo.compress_type = zipinfo.aes_compression
 
+            if not pwd:
+                raise RuntimeError("AES encrypted file requires a password.")
+
         self._zipinfo = zipinfo
 
         super().__init__(fileobj, mode, zipinfo, pwd, close_fileobj)
@@ -179,9 +208,9 @@ def _init_aes_decrypter(self):
 
         self._decrypter = AESZipDecipher(self._zipinfo, self._pwd, header)
 
-        # FIXME:
-        # This should be removed once the upstream zipfile
-        # does the password checking as part of decryptor initialization.
+        # TODO:
+        # This can be removed once the upstream zipfile
+        # does the password checking as part of decryption initialization.
         # For AES the password is validate in AESZipDecipher.
         # This is here to reduce the patch for stdlib.
         if self._zipinfo.flag_bits & zipfile._MASK_USE_DATA_DESCRIPTOR:
@@ -271,3 +300,202 @@ class ZipFileWithAES(zipfile.ZipFile):
 
     _ZipInfo = ZipInfoWithAES
     _ZipExtFile = ZipExtFileWithAES
+    _aes_password = None
+    _aes_strength = AES_256
+    _aes_version = AES_V2
+
+    def setAESEncryption(self, password, version=None):
+        self._aes_password = password
+
+    def _open_to_write(self, zinfo, force_zip64=False):
+        if self._aes_password:
+            zinfo.setAESInfo(version=self._aes_version, strength=self._aes_strength)
+
+        # Stdlib will reset the flag_bits.
+        # We set them again in _writecheck
+        stream = super()._open_to_write(zinfo, force_zip64=force_zip64)
+
+        if not self._aes_password:
+            return stream
+
+        # TODO:
+        # Refactor _open_to_write to alow to reuse more code.
+        # Here we close the zipfilewriter created by _open_to_write
+        # and set a fake file and revert the filelist updated at close.
+        original_fp = self.fp
+        try:
+            self.fp = io.BytesIO()
+            stream.close()
+            self.filelist.pop()
+        finally:
+            self.fp = original_fp
+
+        # The need for zip64 is computes inside _open_to_write and
+        # was already set in the file header.
+        return _ZipWithAESWriteFile(self, zinfo, stream._zip64)
+
+    def _writecheck(self, zinfo):
+        """
+        Pretend that we have normal compression.
+
+        This also fixes the flag_bits that are reset by upstream _open_to_write.
+        """
+        try:
+            if zinfo.aes_version:
+                zinfo.compress_type = zinfo.aes_compression
+                zinfo.flag_bits |= zipfile._MASK_ENCRYPTED
+            super()._writecheck(zinfo)
+        finally:
+            if zinfo.aes_version:
+                zinfo.compress_type = AES_COMPRESSION_TYPE
+
+
+class _ZipWithAESWriteFile(zipfile._ZipWriteFile):
+    def __init__(self, zf, zinfo, zip64):
+        self._zinfo = zinfo
+        self._zip64 = zip64
+        self._zipfile = zf
+        # Here we update stdlib to pass the compression.
+        self._compressor = zipfile._get_compressor(
+            zinfo.aes_compression, zinfo._compresslevel
+        )
+        self._file_size = 0
+        self._compress_size = 0
+        self._crc = 0
+
+        self._salt_length = AES_SALT_LENGTHS[zinfo.aes_strength]
+        key_length = AES_KEY_LENGTHS[zinfo.aes_strength]
+        salt = os.urandom(self._salt_length)
+
+        pwd_verify_length = 2
+        dkLen = 2 * key_length + pwd_verify_length
+        kdf = PBKDF2HMAC(
+            algorithm=hashes.SHA1(),
+            length=dkLen,
+            salt=salt,
+            iterations=1000,
+        )
+        keymaterial = kdf.derive(self._zipfile._aes_password)
+
+        encryption_verify = keymaterial[2 * key_length :]
+
+        self._enckey = keymaterial[:key_length]
+        self._counter = 0
+        encmac_key = keymaterial[key_length : 2 * key_length]
+        self._hmac = hmac.HMAC(encmac_key, hashes.SHA1())
+
+        buf = salt + encryption_verify
+        self._compress_size += len(buf)
+        self._fileobj.write(buf)
+
+    def _encrypt(self, data):
+        blocks = self._getBlocks(data)
+        encrypted_data = []
+        for block in blocks:
+            self._counter += 1
+            cipher = Cipher(
+                algorithms.AES(self._enckey),
+                modes.CTR((self._counter).to_bytes(16, byteorder="little")),
+            )
+            data = cipher.encryptor().update(block)
+            data += cipher.encryptor().finalize()
+            encrypted_data.append(data)
+
+        result = b"".join(encrypted_data)
+        self._hmac.update(result)
+        return result
+
+    @staticmethod
+    def _getBlocks(original):
+        """
+        Return AES blocks.
+        """
+        for i in range(0, len(original), 16):
+            yield original[i : i + 16]
+
+    def write(self, data):
+        if self.closed:
+            raise ValueError("I/O operation on closed file.")
+
+        # Accept any data that supports the buffer protocol
+        if isinstance(data, (bytes, bytearray)):
+            nbytes = len(data)
+        else:
+            data = memoryview(data)
+            nbytes = data.nbytes
+        self._file_size += nbytes
+
+        self._crc = zipfile.crc32(data, self._crc)
+        if self._compressor:
+            data = self._compressor.compress(data)
+
+        # TODO:
+        # This is updated to add encryption.
+        data = self._encrypt(data)
+
+        self._compress_size += len(data)
+        # This is stdlib code.
+        self._fileobj.write(data)
+        return nbytes
+
+    def close(self):
+        if self.closed:
+            return
+
+        try:
+            io.BufferedIOBase.close(self)
+
+            # Flush any data from the compressor, encrypt it and update header info
+            if self._compressor:
+                buf = self._compressor.flush()
+            else:
+                buf = b""
+
+            buf = self._encrypt(buf)
+            buf += struct.pack(f"<{AES_HMAC_SIZE}s", self._hmac.finalize()[:10])
+            self._compress_size += len(buf)
+            self._fileobj.write(buf)
+
+            self._zinfo.compress_size = self._compress_size
+
+            # Below is stdlib code
+            self._zinfo.CRC = self._crc
+            self._zinfo.file_size = self._file_size
+
+            if not self._zip64:
+                if self._file_size > zipfile.ZIP64_LIMIT:
+                    raise RuntimeError("File size too large, try using force_zip64")
+                if self._compress_size > zipfile.ZIP64_LIMIT:
+                    raise RuntimeError(
+                        "Compressed size too large, try using force_zip64"
+                    )
+
+            # Write updated header info
+            if self._zinfo.flag_bits & zipfile._MASK_USE_DATA_DESCRIPTOR:
+                # Write CRC and file sizes after the file data
+                fmt = "<LLQQ" if self._zip64 else "<LLLL"
+                self._fileobj.write(
+                    struct.pack(
+                        fmt,
+                        zipfile._DD_SIGNATURE,
+                        self._zinfo.CRC,
+                        self._zinfo.compress_size,
+                        self._zinfo.file_size,
+                    )
+                )
+                self._zipfile.start_dir = self._fileobj.tell()
+            else:
+                # Seek backwards and write file header (which will now include
+                # correct CRC and file sizes)
+
+                # Preserve current position in file
+                self._zipfile.start_dir = self._fileobj.tell()
+                self._fileobj.seek(self._zinfo.header_offset)
+                self._fileobj.write(self._zinfo.FileHeader(self._zip64))
+                self._fileobj.seek(self._zipfile.start_dir)
+
+            # Successfully written: Add file to our caches
+            self._zipfile.filelist.append(self._zinfo)
+            self._zipfile.NameToInfo[self._zinfo.filename] = self._zinfo
+        finally:
+            self._zipfile._writing = False

tests/test_zipfile_aes.py

@@ -15,6 +15,7 @@
     AES_192,
     AES_256,
     ZipFileWithAES,
+    ZipInfoWithAES,
     zipfile,  # Should be moved to canonical name once we no longer need patching.
 )
 
@@ -308,3 +309,34 @@ def test_decrypt_bad_crc_ae2():
         with pytest.raises(zipfile.BadZipFile) as context:
             sut.read("test.txt", pwd=b"test")
         assert "Bad CRC-32 for file 'test.txt'" == context.value.args[0]
+
+
+def test_encrypt():
+    """
+    By default it will encrypt small files with AES v2 and compress
+    using the stored method.
+    """
+    password = b"test"
+    destination = BytesIO()
+    with ZipFileWithAES(destination, "w") as sut:
+        sut.setAESEncryption(password=password)
+
+        with sut.open("test.txt", mode="w") as stream:
+            stream.write(b"content")
+
+    # stdlib context manager will not close the file when we pass
+    # a file object.
+    source = BytesIO(destination.getvalue())
+
+    with ZipFileWithAES(source, "r") as sut:
+        result = sut.infolist()
+        assert 1 == len(result)
+        assert "test.txt" == result[0].filename
+        assert 7 == result[0].file_size
+
+        assert AES_256 == result[0].aes_strength
+        assert AES_V2 == result[0].aes_version
+        assert zipfile.ZIP_STORED == result[0].aes_compression
+
+        content = sut.read("test.txt", pwd=password)
+        assert b"content" == content