From 4810da2b52d13576f68181393d8998819eaef24f Mon Sep 17 00:00:00 2001
From: Michal Charemza <michal@charemza.name>
Date: Wed, 3 Jan 2024 16:07:07 +0000
Subject: [PATCH] feat: AES-2 encryption

---
 pyproject.toml     |  4 ++
 stream_zip.py      | 91 +++++++++++++++++++++++++++++++---------------
 test_stream_zip.py | 44 ++++++++++++++++++++++
 3 files changed, 109 insertions(+), 30 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 1cf5107..42d6b6b 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -16,6 +16,9 @@ classifiers = [
     "License :: OSI Approved :: MIT License",
     "Topic :: System :: Archiving :: Compression",
 ]
+dependencies = [
+    "pycryptodome>=3.10.1",
+]
 
 [project.optional-dependencies]
 dev = [
@@ -25,6 +28,7 @@ dev = [
     "stream-unzip>=0.0.86"
 ]
 ci = [
+    "pycryptodome==3.10.1",
     "coverage==6.2",
     "pytest==6.2.5",
     "pytest-cov==3.0.0",
diff --git a/stream_zip.py b/stream_zip.py
index b2c6439..d4005cc 100644
--- a/stream_zip.py
+++ b/stream_zip.py
@@ -1,7 +1,13 @@
 from collections import deque
 from struct import Struct
+import secrets
 import zlib
 
+from Crypto.Cipher import AES
+from Crypto.Hash import HMAC, SHA1
+from Crypto.Util import Counter
+from Crypto.Protocol.KDF import PBKDF2
+
 # Private methods
 
 _NO_COMPRESSION_BUFFERED_32 = object()
@@ -63,7 +69,7 @@ def method_compressobj(offset, default_get_compressobj):
     return method_compressobj
 
 
-def stream_zip(files, chunk_size=65536, get_compressobj=lambda: zlib.compressobj(wbits=-zlib.MAX_WBITS, level=9), extended_timestamps=True):
+def stream_zip(files, chunk_size=65536, get_compressobj=lambda: zlib.compressobj(wbits=-zlib.MAX_WBITS, level=9), extended_timestamps=True, password=None):
 
     def evenly_sized(chunks):
         chunk = b''
@@ -119,6 +125,9 @@ def get_zipped_chunks_uneven():
         mod_at_unix_extra_signature = b'UT'
         mod_at_unix_extra_struct = Struct('<2sH1sl')
 
+        aes_extra_signature = b'\x01\x99'
+        aes_extra_struct = Struct('<2sHH2sBH')
+
         modified_at_struct = Struct('<HH')
 
         central_directory = deque()
@@ -188,7 +197,7 @@ def _zip_64_local_header_and_data(compression, name_encoded, mod_at_ms_dos, mod_
 
             # (encryption,) data descriptor and utf-8 file names
             flags = \
-                b'\x88\x08' if aes_extra else \
+                b'\x09\x08' if aes_extra else \
                 b'\x08\x08'
 
             yield from _(local_header_signature)
@@ -206,12 +215,15 @@ def _zip_64_local_header_and_data(compression, name_encoded, mod_at_ms_dos, mod_
             yield from _(name_encoded)
             yield from _(extra)
 
-            uncompressed_size, compressed_size, crc_32 = yield from _zip_data(
+            uncompressed_size, raw_compressed_size, crc_32 = yield from encryption_func(_zip_data(
                 chunks,
                 _get_compress_obj,
                 max_uncompressed_size=0xffffffffffffffff,
                 max_compressed_size=0xffffffffffffffff,
-            )
+            ))
+            compressed_size = \
+                raw_compressed_size + 28 if aes_extra else \
+                raw_compressed_size
 
             yield from _(data_descriptor_signature)
             yield from _(data_descriptor_zip_64_struct.pack(crc_32, compressed_size, uncompressed_size))
@@ -252,7 +264,7 @@ def _zip_32_local_header_and_data(compression, name_encoded, mod_at_ms_dos, mod_
 
             # (encryption,) data descriptor and utf-8 file names
             flags = \
-                b'\x88\x08' if aes_extra else \
+                b'\x09\x08' if aes_extra else \
                 b'\x08\x08'
 
             yield from _(local_header_signature)
@@ -270,12 +282,15 @@ def _zip_32_local_header_and_data(compression, name_encoded, mod_at_ms_dos, mod_
             yield from _(name_encoded)
             yield from _(extra)
 
-            uncompressed_size, compressed_size, crc_32 = yield from encryption_func(_zip_data(
+            uncompressed_size, raw_compressed_size, crc_32 = yield from encryption_func(_zip_data(
                 chunks,
                 _get_compress_obj,
                 max_uncompressed_size=0xffffffff,
                 max_compressed_size=0xffffffff,
             ))
+            compressed_size = \
+                raw_compressed_size + 28 if aes_extra else \
+                raw_compressed_size
 
             yield from _(data_descriptor_signature)
             yield from _(data_descriptor_zip_32_struct.pack(crc_32, compressed_size, uncompressed_size))
@@ -332,18 +347,21 @@ def _no_compression_64_local_header_and_data(compression, name_encoded, mod_at_m
 
             _raise_if_beyond(file_offset, maximum=0xffffffffffffffff, exception_class=OffsetOverflowError)
 
-            chunks, size, crc_32 = _no_compression_buffered_data_size_crc_32(chunks, maximum_size=0xffffffffffffffff)
+            chunks, uncompressed_size, crc_32 = _no_compression_buffered_data_size_crc_32(chunks, maximum_size=0xffffffffffffffff)
+            compressed_size = \
+                uncompressed_size + 28 if aes_extra else \
+                uncompressed_size
 
             extra = zip_64_local_extra_struct.pack(
                 zip_64_extra_signature,
                 16,    # Size of extra
-                size,  # Uncompressed
-                size,  # Compressed
+                uncompressed_size,
+                compressed_size,
             ) + mod_at_unix_extra + aes_extra
 
             # (encryption and) utf-8 file names
             flags = \
-                b'\x80\x08' if aes_extra else \
+                b'\x01\x08' if aes_extra else \
                 b'\x00\x08'
 
             yield from _(local_header_signature)
@@ -367,8 +385,8 @@ def _no_compression_64_local_header_and_data(compression, name_encoded, mod_at_m
             extra = zip_64_central_directory_extra_struct.pack(
                 zip_64_extra_signature,
                 24,    # Size of extra
-                size,  # Uncompressed
-                size,  # Compressed
+                uncompressed_size,
+                compressed_size,
                 file_offset,
             ) + mod_at_unix_extra + aes_extra
             return central_directory_header_struct.pack(
@@ -397,11 +415,14 @@ def _no_compression_32_local_header_and_data(compression, name_encoded, mod_at_m
 
             _raise_if_beyond(file_offset, maximum=0xffffffff, exception_class=OffsetOverflowError)
 
-            chunks, size, crc_32 = _no_compression_buffered_data_size_crc_32(chunks, maximum_size=0xffffffff)
+            chunks, uncompressed_size, crc_32 = _no_compression_buffered_data_size_crc_32(chunks, maximum_size=0xffffffff)
+            compressed_size = \
+                uncompressed_size + 28 if aes_extra else \
+                uncompressed_size
 
             # (encryption and) utf-8 file names
             flags = \
-                b'\x80\x08' if aes_extra else \
+                b'\x01\x08' if aes_extra else \
                 b'\x00\x08'
 
             extra = mod_at_unix_extra + aes_extra
@@ -412,8 +433,8 @@ def _no_compression_32_local_header_and_data(compression, name_encoded, mod_at_m
                 compression,
                 mod_at_ms_dos,
                 crc_32,
-                size,         # Compressed
-                size,         # Uncompressed
+                compressed_size,
+                uncompressed_size,
                 len(name_encoded),
                 len(extra),
             ))
@@ -432,8 +453,8 @@ def _no_compression_32_local_header_and_data(compression, name_encoded, mod_at_m
                compression,
                mod_at_ms_dos,
                crc_32,
-               size,         # Compressed
-               size,         # Uncompressed
+               compressed_size,
+               uncompressed_size,
                len(name_encoded),
                len(extra),
                0,            # File comment length
@@ -467,16 +488,20 @@ def _no_compression_streamed_64_local_header_and_data(compression, name_encoded,
 
             _raise_if_beyond(file_offset, maximum=0xffffffffffffffff, exception_class=OffsetOverflowError)
 
+            compressed_size = \
+                uncompressed_size + 28 if aes_extra else \
+                uncompressed_size
+
             extra = zip_64_local_extra_struct.pack(
                 zip_64_extra_signature,
                 16,                 # Size of extra
-                uncompressed_size,  # Uncompressed
-                uncompressed_size,  # Compressed
+                uncompressed_size,
+                compressed_size,
             ) + mod_at_unix_extra + aes_extra
 
             # (encryption and) utf-8 file names
             flags = \
-                b'\x80\x08' if aes_extra else \
+                b'\x01\x08' if aes_extra else \
                 b'\x00\x08'
 
             yield from _(local_header_signature)
@@ -499,8 +524,8 @@ def _no_compression_streamed_64_local_header_and_data(compression, name_encoded,
             extra = zip_64_central_directory_extra_struct.pack(
                 zip_64_extra_signature,
                 24,                 # Size of extra
-                uncompressed_size,  # Uncompressed
-                uncompressed_size,  # Compressed
+                uncompressed_size,
+                compressed_size,
                 file_offset,
             ) + mod_at_unix_extra
             return central_directory_header_struct.pack(
@@ -531,9 +556,13 @@ def _no_compression_streamed_32_local_header_and_data(compression, name_encoded,
 
             extra = mod_at_unix_extra + aes_extra
 
+            compressed_size = \
+                uncompressed_size + 28 if aes_extra else \
+                uncompressed_size
+
             # (encryption and) utf-8 file names
             flags = \
-                b'\x80\x08' if aes_extra else \
+                b'\x01\x08' if aes_extra else \
                 b'\x00\x08'
 
             yield from _(local_header_signature)
@@ -543,8 +572,8 @@ def _no_compression_streamed_32_local_header_and_data(compression, name_encoded,
                 compression,
                 mod_at_ms_dos,
                 crc_32,
-                uncompressed_size,  # Compressed
-                uncompressed_size,  # Uncompressed
+                compressed_size,
+                uncompressed_size,
                 len(name_encoded),
                 len(extra),
             ))
@@ -562,8 +591,8 @@ def _no_compression_streamed_32_local_header_and_data(compression, name_encoded,
                compression,
                mod_at_ms_dos,
                crc_32,
-               uncompressed_size,  # Compressed
-               uncompressed_size,  # Uncompressed
+               compressed_size,
+               uncompressed_size,
                len(name_encoded),
                len(extra),
                0,                  # File comment length
@@ -616,7 +645,7 @@ def _no_compression_streamed_data(chunks, uncompressed_size, crc_32, maximum_siz
                 (mode << 16) | \
                 (0x10 if name_encoded[-1:] == b'/' else 0x0)  # MS-DOS directory
 
-            data_func, compression = \
+            data_func, raw_compression = \
                 (_zip_64_local_header_and_data, 8) if _method is _ZIP_64 else \
                 (_zip_32_local_header_and_data, 8) if _method is _ZIP_32 else \
                 (_no_compression_64_local_header_and_data, 0) if _method is _NO_COMPRESSION_BUFFERED_64 else \
@@ -624,7 +653,9 @@ def _no_compression_streamed_data(chunks, uncompressed_size, crc_32, maximum_siz
                 (_no_compression_streamed_64_local_header_and_data, 0) if _method is _NO_COMPRESSION_STREAMED_64 else \
                 (_no_compression_streamed_32_local_header_and_data, 0)
 
-            aes_extra, encryption_func = (b'', _no_encryption)
+            compression, aes_extra, encryption_func = \
+                (99, aes_extra_struct.pack(aes_extra_signature, 7, 2, b'AE', 3, raw_compression), _aes_encrypted) if password is not None else \
+                (raw_compression, b'', _no_encryption)
 
             central_directory_header_entry, name_encoded, extra = yield from data_func(compression, name_encoded, mod_at_ms_dos, mod_at_unix_extra, aes_extra, external_attr, uncompressed_size, crc_32, _get_compress_obj, encryption_func, evenly_sized(chunks))
             central_directory_size += len(central_directory_header_signature) + len(central_directory_header_entry) + len(name_encoded) + len(extra)
diff --git a/test_stream_zip.py b/test_stream_zip.py
index 623197e..02afab6 100644
--- a/test_stream_zip.py
+++ b/test_stream_zip.py
@@ -1082,3 +1082,47 @@ def test_unzip_modification_time_extended_timestamps_disabled(method, timezone,
         subprocess.run(['unzip', f'{d}/test.zip', '-d', d], env={'TZ': timezone})
 
         assert os.path.getmtime('my_file') == expected_modified_at.timestamp()
+
+
+def test_password_unzips_with_stream_unzip():
+    now = datetime.strptime('2021-01-01 21:01:12', '%Y-%m-%d %H:%M:%S')
+    mode = stat.S_IFREG | 0o600
+    password = 'my-pass'
+
+    files = (
+        ('file-1', now, mode, ZIP_32, (b'a' * 10000, b'b' * 10000)),
+        ('file-2', now, mode, ZIP_64, (b'a' * 10000, b'b' * 10000)),
+        ('file-3', now, mode, NO_COMPRESSION_64, (b'a' * 10000, b'b' * 10000)),
+        ('file-4', now, mode, NO_COMPRESSION_64(20000, 2664091433), (b'a' * 10000, b'b' * 10000)),
+        ('file-5', now, mode, NO_COMPRESSION_32, (b'a' * 10000, b'b' * 10000)),
+        ('file-6', now, mode, NO_COMPRESSION_32(20000, 2664091433), (b'a' * 10000, b'b' * 10000)),
+    )
+
+    assert [
+        (b'file-1', None, b'a' * 10000 + b'b' * 10000),
+        (b'file-2', None, b'a' * 10000 + b'b' * 10000),
+        (b'file-3', 20000, b'a' * 10000 + b'b' * 10000),
+        (b'file-4', 20000, b'a' * 10000 + b'b' * 10000),
+        (b'file-5', 20000, b'a' * 10000 + b'b' * 10000),
+        (b'file-6', 20000, b'a' * 10000 + b'b' * 10000),
+    ] == [
+        (name, size, b''.join(chunks))
+        for name, size, chunks in stream_unzip(stream_zip(files, password=password), password=password)
+    ]
+
+
+def test_password_bytes_not_deterministic():
+    now = datetime.strptime('2021-01-01 21:01:12', '%Y-%m-%d %H:%M:%S')
+    mode = stat.S_IFREG | 0o600
+    password = 'my-pass'
+
+    files = (
+        ('file-1', now, mode, ZIP_32, (b'a' * 10000, b'b' * 10000)),
+        ('file-2', now, mode, ZIP_64, (b'a' * 10000, b'b' * 10000)),
+        ('file-3', now, mode, NO_COMPRESSION_64, (b'a' * 10000, b'b' * 10000)),
+        ('file-4', now, mode, NO_COMPRESSION_64(20000, 2664091433), (b'a' * 10000, b'b' * 10000)),
+        ('file-5', now, mode, NO_COMPRESSION_32, (b'a' * 10000, b'b' * 10000)),
+        ('file-6', now, mode, NO_COMPRESSION_32(20000, 2664091433), (b'a' * 10000, b'b' * 10000)),
+    )
+
+    assert b''.join(stream_zip(files, password=password)) != b''.join(stream_zip(files, password=password))