temp zip end headers

storage-app/src/shared/archive_helpers.py

@@ -7,6 +7,18 @@
 from gc import collect as gc_collect
 from .settings import ASYNC_PRODUCER_GC_FREQ as GC_FREQ
 from io import BufferedWriter, BytesIO
+from struct import pack as pack_data
+from zipfile import _Extra, ZIP64_LIMIT, ZIP_FILECOUNT_LIMIT
+
+CENTRAL_STRUCT = "<4s4B4HL2L5H2L"
+CENTRAL_STRING = b"PK\001\002"
+ZIP64_VERSION = 45
+END_64_STRUCT_LOC = "<4sLQL"
+END_64_STRING_LOC = b"PK\x06\x07"
+END_64_STRUCT = "<4sQ2H2L4Q"
+END_64_STRING = b"PK\x06\x06"
+END_STRUCT = b"<4s4H2LH"
+END_STRING = b"PK\005\006"
 
 
 class ZipConsumer(Thread):
@@ -20,42 +32,120 @@ def __init__(
         self.queue = queue
         self.path = path
         self.additional = additional
+        self.file_list = []
 
-    def _dump_buffer(self, dest: BufferedWriter, buffer: BytesIO):
-        print("dump", buffer.tell())
-        buffer.seek(0)
-        dest.write(buffer.read())
+    def _dump_buffer(self, dest: BufferedWriter, buffer: BytesIO, zip_buffer: ZipFile):
+        self.file_list += zip_buffer.filelist
+
+        dest.write(buffer.getvalue())
         buffer.seek(0)
         buffer.truncate(0)
+        dest.flush()
 
-    def run(self):
+    def _finalize(self, dest: BufferedWriter):
         buffer = BytesIO()
-        zip_buffer = ZipFile(self.path, "w", ZIP_DEFLATED)
-        # zip_file = open(self.path, "wb")
 
-        first = False
+        with ZipFile(buffer, "w") as zip_buffer:
+            zip_buffer.start_dir = dest.tell()
+            zip_buffer.filelist = self.file_list
+            zip_buffer._write_end_record()
+
+            dest.write(buffer.getvalue())
+
+        dest.close()
+        buffer.close()
+
+    def _write_end_record(self, buffer):
+        for zinfo in self.file_list:
+            dt = zinfo.date_time
+
+            dosdate = (dt[0] - 1980) << 9 | dt[1] << 5 | dt[2]
+            dostime = dt[3] << 11 | dt[4] << 5 | (dt[5] // 2)
+            extra = []
+
+            assert zinfo.file_size <= ZIP64_LIMIT and zinfo.compress_size <= ZIP64_LIMIT
+
+            file_size = zinfo.file_size
+            compress_size = zinfo.compress_size
+
+            if zinfo.header_offset > ZIP64_LIMIT:
+                extra.append(zinfo.header_offset)
+                header_offset = 0xffffffff
+            else: header_offset = zinfo.header_offset
+
+            extra_data = zinfo.extra
+            min_version = 0
+
+            if extra:
+                extra_data = _Extra.strip(extra_data, (1,))
+                extra_data = pack_data('<HH' + 'Q'*len(extra), 1, 8*len(extra), *extra) + extra_data
+
+                min_version = ZIP64_VERSION
+
+            extract_version = max(min_version, zinfo.extract_version)
+            create_version = max(min_version, zinfo.create_version)
+
+            filename, flag_bits = zinfo._encodeFilenameFlags()
+
+            centdir = pack_data(
+                CENTRAL_STRUCT,
+                CENTRAL_STRING,
+                create_version,
+                zinfo.create_system,
+                extract_version,
+                zinfo.reserved,
+                flag_bits,
+                zinfo.compress_type,
+                dostime,
+                dosdate,
+                zinfo.CRC,
+                compress_size,
+                file_size,
+                len(filename),
+                len(extra_data),
+                len(zinfo.comment),
+                0,
+                zinfo.internal_attr,
+                zinfo.external_attr,
+                header_offset
+            )
+
+            buffer.write(centdir, filename, extra_data, zinfo.comment)
+
+    def _write_cent_dir(self, f_size, d_size, buffer):
+        if d_size > ZIP_FILECOUNT_LIMIT or f_size > ZIP64_LIMIT:
+            pack = (END_64_STRUCT, END_64_STRING, 44, 45, 45, 0, 0, d_size, d_size, 0, f_size)
+            buffer.write(pack_data(*pack))
+            buffer.write(pack_data(END_64_STRUCT_LOC, END_64_STRING_LOC, 0, f_size, 1))
+            d_size = min(d_size, 0xFFFF)
+
+        buffer.write(pack_data(END_STRUCT, END_STRING, 0, 0, d_size, d_size, 0, f_size, 0))
+
+    def run(self):
+        buffer = BytesIO()
+        zip_buffer: ZipFile = ZipFile(buffer, "w", ZIP_DEFLATED)
+        zip_file = open(self.path, "wb")
 
         while True:
             if (task := self.queue.get()) is None: break
 
-            if not first:
-                print("consume first")
-                first = True
-
             f_name, f_data = task
             zip_buffer.writestr(f_name, f_data)
 
             del f_data
 
-            # if buffer.tell() > (50 << 20): self._dump_buffer(zip_file, buffer)
+            if buffer.tell() > (100 << 20):
+                self._dump_buffer(zip_file, buffer, zip_buffer)
+                zip_buffer = ZipFile(buffer, "w", ZIP_DEFLATED)
 
         for f_name, f_data in self.additional: zip_buffer.writestr(f_name, f_data)
 
-        # if buffer.tell(): self._dump_buffer(zip_file, buffer)
+        if buffer.tell(): self._dump_buffer(zip_file, buffer, zip_buffer)
 
         zip_buffer.close()
         buffer.close()
-        #zip_file.close()
+
+        self._finalize(zip_file)
 
 
 class FileProducer: