mi
240 lines
8.3 KiB
Python
240 lines
8.3 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import hashlib
|
|
import io
|
|
import json
|
|
import struct
|
|
from collections.abc import AsyncIterator
|
|
from dataclasses import dataclass
|
|
from datetime import UTC, datetime
|
|
from email.utils import parsedate_to_datetime
|
|
from pathlib import Path
|
|
from typing import Protocol
|
|
|
|
from opentelemetry import trace
|
|
from PIL import Image, UnidentifiedImageError
|
|
from pillow_heif import register_heif_opener
|
|
|
|
from app.contracts import Attachment
|
|
|
|
register_heif_opener()
|
|
tracer = trace.get_tracer("message-safety.dependencies")
|
|
|
|
|
|
class ObjectChanged(RuntimeError):
|
|
pass
|
|
|
|
|
|
class DependencyFailure(RuntimeError):
|
|
pass
|
|
|
|
|
|
class ObjectReader(Protocol):
|
|
async def stream(self, attachment: Attachment) -> AsyncIterator[bytes]: ...
|
|
|
|
|
|
class Antivirus(Protocol):
|
|
async def scan(
|
|
self, chunks: AsyncIterator[bytes], *, scan_timeout: float
|
|
) -> AntivirusScanResult: ...
|
|
|
|
async def status(self) -> AntivirusStatus: ...
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class AntivirusScanResult:
|
|
threat: str | None
|
|
signatures_version: str
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class AntivirusStatus:
|
|
engine_version: str
|
|
signatures_version: str
|
|
databases_date: datetime
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class DetectorManifest:
|
|
version: str
|
|
supported: frozenset[str]
|
|
max_size: int
|
|
|
|
@classmethod
|
|
def load(cls, path: Path) -> DetectorManifest:
|
|
raw = path.read_bytes()
|
|
data = json.loads(raw)
|
|
version = "sha256:" + hashlib.sha256(raw).hexdigest()
|
|
return cls(
|
|
version, frozenset(data["supported_mime_types"]), data["hard_limits"]["max_size_bytes"]
|
|
)
|
|
|
|
|
|
def validate_metadata(
|
|
attachment: Attachment, manifest: DetectorManifest, enabled: set[str]
|
|
) -> str | None:
|
|
if attachment.mime_type not in manifest.supported or attachment.mime_type not in enabled:
|
|
return "file.unsupported_mime"
|
|
if attachment.size_bytes > manifest.max_size:
|
|
return "file.size_limit"
|
|
return None
|
|
|
|
|
|
def detect_format(data: bytes, declared: str) -> str | None:
|
|
matches: list[str] = []
|
|
if data.startswith(b"\xff\xd8\xff") and data.endswith(b"\xff\xd9"):
|
|
matches.append("image/jpeg")
|
|
if data.startswith(b"\x89PNG\r\n\x1a\n") and b"IEND" in data[-64:]:
|
|
matches.append("image/png")
|
|
if len(data) >= 12 and data[:4] == b"RIFF" and data[8:12] == b"WEBP":
|
|
matches.append("image/webp")
|
|
if len(data) >= 12 and data[4:8] == b"ftyp":
|
|
brand = data[8:12]
|
|
if brand in {b"heic", b"heix", b"hevc", b"hevx"}:
|
|
matches.append("image/heic")
|
|
if brand in {b"mif1", b"msf1"}:
|
|
matches.append("image/heif")
|
|
if data.startswith(b"%PDF-") and b"%%EOF" in data[-1024:]:
|
|
matches.append("application/pdf")
|
|
if len(matches) != 1:
|
|
return "file.polyglot_or_ambiguous"
|
|
if matches[0] != declared:
|
|
return "file.format_mismatch"
|
|
if declared == "application/pdf":
|
|
lowered = data.lower()
|
|
if b"/encrypt" in lowered:
|
|
return "file.encrypted_content"
|
|
if any(
|
|
token in lowered
|
|
for token in (b"/javascript", b"/openaction", b"/launch", b"/xfa", b"/embeddedfile")
|
|
):
|
|
return "file.active_content"
|
|
else:
|
|
try:
|
|
with Image.open(io.BytesIO(data)) as image:
|
|
width, height = image.size
|
|
if width > 10_000 or height > 10_000 or width * height > 25_000_000:
|
|
return "file.parser_limit"
|
|
image.verify()
|
|
except (UnidentifiedImageError, OSError, ValueError):
|
|
return "file.format_mismatch"
|
|
return None
|
|
|
|
|
|
async def collect_and_hash(
|
|
reader: ObjectReader, attachment: Attachment, *, max_size: int
|
|
) -> tuple[bytes, bytes]:
|
|
digest = hashlib.sha256()
|
|
body = bytearray()
|
|
async for chunk in reader.stream(attachment):
|
|
if len(body) + len(chunk) > max_size:
|
|
raise ObjectChanged("object exceeds bounded size")
|
|
digest.update(chunk)
|
|
body.extend(chunk)
|
|
expected = bytes.fromhex(attachment.checksum.removeprefix("sha256:"))
|
|
if len(body) != attachment.size_bytes or digest.digest() != expected:
|
|
raise ObjectChanged("authoritative object metadata mismatch")
|
|
return bytes(body), digest.digest()
|
|
|
|
|
|
class KeslSocketScanner:
|
|
MAX_HEADER = 4096
|
|
MAX_RESPONSE = 16 * 1024
|
|
|
|
def __init__(self, socket_path: Path, timeout: float = 60.0) -> None:
|
|
self.socket_path = socket_path
|
|
self.timeout = timeout
|
|
|
|
async def _request(
|
|
self,
|
|
header: dict[str, object],
|
|
body: bytes = b"",
|
|
*,
|
|
request_timeout: float | None = None,
|
|
) -> dict[str, object]:
|
|
encoded = json.dumps(header, separators=(",", ":")).encode()
|
|
if len(encoded) > self.MAX_HEADER:
|
|
raise DependencyFailure("KESL request header is too large")
|
|
|
|
async def operation() -> dict[str, object]:
|
|
reader, writer = await asyncio.open_unix_connection(str(self.socket_path))
|
|
try:
|
|
writer.write(struct.pack(">I", len(encoded)) + encoded + body)
|
|
await writer.drain()
|
|
(size,) = struct.unpack(">I", await reader.readexactly(4))
|
|
if size < 2 or size > self.MAX_RESPONSE:
|
|
raise DependencyFailure("invalid KESL broker response size")
|
|
value = json.loads((await reader.readexactly(size)).decode())
|
|
if not isinstance(value, dict):
|
|
raise DependencyFailure("invalid KESL broker response")
|
|
return value
|
|
finally:
|
|
writer.close()
|
|
await writer.wait_closed()
|
|
|
|
try:
|
|
return await asyncio.wait_for(operation(), request_timeout or self.timeout)
|
|
except (OSError, TimeoutError, asyncio.IncompleteReadError, json.JSONDecodeError) as exc:
|
|
raise DependencyFailure("KESL unavailable") from exc
|
|
|
|
async def scan(
|
|
self, chunks: AsyncIterator[bytes], *, scan_timeout: float
|
|
) -> AntivirusScanResult:
|
|
body = bytearray()
|
|
async for chunk in chunks:
|
|
body.extend(chunk)
|
|
with tracer.start_as_current_span(
|
|
"message_safety.antivirus.scan",
|
|
attributes={"message_safety.antivirus.engine": "kesl"},
|
|
):
|
|
value = await self._request(
|
|
{"op": "scan", "size": len(body)},
|
|
bytes(body),
|
|
request_timeout=scan_timeout,
|
|
)
|
|
if value.get("status") != "completed" or value.get("verdict") not in {
|
|
"clean",
|
|
"infected",
|
|
}:
|
|
raise DependencyFailure("KESL scan failed")
|
|
signatures = value.get("signatures_version")
|
|
if not isinstance(signatures, str) or not signatures.startswith("sha256:"):
|
|
raise DependencyFailure("KESL signatures version is missing")
|
|
threat = value.get("threat")
|
|
if value["verdict"] == "infected" and not isinstance(threat, str):
|
|
raise DependencyFailure("KESL infected verdict has no threat")
|
|
return AntivirusScanResult(
|
|
threat=threat if isinstance(threat, str) else None,
|
|
signatures_version=signatures,
|
|
)
|
|
|
|
async def status(self) -> AntivirusStatus:
|
|
with tracer.start_as_current_span(
|
|
"message_safety.antivirus.status",
|
|
attributes={"message_safety.antivirus.engine": "kesl"},
|
|
):
|
|
value = await self._request({"op": "status"})
|
|
if value.get("status") != "ready":
|
|
raise DependencyFailure("KESL is not ready")
|
|
try:
|
|
database_date = datetime.fromisoformat(
|
|
str(value["databases_date"]).replace("Z", "+00:00")
|
|
)
|
|
except ValueError:
|
|
try:
|
|
database_date = parsedate_to_datetime(str(value["databases_date"]))
|
|
except (TypeError, ValueError) as exc:
|
|
raise DependencyFailure("invalid KESL database date") from exc
|
|
if database_date.tzinfo is None:
|
|
database_date = database_date.replace(tzinfo=UTC)
|
|
engine = value.get("engine_version")
|
|
signatures = value.get("signatures_version")
|
|
if not isinstance(engine, str) or not isinstance(signatures, str):
|
|
raise DependencyFailure("incomplete KESL status")
|
|
return AntivirusStatus(engine, signatures, database_date.astimezone(UTC))
|
|
|
|
|
|
async def one_chunk(data: bytes) -> AsyncIterator[bytes]:
|
|
yield data
|