Fix lazy file validation imports

Author: daryllimyt

tests/unit/test_storage_validation.py

@@ -3,7 +3,9 @@
 
 import pytest
 
+import tracecat.storage.validation as storage_validation
 from tracecat.storage.exceptions import (
+    FileContentMismatchError,
     FileExtensionError,
     FileMimeTypeError,
     FileNameError,
@@ -109,6 +111,32 @@ def test_zip_upload_passes(validator_basic: FileSecurityValidator) -> None:
     assert validation_result.content_type == "application/zip"
 
 
+def test_missing_polyfile_dependency_raises_generic_message(
+    monkeypatch: pytest.MonkeyPatch,
+    validator_basic: FileSecurityValidator,
+) -> None:
+    monkeypatch.setattr(storage_validation, "_MAGIC_MATCHER", None)
+    real_import_module = storage_validation.importlib.import_module
+
+    def fake_import_module(name: str, package: str | None = None) -> object:
+        if name == "polyfile.magic":
+            raise ModuleNotFoundError("No module named 'chardet.universaldetector'")
+        return real_import_module(name, package)
+
+    monkeypatch.setattr(
+        storage_validation.importlib, "import_module", fake_import_module
+    )
+
+    with pytest.raises(FileContentMismatchError) as exc:
+        validator_basic.validate_file(
+            content=pdf_bytes(),
+            filename="doc.pdf",
+            declared_mime_type="application/pdf",
+        )
+
+    assert str(exc.value) == "Unknown or unsupported file type"
+
+
 @pytest.mark.parametrize(
     "filename",
     [

tracecat/storage/validation.py

@@ -14,10 +14,11 @@
 https://cheatsheetseries.owasp.org/cheatsheets/File_Upload_Cheat_Sheet.html#list-allowed-extensions
 """
 
+import importlib
 import os
 from collections.abc import Sequence
+from typing import Any
 
-from polyfile.magic import MagicMatcher
 from pydantic import BaseModel
 
 from tracecat import config
@@ -30,6 +31,16 @@
     FileSizeError,
 )
 
+_MAGIC_MATCHER: Any | None = None
+
+
+def _get_magic_matcher() -> Any:
+    global _MAGIC_MATCHER
+    if _MAGIC_MATCHER is None:
+        polyfile_magic = importlib.import_module("polyfile.magic")
+        _MAGIC_MATCHER = polyfile_magic.MagicMatcher.DEFAULT_INSTANCE
+    return _MAGIC_MATCHER
+
 
 class FileValidationResult(BaseModel):
     """Result of file validation."""
@@ -207,7 +218,7 @@ def _validate_magic_number(self, content: bytes, declared_mime_type: str) -> boo
         detected_equiv: set[str] = set()
 
         try:
-            matches = list(MagicMatcher.DEFAULT_INSTANCE.match(content))
+            matches = list(_get_magic_matcher().match(content))
             for match in matches:
                 mimetypes = getattr(match, "mimetypes", None)
                 if mimetypes: