mi
146 lines
4.3 KiB
Python
146 lines
4.3 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import hashlib
|
|
import io
|
|
import json
|
|
import struct
|
|
from pathlib import Path
|
|
from uuid import UUID
|
|
|
|
import pytest
|
|
from PIL import Image
|
|
|
|
from app.contracts import Attachment
|
|
from app.file_pipeline import (
|
|
KeslSocketScanner,
|
|
ObjectChanged,
|
|
collect_and_hash,
|
|
detect_format,
|
|
one_chunk,
|
|
)
|
|
|
|
|
|
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)
|
|
|
|
|
|
class FakeWriter:
|
|
def __init__(self) -> None:
|
|
self.request = bytearray()
|
|
|
|
def write(self, value: bytes) -> None:
|
|
self.request.extend(value)
|
|
|
|
async def drain(self) -> None:
|
|
return None
|
|
|
|
def close(self) -> None:
|
|
return None
|
|
|
|
async def wait_closed(self) -> None:
|
|
return None
|
|
|
|
|
|
def framed(value: dict[str, object]):
|
|
body = json.dumps(value).encode()
|
|
reader = __import__("asyncio").StreamReader()
|
|
reader.feed_data(struct.pack(">I", len(body)) + body)
|
|
reader.feed_eof()
|
|
return reader
|
|
|
|
|
|
async def test_kesl_socket_clean_infected_and_status(monkeypatch, tmp_path: Path) -> None:
|
|
responses = [
|
|
{
|
|
"status": "completed",
|
|
"verdict": "clean",
|
|
"threat": None,
|
|
"engine_version": "12.4",
|
|
"signatures_version": "sha256:" + "a" * 64,
|
|
},
|
|
{
|
|
"status": "completed",
|
|
"verdict": "infected",
|
|
"threat": "EICAR-Test-File",
|
|
"engine_version": "12.4",
|
|
"signatures_version": "sha256:" + "b" * 64,
|
|
},
|
|
{
|
|
"status": "ready",
|
|
"engine_version": "12.4",
|
|
"databases_date": "2026-09-07T11:25:00+00:00",
|
|
"signatures_version": "sha256:" + "c" * 64,
|
|
},
|
|
]
|
|
|
|
async def connect(_):
|
|
return framed(responses.pop(0)), FakeWriter()
|
|
|
|
monkeypatch.setattr(asyncio, "open_unix_connection", connect, raising=False)
|
|
scanner = KeslSocketScanner(tmp_path / "scan.sock")
|
|
clean = await scanner.scan(one_chunk(b"clean"), scan_timeout=1)
|
|
infected = await scanner.scan(one_chunk(b"eicar"), scan_timeout=1)
|
|
status = await scanner.status()
|
|
assert clean.threat is None
|
|
assert infected.threat == "EICAR-Test-File"
|
|
assert status.engine_version == "12.4"
|