Merge pull request #2045 from kili-technology/feature/lab-4341-aakd-i-need-to-validate-file-extensions-in-import-service

feat(lab-4341): check file type before importing

Author: florianlega

src/kili/services/asset_import/__init__.py

@@ -63,4 +63,5 @@ def import_assets(  # pylint: disable=too-many-arguments
     asset_importer = importer_by_type[input_type](*importer_params)
     casted_assets = cast(list[AssetLike], assets)
     asset_importer.check_asset_contents(casted_assets)
+    asset_importer.check_file_extensions(casted_assets, input_type)
     return asset_importer.import_assets(assets=casted_assets, input_type=input_type)

src/kili/services/asset_import/base.py

@@ -4,12 +4,13 @@
 import logging
 import mimetypes
 import os
+import urllib.parse
 import warnings
 from collections.abc import Callable
 from concurrent.futures import ThreadPoolExecutor
 from itertools import repeat
 from json import dumps, loads
-from pathlib import Path
+from pathlib import Path, PurePosixPath
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -34,6 +35,7 @@
 from kili.domain.project import InputType, ProjectId
 from kili.domain.types import ListOrTuple
 from kili.services.asset_import.constants import (
+    ALLOWED_EXTENSIONS_BY_INPUT_TYPE,
     IMPORT_BATCH_SIZE,
     project_compatible_mimetypes,
 )
@@ -524,6 +526,45 @@ def check_asset_contents(assets: list[AssetLike]) -> None:
                     " multi_layer_content and empty json_content"
                 )
 
+    @staticmethod
+    def check_file_extensions(assets: list[AssetLike], input_type: str) -> None:
+        """Validate file extensions against allowed extensions for the given input type.
+
+        Skips validation when content has no extension (e.g. raw text, extensionless URLs).
+        Raises ImportValidationError on the first asset with a disallowed extension.
+        """
+        allowed = ALLOWED_EXTENSIONS_BY_INPUT_TYPE.get(input_type)
+        if allowed is None:
+            return
+
+        def get_ext(path_or_url: str) -> str:
+            path = urllib.parse.urlparse(path_or_url).path if is_url(path_or_url) else path_or_url
+            return PurePosixPath(path).suffix.lower()
+
+        for asset in assets:
+            content = asset.get("content")
+            if isinstance(content, str):
+                ext = get_ext(content)
+                if ext and ext not in allowed:
+                    raise ImportValidationError(
+                        f"File extension '{ext}' is not allowed for {input_type} projects"
+                        f" (asset external_id='{asset.get('external_id', 'unknown')}')."
+                        f" Allowed extensions: {', '.join(sorted(allowed))}"
+                    )
+
+            multi_layer = asset.get("multi_layer_content")
+            if multi_layer:
+                for layer in multi_layer:
+                    path = layer.get("path") or layer.get("url") or layer.get("content", "")
+                    if path:
+                        ext = get_ext(path)
+                        if ext and ext not in allowed:
+                            raise ImportValidationError(
+                                f"File extension '{ext}' is not allowed for {input_type} projects"
+                                f" (asset external_id='{asset.get('external_id', 'unknown')}')."
+                                f" Allowed extensions: {', '.join(sorted(allowed))}"
+                            )
+
     def _can_upload_from_local_data(self) -> bool:
         user_me = self.kili.kili_api_gateway.get_current_user(fields=("email",))
         options = QueryOptions(first=1, disable_tqdm=True)

src/kili/services/asset_import/constants.py

@@ -10,3 +10,28 @@
 MB_SIZE = 1024**2
 LARGE_IMAGE_THRESHOLD_SIZE = 30 * MB_SIZE
 MAX_WIDTH_OR_HEIGHT_NON_TILED = 10000
+
+ALLOWED_EXTENSIONS_BY_INPUT_TYPE: dict[str, frozenset[str]] = {
+    "AUDIO": frozenset({".flac", ".mp3", ".mp4", ".wav"}),
+    "GEOSPATIAL": frozenset({".tif", ".tiff", ".jp2", ".ntf", ".nitf"}),
+    "IMAGE": frozenset(
+        {
+            ".jpeg",
+            ".jpg",
+            ".png",
+            ".bmp",
+            ".gif",
+            ".webp",
+            ".ico",
+            ".tif",
+            ".tiff",
+            ".jp2",
+            ".ntf",
+            ".nitf",
+        }
+    ),
+    "LLM_RLHF": frozenset({".json"}),
+    "PDF": frozenset({".pdf"}),
+    "TEXT": frozenset({".txt", ".csv"}),
+    "VIDEO": frozenset({".mp4", ".mkv", ".3gp", ".avi", ".m4v", ".mov", ".webm"}),
+}

tests/unit/services/asset_import/test_import_common.py

@@ -6,7 +6,7 @@
 from kili.core.graphql.operations.asset.mutations import GQL_APPEND_MANY_ASSETS
 from kili.domain.project import ProjectId
 from kili.services.asset_import import import_assets
-from kili.services.asset_import.exceptions import MimeTypeError
+from kili.services.asset_import.exceptions import ImportValidationError
 from tests.unit.services.asset_import.base import ImportTestCase
 from tests.unit.services.asset_import.mocks import (
     mocked_request_signed_urls,
@@ -20,11 +20,12 @@
 @patch("kili.utils.bucket.upload_data_via_rest", mocked_upload_data_via_rest)
 class TestContentType(ImportTestCase):
     def test_cannot_upload_an_image_to_video_project(self, *_):
-        self.kili.kili_api_gateway.get_project.return_value = {"inputType": "VIDEO_LEGACY"}
+        self.kili.kili_api_gateway.get_project.return_value = {"inputType": "VIDEO"}
         url = "https://storage.googleapis.com/label-public-staging/car/car_1.jpg"
         path_image = self.downloader(url)
         assets = [{"content": path_image, "external_id": "image"}]
-        with pytest.raises(MimeTypeError):
+        # Extension check runs before MIME type check, so ImportValidationError is raised first
+        with pytest.raises(ImportValidationError):
             import_assets(self.kili, ProjectId(self.project_id), assets, disable_tqdm=True)
 
     def test_cannot_import_files_not_found_to_an_image_project(self, *_):
@@ -84,3 +85,184 @@ def test_import_assets_verify(self, mocked_verify_batch_imported, *_):
         mocked_verify_batch_imported.assert_not_called()
         import_assets(self.kili, ProjectId("project_id"), assets, verify=True)
         mocked_verify_batch_imported.assert_called_once()
+
+
+@patch("kili.utils.bucket.generate_unique_id", mocked_unique_id)
+@patch("kili.utils.bucket.request_signed_urls", mocked_request_signed_urls)
+@patch("kili.utils.bucket.upload_data_via_rest", mocked_upload_data_via_rest)
+class TestFileExtensionValidation(ImportTestCase):
+    """Tests that the import service validates file extensions before uploading."""
+
+    # --- IMAGE ---
+    def test_image_project_rejects_video_extension(self, *_):
+        self.kili.kili_api_gateway.get_project.return_value = {"inputType": "IMAGE"}
+        assets = [{"content": "https://example.com/clip.mp4", "external_id": "wrong"}]
+        with pytest.raises(ImportValidationError, match=r"\.mp4"):
+            import_assets(self.kili, ProjectId(self.project_id), assets)
+
+    def test_image_project_rejects_audio_extension(self, *_):
+        self.kili.kili_api_gateway.get_project.return_value = {"inputType": "IMAGE"}
+        assets = [{"content": "https://example.com/sound.mp3", "external_id": "wrong"}]
+        with pytest.raises(ImportValidationError, match=r"\.mp3"):
+            import_assets(self.kili, ProjectId(self.project_id), assets)
+
+    def test_image_project_accepts_jpg_extension(self, *_):
+        self.kili.kili_api_gateway.get_project.return_value = {"inputType": "IMAGE"}
+        assets = [{"content": "https://example.com/image.jpg", "external_id": "ok", "id": "uid"}]
+        # Should not raise ImportValidationError
+        import_assets(self.kili, ProjectId(self.project_id), assets)
+
+    def test_image_project_accepts_tif_extension(self, *_):
+        self.kili.kili_api_gateway.get_project.return_value = {"inputType": "IMAGE"}
+        assets = [{"content": "https://example.com/geo.tif", "external_id": "ok", "id": "uid"}]
+        import_assets(self.kili, ProjectId(self.project_id), assets)
+
+    # --- VIDEO ---
+    def test_video_project_rejects_image_extension(self, *_):
+        self.kili.kili_api_gateway.get_project.return_value = {"inputType": "VIDEO"}
+        assets = [{"content": "https://example.com/photo.jpg", "external_id": "wrong"}]
+        with pytest.raises(ImportValidationError, match=r"\.jpg"):
+            import_assets(self.kili, ProjectId(self.project_id), assets)
+
+    def test_video_project_rejects_pdf_extension(self, *_):
+        self.kili.kili_api_gateway.get_project.return_value = {"inputType": "VIDEO"}
+        assets = [{"content": "https://example.com/doc.pdf", "external_id": "wrong"}]
+        with pytest.raises(ImportValidationError, match=r"\.pdf"):
+            import_assets(self.kili, ProjectId(self.project_id), assets)
+
+    def test_video_project_accepts_mp4_extension(self, *_):
+        self.kili.kili_api_gateway.get_project.return_value = {"inputType": "VIDEO"}
+        assets = [{"content": "https://example.com/vid.mp4", "external_id": "ok"}]
+        import_assets(self.kili, ProjectId(self.project_id), assets)
+
+    def test_video_project_accepts_mkv_extension(self, *_):
+        self.kili.kili_api_gateway.get_project.return_value = {"inputType": "VIDEO"}
+        assets = [{"content": "https://example.com/vid.mkv", "external_id": "ok"}]
+        import_assets(self.kili, ProjectId(self.project_id), assets)
+
+    # --- AUDIO ---
+    def test_audio_project_rejects_image_extension(self, *_):
+        self.kili.kili_api_gateway.get_project.return_value = {"inputType": "AUDIO"}
+        assets = [{"content": "https://example.com/photo.jpg", "external_id": "wrong"}]
+        with pytest.raises(ImportValidationError, match=r"\.jpg"):
+            import_assets(self.kili, ProjectId(self.project_id), assets)
+
+    def test_audio_project_rejects_pdf_extension(self, *_):
+        self.kili.kili_api_gateway.get_project.return_value = {"inputType": "AUDIO"}
+        assets = [{"content": "https://example.com/doc.pdf", "external_id": "wrong"}]
+        with pytest.raises(ImportValidationError, match=r"\.pdf"):
+            import_assets(self.kili, ProjectId(self.project_id), assets)
+
+    def test_audio_project_accepts_mp3_extension(self, *_):
+        self.kili.kili_api_gateway.get_project.return_value = {"inputType": "AUDIO"}
+        assets = [{"content": "https://example.com/audio.mp3", "external_id": "ok"}]
+        import_assets(self.kili, ProjectId(self.project_id), assets)
+
+    def test_audio_project_accepts_wav_extension(self, *_):
+        self.kili.kili_api_gateway.get_project.return_value = {"inputType": "AUDIO"}
+        assets = [{"content": "https://example.com/audio.wav", "external_id": "ok"}]
+        import_assets(self.kili, ProjectId(self.project_id), assets)
+
+    # --- PDF ---
+    def test_pdf_project_rejects_image_extension(self, *_):
+        self.kili.kili_api_gateway.get_project.return_value = {"inputType": "PDF"}
+        assets = [{"content": "https://example.com/photo.jpg", "external_id": "wrong"}]
+        with pytest.raises(ImportValidationError, match=r"\.jpg"):
+            import_assets(self.kili, ProjectId(self.project_id), assets)
+
+    def test_pdf_project_rejects_video_extension(self, *_):
+        self.kili.kili_api_gateway.get_project.return_value = {"inputType": "PDF"}
+        assets = [{"content": "https://example.com/clip.mp4", "external_id": "wrong"}]
+        with pytest.raises(ImportValidationError, match=r"\.mp4"):
+            import_assets(self.kili, ProjectId(self.project_id), assets)
+
+    def test_pdf_project_accepts_pdf_extension(self, *_):
+        self.kili.kili_api_gateway.get_project.return_value = {"inputType": "PDF"}
+        assets = [{"content": "https://example.com/doc.pdf", "external_id": "ok"}]
+        import_assets(self.kili, ProjectId(self.project_id), assets)
+
+    # --- GEOSPATIAL ---
+    def test_geospatial_project_rejects_image_extension(self, *_):
+        self.kili.kili_api_gateway.get_project.return_value = {"inputType": "GEOSPATIAL"}
+        assets = [{"content": "https://example.com/photo.jpg", "external_id": "wrong"}]
+        with pytest.raises(ImportValidationError, match=r"\.jpg"):
+            import_assets(self.kili, ProjectId(self.project_id), assets)
+
+    def test_geospatial_project_accepts_tif_extension(self, *_):
+        self.kili.kili_api_gateway.get_project.return_value = {"inputType": "GEOSPATIAL"}
+        assets = [{"content": "https://example.com/geo.tif", "external_id": "ok"}]
+        import_assets(self.kili, ProjectId(self.project_id), assets)
+
+    def test_geospatial_project_accepts_jp2_extension(self, *_):
+        self.kili.kili_api_gateway.get_project.return_value = {"inputType": "GEOSPATIAL"}
+        assets = [{"content": "https://example.com/geo.jp2", "external_id": "ok"}]
+        import_assets(self.kili, ProjectId(self.project_id), assets)
+
+    def test_geospatial_project_rejects_wrong_extension_in_multi_layer(self, *_):
+        self.kili.kili_api_gateway.get_project.return_value = {"inputType": "GEOSPATIAL"}
+        assets = [
+            {
+                "multi_layer_content": [
+                    {"path": "/local/layer.jpg", "name": "layer1"},
+                ],
+                "external_id": "wrong_multi_layer",
+            }
+        ]
+        with pytest.raises(ImportValidationError, match=r"\.jpg"):
+            import_assets(self.kili, ProjectId(self.project_id), assets)
+
+    def test_geospatial_project_accepts_correct_extension_in_multi_layer(self, *_):
+        self.kili.kili_api_gateway.get_project.return_value = {"inputType": "GEOSPATIAL"}
+        assets = [
+            {
+                "multi_layer_content": [
+                    {"path": "/local/layer.tif", "name": "layer1"},
+                    {"path": "/local/layer2.tiff", "name": "layer2"},
+                ],
+                "external_id": "ok_multi_layer",
+            }
+        ]
+        # Should not raise ImportValidationError (may raise later on file access)
+        with pytest.raises(Exception) as exc_info:
+            import_assets(self.kili, ProjectId(self.project_id), assets)
+        assert not isinstance(exc_info.value, ImportValidationError)
+
+    # --- TEXT ---
+    def test_text_project_rejects_video_extension(self, *_):
+        self.kili.kili_api_gateway.get_project.return_value = {"inputType": "TEXT"}
+        assets = [{"content": "https://example.com/clip.mp4", "external_id": "wrong"}]
+        with pytest.raises(ImportValidationError, match=r"\.mp4"):
+            import_assets(self.kili, ProjectId(self.project_id), assets)
+
+    def test_text_project_accepts_txt_extension(self, *_):
+        self.kili.kili_api_gateway.get_project.return_value = {"inputType": "TEXT"}
+        assets = [{"content": "https://example.com/file.txt", "external_id": "ok"}]
+        import_assets(self.kili, ProjectId(self.project_id), assets)
+
+    def test_text_project_skips_validation_for_raw_text(self, *_):
+        self.kili.kili_api_gateway.get_project.return_value = {"inputType": "TEXT"}
+        assets = [{"content": "this is raw text with no extension", "external_id": "ok"}]
+        import_assets(self.kili, ProjectId(self.project_id), assets)
+
+    # --- LLM_RLHF ---
+    def test_llm_project_rejects_non_json_extension(self, *_):
+        self.kili.kili_api_gateway.get_project.return_value = {"inputType": "LLM_RLHF"}
+        assets = [{"content": "https://example.com/data.txt", "external_id": "wrong"}]
+        with pytest.raises(ImportValidationError, match=r"\.txt"):
+            import_assets(self.kili, ProjectId(self.project_id), assets)
+
+    def test_llm_project_accepts_json_extension(self, *_):
+        self.kili.kili_api_gateway.get_project.return_value = {"inputType": "LLM_RLHF"}
+        assets = [{"content": "https://example.com/data.json", "external_id": "ok"}]
+        import_assets(self.kili, ProjectId(self.project_id), assets)
+
+    # --- no extension in URL skips validation ---
+    def test_extensionless_url_skips_validation_for_image_project(self, *_):
+        self.kili.kili_api_gateway.get_project.return_value = {"inputType": "IMAGE"}
+        assets = [{"content": "https://example.com/no-ext", "external_id": "ok", "id": "uid"}]
+        import_assets(self.kili, ProjectId(self.project_id), assets)
+
+    def test_extensionless_url_skips_validation_for_video_project(self, *_):
+        self.kili.kili_api_gateway.get_project.return_value = {"inputType": "VIDEO"}
+        assets = [{"content": "https://hosted-data", "external_id": "ok"}]
+        import_assets(self.kili, ProjectId(self.project_id), assets)