diff --git a/stream_zip.py b/stream_zip.py index 991600c..edf26f2 100644 --- a/stream_zip.py +++ b/stream_zip.py @@ -136,7 +136,10 @@ def _raise_if_beyond(offset, maximum, exception_class): if offset > maximum: raise exception_class() - def _zip_64_local_header_and_data(name_encoded, mod_at_ms_dos, mod_at_unix_extra, aes_extra, external_attr, uncompressed_size, crc_32, _get_compress_obj, chunks): + def _no_encryption(chunks): + return (yield from chunks) + + def _zip_64_local_header_and_data(name_encoded, mod_at_ms_dos, mod_at_unix_extra, aes_extra, external_attr, uncompressed_size, crc_32, _get_compress_obj, encryption_func, chunks): file_offset = offset _raise_if_beyond(file_offset, maximum=0xffffffffffffffff, exception_class=OffsetOverflowError) @@ -205,7 +208,7 @@ def _zip_64_local_header_and_data(name_encoded, mod_at_ms_dos, mod_at_unix_extra 0xffffffff, # Offset of local header - since zip64 ), name_encoded, extra - def _zip_32_local_header_and_data(name_encoded, mod_at_ms_dos, mod_at_unix_extra, aes_extra, external_attr, uncompressed_size, crc_32, _get_compress_obj, chunks): + def _zip_32_local_header_and_data(name_encoded, mod_at_ms_dos, mod_at_unix_extra, aes_extra, external_attr, uncompressed_size, crc_32, _get_compress_obj, encryption_func, chunks): file_offset = offset _raise_if_beyond(file_offset, maximum=0xffffffff, exception_class=OffsetOverflowError) @@ -232,12 +235,12 @@ def _zip_32_local_header_and_data(name_encoded, mod_at_ms_dos, mod_at_unix_extra yield from _(name_encoded) yield from _(extra) - uncompressed_size, compressed_size, crc_32 = yield from _zip_data( + uncompressed_size, compressed_size, crc_32 = yield from encryption_func(_zip_data( chunks, _get_compress_obj, max_uncompressed_size=0xffffffff, max_compressed_size=0xffffffff, - ) + )) yield from _(data_descriptor_signature) yield from _(data_descriptor_zip_32_struct.pack(crc_32, compressed_size, uncompressed_size)) @@ -289,7 +292,7 @@ def _zip_data(chunks, _get_compress_obj, max_uncompressed_size, max_compressed_s return uncompressed_size, compressed_size, crc_32 - def _no_compression_64_local_header_and_data(name_encoded, mod_at_ms_dos, mod_at_unix_extra, aes_extra, external_attr, uncompressed_size, crc_32, _get_compress_obj, chunks): + def _no_compression_64_local_header_and_data(name_encoded, mod_at_ms_dos, mod_at_unix_extra, aes_extra, external_attr, uncompressed_size, crc_32, _get_compress_obj, encryption_func, chunks): file_offset = offset _raise_if_beyond(file_offset, maximum=0xffffffffffffffff, exception_class=OffsetOverflowError) @@ -323,7 +326,7 @@ def _no_compression_64_local_header_and_data(name_encoded, mod_at_ms_dos, mod_at yield from _(name_encoded) yield from _(extra) - for chunk in chunks: + for chunk in encryption_func(chunks): yield from _(chunk) extra = zip_64_central_directory_extra_struct.pack( @@ -354,7 +357,7 @@ def _no_compression_64_local_header_and_data(name_encoded, mod_at_ms_dos, mod_at ), name_encoded, extra - def _no_compression_32_local_header_and_data(name_encoded, mod_at_ms_dos, mod_at_unix_extra, aes_extra, external_attr, uncompressed_size, crc_32, _get_compress_obj, chunks): + def _no_compression_32_local_header_and_data(name_encoded, mod_at_ms_dos, mod_at_unix_extra, aes_extra, external_attr, uncompressed_size, crc_32, _get_compress_obj, encryption_func, chunks): file_offset = offset _raise_if_beyond(file_offset, maximum=0xffffffff, exception_class=OffsetOverflowError) @@ -382,7 +385,7 @@ def _no_compression_32_local_header_and_data(name_encoded, mod_at_ms_dos, mod_at yield from _(name_encoded) yield from _(extra) - for chunk in chunks: + for chunk in encryption_func(chunks): yield from _(chunk) return central_directory_header_struct.pack( @@ -424,7 +427,7 @@ def _chunks(): return chunks, size, crc_32 - def _no_compression_streamed_64_local_header_and_data(name_encoded, mod_at_ms_dos, mod_at_unix_extra, aes_extra, external_attr, uncompressed_size, crc_32, _get_compress_obj, chunks): + def _no_compression_streamed_64_local_header_and_data(name_encoded, mod_at_ms_dos, mod_at_unix_extra, aes_extra, external_attr, uncompressed_size, crc_32, _get_compress_obj, encryption_func, chunks): file_offset = offset _raise_if_beyond(file_offset, maximum=0xffffffffffffffff, exception_class=OffsetOverflowError) @@ -456,7 +459,7 @@ def _no_compression_streamed_64_local_header_and_data(name_encoded, mod_at_ms_do yield from _(name_encoded) yield from _(extra) - yield from _no_compression_streamed_data(chunks, uncompressed_size, crc_32, 0xffffffffffffffff) + yield from encryption_func(_no_compression_streamed_data(chunks, uncompressed_size, crc_32, 0xffffffffffffffff)) extra = zip_64_central_directory_extra_struct.pack( zip_64_extra_signature, @@ -486,7 +489,7 @@ def _no_compression_streamed_64_local_header_and_data(name_encoded, mod_at_ms_do ), name_encoded, extra - def _no_compression_streamed_32_local_header_and_data(name_encoded, mod_at_ms_dos, mod_at_unix_extra, aes_extra, external_attr, uncompressed_size, crc_32, _get_compress_obj, chunks): + def _no_compression_streamed_32_local_header_and_data(name_encoded, mod_at_ms_dos, mod_at_unix_extra, aes_extra, external_attr, uncompressed_size, crc_32, _get_compress_obj, encryption_func, chunks): file_offset = offset _raise_if_beyond(file_offset, maximum=0xffffffff, exception_class=OffsetOverflowError) @@ -513,7 +516,7 @@ def _no_compression_streamed_32_local_header_and_data(name_encoded, mod_at_ms_do yield from _(name_encoded) yield from _(extra) - yield from _no_compression_streamed_data(chunks, uncompressed_size, crc_32, 0xffffffff) + yield from encryption_func(_no_compression_streamed_data(chunks, uncompressed_size, crc_32, 0xffffffff)) return central_directory_header_struct.pack( 20, # Version made by @@ -586,9 +589,9 @@ def _no_compression_streamed_data(chunks, uncompressed_size, crc_32, maximum_siz _no_compression_streamed_64_local_header_and_data if _method is _NO_COMPRESSION_STREAMED_64 else \ _no_compression_streamed_32_local_header_and_data - aes_extra = b'' + aes_extra, encryption_func = (b'', _no_encryption) - central_directory_header_entry, name_encoded, extra = yield from data_func(name_encoded, mod_at_ms_dos, mod_at_unix_extra, aes_extra, external_attr, uncompressed_size, crc_32, _get_compress_obj, evenly_sized(chunks)) + central_directory_header_entry, name_encoded, extra = yield from data_func(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) central_directory.append((central_directory_header_entry, name_encoded, extra))