[5620] fix(api): classify plain ZIP attachments correctly

api/oss/src/core/sessions/attachments/media.py

@@ -1,12 +1,12 @@
-from typing import Optional
+from io import BytesIO
+from zipfile import BadZipFile, ZipFile
 
 import puremagic
 from puremagic.main import PureError
 
 from oss.src.core.sessions.attachments.dtos import AttachmentKind, AttachmentMedia
 from oss.src.core.sessions.attachments.types import AttachmentInvalid
 
-
 _NATIVE_IMAGE_TYPES = {
     "image/gif",
     "image/jpeg",
@@ -43,9 +43,13 @@
     b"V_VP9",
 )
 _CONTAINER_SCAN_BYTES = 1024 * 1024
+_ZIP_MEDIA_TYPE = "application/zip"
+_DOCX_MEDIA_TYPE = (
+    "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
+)
 
 
-def _sniff(data: bytes) -> Optional[str]:
+def _sniff(data: bytes) -> str | None:
     """The only puremagic call; classification is policy, not a security boundary."""
     try:
         media_type = puremagic.from_string(data, mime=True)
@@ -71,14 +75,39 @@ def _is_m4a(*, data: bytes) -> bool:
     return any(brand in _M4A_BRANDS for brand in brands)
 
 
+def _canonical_zip_media_type(
+    *,
+    data: bytes,
+    inspected_media_type: str | None,
+) -> str | None:
+    if inspected_media_type not in {_ZIP_MEDIA_TYPE, _DOCX_MEDIA_TYPE}:
+        return inspected_media_type
+
+    try:
+        with ZipFile(BytesIO(data)) as archive:
+            names = set(archive.namelist())
+    except (BadZipFile, ValueError):
+        return inspected_media_type
+
+    if "[Content_Types].xml" in names and "word/document.xml" in names:
+        return _DOCX_MEDIA_TYPE
+
+    return _ZIP_MEDIA_TYPE
+
+
 def _canonical_container_media_type(
     *,
     data: bytes,
-    inspected_media_type: Optional[str],
-) -> Optional[str]:
+    inspected_media_type: str | None,
+) -> str | None:
     if _is_m4a(data=data):
         return "audio/mp4"
 
+    inspected_media_type = _canonical_zip_media_type(
+        data=data,
+        inspected_media_type=inspected_media_type,
+    )
+
     scan = data[:_CONTAINER_SCAN_BYTES]
     if scan.startswith(b"OggS"):
         if _contains_any(data=scan, signatures=_OGG_VIDEO_CODECS):
@@ -109,9 +138,7 @@ def _kind_for(*, media_type: str) -> AttachmentKind:
     return AttachmentKind.OTHER
 
 
-def classify(
-    *, data: bytes, declared_media_type: Optional[str] = None
-) -> AttachmentMedia:
+def classify(*, data: bytes, declared_media_type: str | None = None) -> AttachmentMedia:
     if not data:
         raise AttachmentInvalid()
 

api/oss/tests/pytest/unit/sessions/test_attachment_media.py

@@ -2,7 +2,6 @@
 from zipfile import ZipFile
 
 import pytest
-
 from oss.src.core.sessions.attachments.dtos import AttachmentKind
 from oss.src.core.sessions.attachments.media import classify
 from oss.src.core.sessions.attachments.types import AttachmentInvalid
@@ -45,21 +44,78 @@ def test_empty_or_unrecognizable_non_utf8_bytes_are_invalid(data):
         classify(data=data, declared_media_type="application/octet-stream")
 
 
-def test_recognized_zip_is_accepted_as_other():
+def _zip_bytes(*, files):
     buffer = BytesIO()
     with ZipFile(buffer, "w") as archive:
-        archive.writestr("file.txt", "hello")
+        for path, content in files:
+            archive.writestr(path, content)
+    return buffer.getvalue()
+
+
+def test_plain_zip_is_classified_as_zip():
+    data = _zip_bytes(files=[("file.txt", "hello")])
 
     result = classify(
-        data=buffer.getvalue(),
+        data=data,
         declared_media_type="application/zip",
     )
 
-    assert result.media_type.startswith("application/")
+    assert result.media_type == "application/zip"
     assert result.kind == AttachmentKind.OTHER
     assert result.native_image is False
 
 
+def test_plain_zip_structure_overrides_declared_docx_type():
+    result = classify(
+        data=_zip_bytes(files=[("notes.txt", "hello")]),
+        declared_media_type=(
+            "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
+        ),
+    )
+
+    assert result.media_type == "application/zip"
+    assert result.kind == AttachmentKind.OTHER
+
+
+@pytest.mark.parametrize(
+    "declared_media_type",
+    [
+        "application/zip",
+        "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
+    ],
+)
+def test_docx_structure_is_classified_as_docx(declared_media_type):
+    result = classify(
+        data=_zip_bytes(
+            files=[
+                ("[Content_Types].xml", "<Types />"),
+                ("word/document.xml", "<document />"),
+            ]
+        ),
+        declared_media_type=declared_media_type,
+    )
+
+    assert result.media_type == (
+        "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
+    )
+    assert result.kind == AttachmentKind.OTHER
+
+
+def test_malformed_zip_signature_keeps_inspected_type(monkeypatch):
+    monkeypatch.setattr(
+        "oss.src.core.sessions.attachments.media.puremagic.from_string",
+        lambda *_args, **_kwargs: "application/zip",
+    )
+
+    result = classify(
+        data=b"PK\x03\x04not-a-valid-archive",
+        declared_media_type="application/zip",
+    )
+
+    assert result.media_type == "application/zip"
+    assert result.kind == AttachmentKind.OTHER
+
+
 def test_svg_is_a_workspace_document_not_a_native_image():
     result = classify(
         data=b'<svg xmlns="http://www.w3.org/2000/svg"></svg>',