Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions docker-compose.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
2 changes: 1 addition & 1 deletion kirovy/models/cnc_map.py
Original file line number Diff line number Diff line change
Expand Up @@ -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::
Expand Down
14 changes: 14 additions & 0 deletions kirovy/models/file_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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)
38 changes: 38 additions & 0 deletions kirovy/zip_storage.py
Original file line number Diff line number Diff line change
@@ -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
2 changes: 1 addition & 1 deletion tests/models/test_cnc_map.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
19 changes: 16 additions & 3 deletions tests/test_views/test_map_upload.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,6 @@
import pathlib
import zipfile
import io

from django.core.files.uploadedfile import UploadedFile
from rest_framework import status
Expand Down Expand Up @@ -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"])
Expand Down