Проект разделен на два репозитория
This commit is contained in:
mi
@@ -0,0 +1,171 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import io
|
||||
import json
|
||||
import struct
|
||||
from collections.abc import AsyncIterator
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Protocol
|
||||
|
||||
from PIL import Image, UnidentifiedImageError
|
||||
from pillow_heif import register_heif_opener
|
||||
|
||||
from app.contracts import Attachment
|
||||
|
||||
register_heif_opener()
|
||||
|
||||
|
||||
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]) -> str | None: ...
|
||||
|
||||
|
||||
@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 ClamAvInstream:
|
||||
def __init__(self, host: str, port: int, timeout: float = 45.0) -> None:
|
||||
self.host, self.port, self.timeout = host, port, timeout
|
||||
|
||||
async def scan(self, chunks: AsyncIterator[bytes]) -> str | None:
|
||||
async def operation() -> str | None:
|
||||
reader, writer = await asyncio.open_connection(self.host, self.port)
|
||||
try:
|
||||
writer.write(b"zINSTREAM\0")
|
||||
async for chunk in chunks:
|
||||
writer.write(struct.pack(">I", len(chunk)) + chunk)
|
||||
await writer.drain()
|
||||
writer.write(struct.pack(">I", 0))
|
||||
await writer.drain()
|
||||
result = await reader.readuntil(b"\0")
|
||||
text = result.rstrip(b"\0").decode("utf-8", "replace")
|
||||
if text.endswith(" OK"):
|
||||
return None
|
||||
if text.endswith(" FOUND"):
|
||||
return text.rsplit(": ", 1)[-1].removesuffix(" FOUND")
|
||||
raise DependencyFailure("invalid ClamAV response")
|
||||
finally:
|
||||
writer.close()
|
||||
await writer.wait_closed()
|
||||
|
||||
try:
|
||||
return await asyncio.wait_for(operation(), self.timeout)
|
||||
except (OSError, TimeoutError) as exc:
|
||||
raise DependencyFailure("ClamAV unavailable") from exc
|
||||
|
||||
async def signatures_version(self) -> str:
|
||||
try:
|
||||
reader, writer = await asyncio.wait_for(
|
||||
asyncio.open_connection(self.host, self.port), 2.0
|
||||
)
|
||||
try:
|
||||
writer.write(b"zVERSION\0")
|
||||
await writer.drain()
|
||||
raw = await asyncio.wait_for(reader.readuntil(b"\0"), 2.0)
|
||||
finally:
|
||||
writer.close()
|
||||
await writer.wait_closed()
|
||||
except (OSError, TimeoutError) as exc:
|
||||
raise DependencyFailure("ClamAV unavailable") from exc
|
||||
value = raw.rstrip(b"\0")
|
||||
if not value.startswith(b"ClamAV ") or len(value) > 512:
|
||||
raise DependencyFailure("invalid ClamAV version response")
|
||||
return "sha256:" + hashlib.sha256(value).hexdigest()
|
||||
|
||||
|
||||
async def one_chunk(data: bytes) -> AsyncIterator[bytes]:
|
||||
yield data
|
||||
Reference in New Issue
Block a user