diff --git a/sketch_map_tool/config.py b/sketch_map_tool/config.py index 543ba14e..b936819d 100644 --- a/sketch_map_tool/config.py +++ b/sketch_map_tool/config.py @@ -28,6 +28,8 @@ def load_config_default() -> Dict[str, str]: "wms-layers": "heigit:osm-carto@2xx", "wms-read-timeout": 600, "max-nr-simultaneous-uploads": 100, + "max_single_file_size": 838860800, # 100mb + "max_pixel_per_image": 10e8, # 10.000*10.000 } @@ -51,6 +53,8 @@ def load_config_from_env() -> Dict[str, str]: "wms-layers": os.getenv("SMT-WMS-LAYERS"), "wms-read-timeout": os.getenv("SMT-WMS-READ-TIMEOUT"), "max-nr-simultaneous-uploads": os.getenv("SMT-MAX-NR-SIM-UPLOADS"), + "max_single_file_size": os.getenv("MAX-SINGLE-FILE-SIZE"), # 100mb + "max_pixel_per_image": os.getenv("MAX-PIXEL-PER-IMAGE"), } return {k: v for k, v in cfg.items() if v is not None} diff --git a/sketch_map_tool/database/client_flask.py b/sketch_map_tool/database/client_flask.py index 775baff0..0efae31f 100644 --- a/sketch_map_tool/database/client_flask.py +++ b/sketch_map_tool/database/client_flask.py @@ -9,7 +9,11 @@ from sketch_map_tool.config import get_config_value from sketch_map_tool.definitions import REQUEST_TYPES -from sketch_map_tool.exceptions import FileNotFoundError_, UUIDNotFoundError +from sketch_map_tool.exceptions import ( + FileNotFoundError_, + UUIDNotFoundError, +) +from sketch_map_tool.validators import validate_uploaded_sketchmap def open_connection(): @@ -31,7 +35,7 @@ def _insert_id_map(uuid: str, map_: dict): create_query = """ CREATE TABLE IF NOT EXISTS uuid_map( uuid uuid PRIMARY KEY, - ts timestamptz; + ts timestamptz, map json NOT NULL ) """ @@ -96,7 +100,8 @@ def insert_files(files) -> list[int]: curs.execute(create_query) ids = [] for file in files: - curs.execute(insert_query, (secure_filename(file.filename), file.read())) + content = validate_uploaded_sketchmap(file) + curs.execute(insert_query, (secure_filename(file.filename), content)) ids.append(curs.fetchone()[0]) return ids diff --git a/sketch_map_tool/validators.py b/sketch_map_tool/validators.py index 4ebe755e..49463074 100644 --- a/sketch_map_tool/validators.py +++ b/sketch_map_tool/validators.py @@ -1,7 +1,13 @@ +import os +from io import BytesIO from typing import get_args from uuid import UUID +import PIL.Image as Image + +from sketch_map_tool import get_config_value from sketch_map_tool.definitions import REQUEST_TYPES +from sketch_map_tool.exceptions import UploadLimitsExceededError from sketch_map_tool.models import LiteratureReference @@ -14,6 +20,28 @@ def validate_type(type_: REQUEST_TYPES): ) +def validate_uploaded_sketchmap(file): + max_single_file_size = int(get_config_value("max_single_file_size")) + max_pixel_per_image = int(get_config_value("max_pixel_per_image")) + + file_length = file.seek(0, os.SEEK_END) + if file_length > max_single_file_size: + raise UploadLimitsExceededError( + f"You can only upload pictures " + f"up to a filesize of {max_single_file_size}." + ) + file.seek(0, os.SEEK_SET) + content = file.read() + img = Image.open(BytesIO(content)) + total_pxl_cnt = img.size[0] * img.size[1] + if total_pxl_cnt > max_pixel_per_image: + raise UploadLimitsExceededError( + f"You can only upload pictures up to " + f"a total pixel count of {max_pixel_per_image}." + ) + return content + + def validate_uuid(uuid: str): """validation function for endpoint parameter """ try: @@ -26,13 +54,16 @@ def validate_literature_reference(literatur_reference: LiteratureReference): """Validate literatur reference to not include empty strings.""" if literatur_reference.citation == "": raise ValueError( - "Literature reference JSON fields should not contain empty strings as values." + "Literature reference JSON fields " + "should not contain empty strings as values." ) if literatur_reference.img_src == "": raise ValueError( - "Literature reference JSON fields should not contain empty strings as values." + "Literature reference JSON fields should " + "not contain empty strings as values." ) if literatur_reference.url == "": raise ValueError( - "Literature reference JSON fields should not contain empty strings as values." + "Literature reference JSON fields should " + "not contain empty strings as values." ) diff --git a/tests/integration/test_database_client_flask.py b/tests/integration/test_database_client_flask.py index 4a3d9f53..f98c3206 100644 --- a/tests/integration/test_database_client_flask.py +++ b/tests/integration/test_database_client_flask.py @@ -5,7 +5,7 @@ from psycopg2.extensions import connection from sketch_map_tool.database import client_flask -from sketch_map_tool.exceptions import FileNotFoundError_ +from sketch_map_tool.exceptions import FileNotFoundError_, UploadLimitsExceededError def test_open_close_connection(flask_app): @@ -62,6 +62,20 @@ def test_insert_files(flask_app, files): client_flask.delete_file(i) +def test_insert_too_large_files(flask_app, files, monkeypatch): + with flask_app.app_context(): + with pytest.raises(UploadLimitsExceededError): + monkeypatch.setenv("MAX-SINGLE-FILE-SIZE", "10") + client_flask.insert_files(files) + + +def test_insert_too_many_pixels_files(flask_app, files, monkeypatch): + with flask_app.app_context(): + with pytest.raises(UploadLimitsExceededError): + monkeypatch.setenv("MAX-PIXEL-PER-IMAGE", "10") + client_flask.insert_files(files) + + def test_delete_file(flask_app, files): with flask_app.app_context(): ids = client_flask.insert_files(files)