diff --git a/app/routers/maps.py b/app/routers/maps.py index 8217db6..9e3b7db 100644 --- a/app/routers/maps.py +++ b/app/routers/maps.py @@ -3,6 +3,7 @@ import io import os.path import pathlib import zipfile +import binascii from typing import List, Optional from fastapi import APIRouter, HTTPException, Query from ..schemas import * @@ -42,34 +43,39 @@ async def get_map_by_crc(human_crc: Optional[str] = Query(None, min_length=8, ma @router.post("/maps", tags=["maps"], response_model=MapOut) async def upload_map(map_in: MapIn): - allowed_name = re.compile("[a-zA-Z-_.\d\(\)\[\] ]+.gck") + allowed_name = re.compile(r"[a-zA-Z-_.\d()\[\] ]+.gck") filename = map_in.name if not allowed_name.fullmatch(filename): raise HTTPException(status_code=400, detail="Invalid filename") - map_bytes = base64.b64decode(map_in.b64_data.encode("utf8")) - if len(map_bytes) > MAX_UPLOAD_SIZE: + try: + map_bytes = base64.b64decode(map_in.b64_data.encode("utf8")) + except binascii.Error: + raise HTTPException(status_code=400, detail="Invalid base64 data") + + if len(map_in.b64_data) > MAX_UPLOAD_SIZE: raise HTTPException(status_code=400, detail="File too big") + crc = crc32(map_bytes) + for existing_map in MAPS: + if existing_map["crc"] == crc: + return MapOut(**existing_map) + map_io = io.BytesIO(map_bytes) try: zipfile.ZipFile(map_io) except zipfile.BadZipfile: raise HTTPException(status_code=400, detail="File is not a valid map") - crc = crc32(map_bytes) - existing_map = [gmap for gmap in MAPS if gmap["crc"] == crc] - if existing_map: - return MapOut(**existing_map[0]) - else: - uploaded_filename = filename - i = 0 - while os.path.exists(f"{config['upload_path']}{uploaded_filename}"): - file_wo_ext = ".".join(filename.split('.')[0:-1]) - uploaded_filename = f"{file_wo_ext}-{i}.gck" - i += 1 - with open(f"{config['upload_path']}{uploaded_filename}", "wb") as fp: - fp.write(map_bytes) - gmap = map_file_to_dict(pathlib.Path(f"{config['upload_path']}{uploaded_filename}")) - MAPS.append(gmap) - return MapOut(**gmap) + # upload that map + uploaded_filename = filename + i = 0 + while os.path.exists(f"{config['upload_path']}{uploaded_filename}"): + file_wo_ext = ".".join(filename.split('.')[0:-1]) + uploaded_filename = f"{file_wo_ext}-{i}.gck" + i += 1 + with open(f"{config['upload_path']}{uploaded_filename}", "wb") as fp: + fp.write(map_bytes) + gmap = map_file_to_dict(pathlib.Path(f"{config['upload_path']}{uploaded_filename}")) + MAPS.append(gmap) + return MapOut(**gmap)