Внедрение KESL на ВМ2 + замена CLAMAV на KESL
This commit is contained in:
mi
@@ -7,6 +7,8 @@ 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
|
||||
|
||||
@@ -33,7 +35,24 @@ class ObjectReader(Protocol):
|
||||
|
||||
|
||||
class Antivirus(Protocol):
|
||||
async def scan(self, chunks: AsyncIterator[bytes]) -> str | None: ...
|
||||
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)
|
||||
@@ -119,65 +138,101 @@ async def collect_and_hash(
|
||||
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
|
||||
class KeslSocketScanner:
|
||||
MAX_HEADER = 4096
|
||||
MAX_RESPONSE = 16 * 1024
|
||||
|
||||
async def scan(self, chunks: AsyncIterator[bytes]) -> str | None:
|
||||
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.clamav.scan",
|
||||
attributes={"server.address.type": "clamav"},
|
||||
"message_safety.antivirus.scan",
|
||||
attributes={"message_safety.antivirus.engine": "kesl"},
|
||||
):
|
||||
return await self._scan(chunks)
|
||||
|
||||
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:
|
||||
with tracer.start_as_current_span("message_safety.clamav.version"):
|
||||
return await self._signatures_version()
|
||||
|
||||
async def _signatures_version(self) -> str:
|
||||
try:
|
||||
reader, writer = await asyncio.wait_for(
|
||||
asyncio.open_connection(self.host, self.port), 2.0
|
||||
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:
|
||||
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()
|
||||
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]:
|
||||
|
||||
Reference in New Issue
Block a user