Skip to content

Commit

Permalink
refactor: start of aes extra field support
Browse files Browse the repository at this point in the history
For now it's always the empty bytes, but will add to it in later commits
  • Loading branch information
michalc committed Jan 3, 2024
1 parent 2c63c5b commit a7d1884
Showing 1 changed file with 17 additions and 15 deletions.
32 changes: 17 additions & 15 deletions stream_zip.py
Original file line number Diff line number Diff line change
Expand Up @@ -136,7 +136,7 @@ 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, external_attr, uncompressed_size, crc_32, _get_compress_obj, 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, chunks):
file_offset = offset

_raise_if_beyond(file_offset, maximum=0xffffffffffffffff, exception_class=OffsetOverflowError)
Expand All @@ -146,7 +146,7 @@ def _zip_64_local_header_and_data(name_encoded, mod_at_ms_dos, mod_at_unix_extra
16, # Size of extra
0, # Uncompressed size - since data descriptor
0, # Compressed size - since data descriptor
) + mod_at_unix_extra
) + mod_at_unix_extra + aes_extra
yield from _(local_header_signature)
yield from _(local_header_struct.pack(
45, # Version
Expand Down Expand Up @@ -178,7 +178,7 @@ def _zip_64_local_header_and_data(name_encoded, mod_at_ms_dos, mod_at_unix_extra
uncompressed_size,
compressed_size,
file_offset,
) + mod_at_unix_extra
) + mod_at_unix_extra + aes_extra
return central_directory_header_struct.pack(
45, # Version made by
3, # System made by (UNIX)
Expand All @@ -199,12 +199,12 @@ 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, 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, chunks):
file_offset = offset

_raise_if_beyond(file_offset, maximum=0xffffffff, exception_class=OffsetOverflowError)

extra = mod_at_unix_extra
extra = mod_at_unix_extra + aes_extra
yield from _(local_header_signature)
yield from _(local_header_struct.pack(
20, # Version
Expand Down Expand Up @@ -277,7 +277,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, 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, chunks):
file_offset = offset

_raise_if_beyond(file_offset, maximum=0xffffffffffffffff, exception_class=OffsetOverflowError)
Expand All @@ -289,7 +289,7 @@ def _no_compression_64_local_header_and_data(name_encoded, mod_at_ms_dos, mod_at
16, # Size of extra
size, # Uncompressed
size, # Compressed
) + mod_at_unix_extra
) + mod_at_unix_extra + aes_extra
yield from _(local_header_signature)
yield from _(local_header_struct.pack(
45, # Version
Expand All @@ -314,7 +314,7 @@ def _no_compression_64_local_header_and_data(name_encoded, mod_at_ms_dos, mod_at
size, # Uncompressed
size, # Compressed
file_offset,
) + mod_at_unix_extra
) + mod_at_unix_extra + aes_extra
return central_directory_header_struct.pack(
45, # Version made by
3, # System made by (UNIX)
Expand All @@ -336,14 +336,14 @@ 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, 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, chunks):
file_offset = offset

_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)

extra = mod_at_unix_extra
extra = mod_at_unix_extra + aes_extra
yield from _(local_header_signature)
yield from _(local_header_struct.pack(
20, # Version
Expand Down Expand Up @@ -401,7 +401,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, 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, chunks):
file_offset = offset

_raise_if_beyond(file_offset, maximum=0xffffffffffffffff, exception_class=OffsetOverflowError)
Expand All @@ -411,7 +411,7 @@ def _no_compression_streamed_64_local_header_and_data(name_encoded, mod_at_ms_do
16, # Size of extra
uncompressed_size, # Uncompressed
uncompressed_size, # Compressed
) + mod_at_unix_extra
) + mod_at_unix_extra + aes_extra
yield from _(local_header_signature)
yield from _(local_header_struct.pack(
45, # Version
Expand Down Expand Up @@ -457,12 +457,12 @@ 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, 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, chunks):
file_offset = offset

_raise_if_beyond(file_offset, maximum=0xffffffff, exception_class=OffsetOverflowError)

extra = mod_at_unix_extra
extra = mod_at_unix_extra + aes_extra
yield from _(local_header_signature)
yield from _(local_header_struct.pack(
20, # Version
Expand Down Expand Up @@ -551,7 +551,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

central_directory_header_entry, name_encoded, extra = yield from data_func(name_encoded, mod_at_ms_dos, mod_at_unix_extra, external_attr, uncompressed_size, crc_32, _get_compress_obj, evenly_sized(chunks))
aes_extra = b''

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_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))

Expand Down

0 comments on commit a7d1884

Please sign in to comment.