diff --git a/docker-compose.yml b/docker-compose.yml index efa7507..d7f52b0 100644 --- a/docker-compose.yml +++ b/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: diff --git a/kirovy/models/cnc_map.py b/kirovy/models/cnc_map.py index fc11979..9781649 100644 --- a/kirovy/models/cnc_map.py +++ b/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:: diff --git a/kirovy/models/file_base.py b/kirovy/models/file_base.py index 1fa5d33..2217530 100644 --- a/kirovy/models/file_base.py +++ b/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) diff --git a/kirovy/zip_storage.py b/kirovy/zip_storage.py new file mode 100644 index 0000000..78d65b4 --- /dev/null +++ b/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 diff --git a/tests/models/test_cnc_map.py b/tests/models/test_cnc_map.py index 71012f1..5c27fe7 100644 --- a/tests/models/test_cnc_map.py +++ b/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. diff --git a/tests/test_views/test_map_upload.py b/tests/test_views/test_map_upload.py index c2cbbf0..e5fb9b9 100644 --- a/tests/test_views/test_map_upload.py +++ b/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"])