feat(flask): validate number of pixel of uploaded files

* feat: size restrictions by pixel and file size

* test(validators) add unit test

* refactor: move validation of files to routes.py

---------

Co-authored-by: Matthias Schaub <[email protected]>

sketch_map_tool/config.py

@@ -28,6 +28,7 @@ def load_config_default() -> Dict[str, str]:
         "wms-layers": "heigit:osm-carto@2xx",
         "wms-read-timeout": 600,
         "max-nr-simultaneous-uploads": 100,
+        "max_pixel_per_image": 10e8,  # 10.000*10.000
     }
 
 
@@ -51,6 +52,7 @@ 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_pixel_per_image": os.getenv("MAX-PIXEL-PER-IMAGE"),
     }
     return {k: v for k, v in cfg.items() if v is not None}
 

sketch_map_tool/database/client_flask.py

@@ -8,7 +8,10 @@
 
 from sketch_map_tool.config import get_config_value
 from sketch_map_tool.definitions import REQUEST_TYPES
-from sketch_map_tool.exceptions import CustomFileNotFoundError, UUIDNotFoundError
+from sketch_map_tool.exceptions import (
+    CustomFileNotFoundError,
+    UUIDNotFoundError,
+)
 
 
 def open_connection():
@@ -89,7 +92,6 @@ def insert_files(files) -> list[int]:
     insert_query = "INSERT INTO blob(file_name, file) VALUES (%s, %s) RETURNING id"
     db_conn = open_connection()
     with db_conn.cursor() as curs:
-        # executemany and fetchall does not work together
         curs.execute(create_query)
         ids = []
         for file in files:

sketch_map_tool/routes.py

@@ -9,7 +9,6 @@
 
 from sketch_map_tool import celery_app, definitions, tasks, upload_processing
 from sketch_map_tool import flask_app as app
-from sketch_map_tool.config import get_config_value
 from sketch_map_tool.database import client_flask as db_client_flask
 from sketch_map_tool.definitions import REQUEST_TYPES
 from sketch_map_tool.exceptions import (
@@ -23,7 +22,11 @@
 from sketch_map_tool.helpers import to_array
 from sketch_map_tool.models import Bbox, PaperFormat, Size
 from sketch_map_tool.tasks import digitize_sketches, georeference_sketch_maps
-from sketch_map_tool.validators import validate_type, validate_uuid
+from sketch_map_tool.validators import (
+    validate_type,
+    validate_uploaded_sketchmaps,
+    validate_uuid,
+)
 
 
 @app.get("/")
@@ -105,13 +108,8 @@ def digitize_results_post() -> Response:
     if "file" not in request.files:
         return redirect(url_for("digitize"))
     files = request.files.getlist("file")
-    max_nr_simultaneous_uploads = int(get_config_value("max-nr-simultaneous-uploads"))
-    if len(files) > max_nr_simultaneous_uploads:
-        raise UploadLimitsExceededError(
-            f"You can only upload up to {max_nr_simultaneous_uploads} files at once."
-        )
+    validate_uploaded_sketchmaps(files)
     ids = db_client_flask.insert_files(files)
-
     file_names = [db_client_flask.select_file_name(i) for i in ids]
     args = [
         upload_processing.read_qr_code(to_array(db_client_flask.select_file(_id)))

sketch_map_tool/validators.py

@@ -1,7 +1,12 @@
 from typing import get_args
 from uuid import UUID
 
+import PIL.Image as Image
+from werkzeug.datastructures import FileStorage
+
+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 +19,27 @@ def validate_type(type_: REQUEST_TYPES):
         )
 
 
+def validate_uploaded_sketchmaps(files: list[FileStorage]):
+    """Validation function for uploaded files."""
+
+    max_nr_simultaneous_uploads = int(get_config_value("max-nr-simultaneous-uploads"))
+    max_pixel_per_image = int(get_config_value("max_pixel_per_image"))
+
+    if len(files) > max_nr_simultaneous_uploads:
+        raise UploadLimitsExceededError(
+            f"You can only upload up to {max_nr_simultaneous_uploads} files at once."
+        )
+
+    for file in files:
+        img = Image.open(file)
+        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}."
+            )
+
+
 def validate_uuid(uuid: str):
     """validation function for endpoint parameter <uuid>"""
     try:
@@ -22,20 +48,20 @@ def validate_uuid(uuid: str):
         raise ValueError("The provided URL does not contain a valid UUID") from error
 
 
-def validate_literature_reference(literatur_reference: LiteratureReference):
-    """Validate literatur reference to not include empty strings."""
-    if literatur_reference.citation == "":
+def validate_literature_reference(literature_reference: LiteratureReference):
+    """Validate literature reference to not include empty strings."""
+    if literature_reference.citation == "":
         raise ValueError(
             "Literature reference JSON fields "
-            "should not contain empty strings as values."
+            + "should not contain empty strings as values."
         )
-    if literatur_reference.img_src == "":
+    if literature_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 == "":
+    if literature_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."
         )

tests/integration/test_database_client_flask.py

@@ -5,7 +5,9 @@
 from psycopg2.extensions import connection
 
 from sketch_map_tool.database import client_flask
-from sketch_map_tool.exceptions import CustomFileNotFoundError
+from sketch_map_tool.exceptions import (
+    CustomFileNotFoundError,
+)
 
 
 def test_open_close_connection(flask_app):

tests/unit/test_config.py

@@ -27,6 +27,7 @@ def config_keys():
         "wms-layers",
         "wms-read-timeout",
         "max-nr-simultaneous-uploads",
+        "max_pixel_per_image",
     )
 
 
@@ -93,8 +94,12 @@ def test_get_config(config_keys):
 def test_get_config_value(config_keys):
     for key in config_keys:
         val = config.get_config_value(key)
-        if key in ["wms-read-timeout", "max-nr-simultaneous-uploads"]:
-            assert isinstance(val, int)
+        if key in [
+            "wms-read-timeout",
+            "max-nr-simultaneous-uploads",
+            "max_pixel_per_image",
+        ]:
+            assert isinstance(val, int | float)
         else:
             assert isinstance(val, str)
 

tests/unit/test_validators.py

@@ -1,6 +1,11 @@
 import pytest
 
-from sketch_map_tool.validators import validate_type, validate_uuid
+from sketch_map_tool.exceptions import UploadLimitsExceededError
+from sketch_map_tool.validators import (
+    validate_type,
+    validate_uploaded_sketchmaps,
+    validate_uuid,
+)
 
 
 @pytest.mark.parametrize(
@@ -17,6 +22,12 @@ def test_validate_type_invalid(type_):
         validate_type(type_)
 
 
+def test_validate_file_by_pixel_count(files, monkeypatch):
+    with pytest.raises(UploadLimitsExceededError):
+        monkeypatch.setenv("MAX-PIXEL-PER-IMAGE", "10")
+        validate_uploaded_sketchmaps(files)
+
+
 def test_validate_uui(uuid):
     validate_uuid(uuid)