Проект разделен на два репозитория

VM2_services/codebase/services/message-safety/tests/test_files.py

@@ -0,0 +1,73 @@
+from __future__ import annotations
+
+import hashlib
+import io
+from uuid import UUID
+
+import pytest
+from PIL import Image
+
+from app.contracts import Attachment
+from app.file_pipeline import ObjectChanged, collect_and_hash, detect_format
+
+
+def image_bytes(format_name: str) -> bytes:
+    output = io.BytesIO()
+    Image.new("RGB", (2, 2), "white").save(output, format=format_name)
+    return output.getvalue()
+
+
+@pytest.mark.parametrize(
+    "format_name,mime",
+    [("JPEG", "image/jpeg"), ("PNG", "image/png"), ("WEBP", "image/webp")],
+)
+def test_bounded_image_detector(format_name: str, mime: str) -> None:
+    assert detect_format(image_bytes(format_name), mime) is None
+    assert detect_format(image_bytes(format_name), "application/pdf") == "file.format_mismatch"
+
+
+def test_pdf_active_encrypted_and_malformed() -> None:
+    clean = b"%PDF-1.7\n1 0 obj <<>> endobj\nstartxref\n0\n%%EOF"
+    assert detect_format(clean, "application/pdf") is None
+    assert (
+        detect_format(clean.replace(b"<<>>", b"<</Encrypt 2 0 R>>"), "application/pdf")
+        == "file.encrypted_content"
+    )
+    assert (
+        detect_format(clean.replace(b"<<>>", b"<</JavaScript 2 0 R>>"), "application/pdf")
+        == "file.active_content"
+    )
+    assert detect_format(b"%PDF-1.7 no eof", "application/pdf") == "file.polyglot_or_ambiguous"
+
+
+class Reader:
+    def __init__(self, data: bytes) -> None:
+        self.data = data
+
+    async def stream(self, attachment):
+        yield self.data[:2]
+        yield self.data[2:]
+
+
+def attachment(data: bytes, *, size: int | None = None) -> Attachment:
+    return Attachment(
+        attachment_id=UUID("00000000-0000-4000-8000-000000000003"),
+        quarantine_object_key=(
+            "quarantine/users/00000000-0000-4000-8000-000000000001/"
+            "dialogs/00000000-0000-4000-8000-000000000002/"
+            "00000000-0000-4000-8000-000000000003"
+        ),
+        quarantine_version_id="v1",
+        quarantine_etag='"etag"',
+        mime_type="application/pdf",
+        size_bytes=size or len(data),
+        checksum="sha256:" + hashlib.sha256(data).hexdigest(),
+    )
+
+
+async def test_authoritative_stream_hash_and_size() -> None:
+    data = b"content"
+    body, digest = await collect_and_hash(Reader(data), attachment(data), max_size=100)
+    assert body == data and digest == hashlib.sha256(data).digest()
+    with pytest.raises(ObjectChanged):
+        await collect_and_hash(Reader(data), attachment(data, size=len(data) + 1), max_size=100)