mi
74 lines
2.4 KiB
Python
74 lines
2.4 KiB
Python
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)
|