From 45401e325554cab0581f4e7dc9caa2acb29b9458 Mon Sep 17 00:00:00 2001
From: Thomas Forgione <thomas@forgione.fr>
Date: Thu, 1 Aug 2024 09:55:17 +0200
Subject: [PATCH] Factorize response, compute content length

---
 archive.py | 72 ++++++++++++++++++++++++++++++++++++++++++++----------
 1 file changed, 59 insertions(+), 13 deletions(-)

diff --git a/archive.py b/archive.py
index 3a28610..d581aef 100644
--- a/archive.py
+++ b/archive.py
@@ -4,6 +4,8 @@ from flask import Response
 import functools
 import os
 import zlib
+from typing import Optional
+import time
 
 # Chunks for crc 32 computation
 CRC32_CHUNK_SIZE = 65_536
@@ -71,9 +73,31 @@ class ArchiveSender:
     def add_file(self, filename: str, filepath: str):
         self.files[filename] = filepath
 
-    def response(self):
+    def content_length(self) -> Optional[int]:
+        return None
+
+    def generator(self):
         raise NotImplementedError("Abstract method")
 
+    def mime_type(self) -> str:
+        raise NotImplementedError("Abstract method")
+
+    def archive_name(self) -> str:
+        raise NotImplementedError("Abstract method")
+
+    def response(self):
+        headers = {'Content-Disposition': f'attachment; filename="{self.archive_name()}"'}
+
+        length = self.content_length()
+        if length is not None:
+            headers['Content-Length'] = str(length)
+
+        return Response(
+            self.generator(),
+            mimetype=self.mime_type(),
+            headers=headers,
+        )
+
 
 class TarSender(ArchiveSender):
     def generator(self):
@@ -98,12 +122,22 @@ class TarSender(ArchiveSender):
                     yield b'\x00' * (512 - bytes_sent % 512)
         return generate()
 
-    def response(self):
-        return Response(
-            self.generator(),
-            mimetype='application/x-tar',
-            headers={'Content-Disposition': 'attachment; filename="archive.tar"'}
-        )
+    def mime_type(self) -> str:
+        return 'application/x-tar'
+
+    def archive_name(self) -> str:
+        return 'archive.tar'
+
+    def content_length(self) -> int:
+        length = 0
+
+        for file in self.files.values():
+            stat = os.stat(file)
+
+            # Add size of header, and size of content ceiled to 512 bytes
+            length += 512 + stat.st_size + (512 - stat.st_size % 512)
+
+        return length
 
 
 def crc32(filename) -> int:
@@ -263,6 +297,7 @@ class ZipSender(ArchiveSender):
 
                         current_byte += len(bytes)
                         yield bytes
+                        time.sleep(1)
 
             central_directory_size = 0
             centra_directory_offset = current_byte
@@ -277,9 +312,20 @@ class ZipSender(ArchiveSender):
 
         return generate()
 
-    def response(self):
-        return Response(
-            self.generator(),
-            mimetype='application/zip',
-            headers={'Content-Disposition': 'attachment; filename="archive.zip"'}
-        )
+    def content_length(self) -> int:
+        length = 0
+
+        for name, file in self.files.items():
+            stat = os.stat(file)
+
+            # Add size of local file header, central directory file header and file size
+            length += 76 + 2 * len(name) + stat.st_size
+
+        # Add size of end of central directory
+        return length + 22
+
+    def mime_type(self) -> str:
+        return 'application/zip'
+
+    def archive_name(self) -> str:
+        return 'archive.zip'