feat: Zip uploaded maps (#11)

Author: alexlambson

docker-compose.yml

@@ -64,6 +64,7 @@ services:
       dockerfile: test.DockerFile
     volumes:
       - .:/cncnet-map-api
+      - ./data/tmp:/tmp/pytest-of-root  # For inspecting files during pytests
     env_file:
       - .env
     environment:

kirovy/models/cnc_map.py

@@ -158,7 +158,7 @@ def set_ban(self, is_banned: bool, banned_by: cnc_user.CncUser) -> None:
         self.save(update_fields=["is_banned"])
 
 
-class CncMapFile(file_base.CncNetFileBaseModel):
+class CncMapFile(file_base.CncNetZippedFileBaseModel):
     """Represents the actual map file that a Command & Conquer game reads.
 
     .. warning::

kirovy/models/file_base.py

@@ -6,6 +6,7 @@
 from kirovy.models.cnc_base_model import CncNetBaseModel
 from kirovy.models import cnc_game as game_models
 from kirovy.utils import file_utils
+from kirovy.zip_storage import ZipFileStorage
 
 
 def _generate_upload_to(instance: "CncNetFileBaseModel", filename: t.Union[str, pathlib.Path]) -> pathlib.Path:
@@ -116,3 +117,16 @@ def generate_upload_to(instance: "CncNetFileBaseModel", filename: str) -> pathli
             str(instance.id),
             filename.name,
         )
+
+
+class CncNetZippedFileBaseModel(CncNetFileBaseModel):
+    """A base file class that will zip and unzip the ``file`` attribute.
+
+    Do **not** use this class for files that will be directly accessed via hyperlink.
+    e.g. don't use this class for images.
+    """
+
+    class Meta:
+        abstract = True
+
+    file = models.FileField(null=False, upload_to=_generate_upload_to, storage=ZipFileStorage)

kirovy/zip_storage.py

@@ -0,0 +1,38 @@
+import io
+import pathlib
+import zipfile
+
+from django.core.files import File
+from django.core.files.storage import FileSystemStorage
+
+
+class ZipFileStorage(FileSystemStorage):
+    def save(self, name: str, content: File, max_length: int | None = None):
+        if is_zipfile(content):
+            return super().save(name, content, max_length)
+
+        internal_filename = pathlib.Path(name).name
+        zip_buffer = io.BytesIO()
+        with zipfile.ZipFile(zip_buffer, "w", allowZip64=False, compresslevel=4) as zf:
+            content.seek(0)
+            zf.writestr(internal_filename, content.read())
+
+        return super().save(f"{name}.zip", zip_buffer, max_length=max_length)
+
+
+def is_zipfile(file_or_path: File | pathlib.Path) -> bool:
+    """Checks if a file is a zip file.
+
+    :param file_or_path:
+        The path to a file, or the file itself, to check.
+    :returns:
+        ``True`` if the file is a zip file.
+    """
+    try:
+        with zipfile.ZipFile(file_or_path, "r") as zf:
+            # zf.getinfo("")  # check if zipfile is valid.
+            return True
+    except zipfile.BadZipfile:
+        return False
+    except Exception as e:
+        raise e

tests/models/test_cnc_map.py

@@ -24,7 +24,7 @@ def test_cnc_map_generate_upload_to(game_yuri, extension_map, file_map_desert, c
         "yr",
         "worlds",
         cnc_map.id.hex,
-        f"yr_{cnc_map.id.hex}_v01.map",
+        f"yr_{cnc_map.id.hex}_v01.map.zip",  # We use the ZipFileStorage for maps.
     )
     saved_map = CncMapFile(
         height=117,  # doesn't matter for this test.

tests/test_views/test_map_upload.py

@@ -1,4 +1,6 @@
 import pathlib
+import zipfile
+import io
 
 from django.core.files.uploadedfile import UploadedFile
 from rest_framework import status
@@ -26,18 +28,29 @@ def test_map_file_upload_happy_path(client_user, file_map_desert, game_yuri, ext
 
     uploaded_file_url: str = response.data["result"]["cnc_map_file"]
     uploaded_image_url: str = response.data["result"]["extracted_preview_file"]
+
+    # We need to strip the url path off of the files,
+    # then check the tmp directory to make sure the uploaded files were saved
     strip_media_url = f"/{settings.MEDIA_URL}"
-    uploaded_file = pathlib.Path(tmp_media_root) / uploaded_file_url.lstrip(strip_media_url)
+    uploaded_zipped_file = pathlib.Path(tmp_media_root) / uploaded_file_url.lstrip(strip_media_url)
     uploaded_image = pathlib.Path(tmp_media_root) / uploaded_image_url.lstrip(strip_media_url)
-    assert uploaded_file.exists()
+    assert uploaded_zipped_file.exists()
     assert uploaded_image.exists()
 
+    # We need to unzip the map so that we can actually verify the saved map contents.
+    map_filename = pathlib.Path(uploaded_file_url).name.replace(".zip", "")
+    # Extract the map from the zipfile, and convert to something that the map parser understands
+    _unzip_io = io.BytesIO()
+    _unzip_io.write(zipfile.ZipFile(uploaded_zipped_file).read(map_filename))
+    _unzip_io.seek(0)
+    uploaded_file = UploadedFile(_unzip_io)
+
     file_response = client_user.get(uploaded_file_url)
     image_response = client_user.get(uploaded_image_url)
     assert file_response.status_code == status.HTTP_200_OK
     assert image_response.status_code == status.HTTP_200_OK
 
-    parser = CncGen2MapParser(UploadedFile(open(uploaded_file, "rb")))
+    parser = CncGen2MapParser(uploaded_file)
     assert parser.ini.get("CnCNet", "ID") == str(response.data["result"]["cnc_map_id"])
 
     map_object = CncMap.objects.get(id=response.data["result"]["cnc_map_id"])