PacketStation/server/packetstation/domain/scan.py

"""Virus scan through clamd (ClamAV), optional. Off when CLAMAV_HOST is empty.

The file is streamed to clamd with zINSTREAM: the command, then chunks each led by a 4-byte big-endian length, then a
zero length. clamd answers `stream: OK` or `stream: <signature> FOUND`. When a scanner is configured and cannot be
reached, the upload is refused (fail closed).
"""

from __future__ import annotations

import logging
import socket
import struct
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from ..api.errors import ApiError

log = logging.getLogger(__name__)

CHUNK = 256 * 1024
REPLY_MAX = 4096
MSG_INFECTED = "That file failed the virus scan and was not stored."
MSG_UNAVAILABLE = "The virus scanner is not reachable. Try again in a minute."


class ScanError(Exception):
    """The scanner could not be reached or answered badly."""


@dataclass(frozen=True)
class ScanResult:
    clean: bool
    signature: str | None


class ClamdScanner:
    def __init__(self, host: str, port: int = 3310, timeout: float = 30.0):
        self.host = host
        self.port = int(port)
        self.timeout = float(timeout)

    def _connect(self, timeout: float | None = None) -> socket.socket:
        try:
            return socket.create_connection((self.host, self.port), timeout=timeout or self.timeout)
        except OSError as e:
            raise ScanError(f"can't reach clamd at {self.host}:{self.port}: {type(e).__name__}") from None

    @staticmethod
    def _reply(sock: socket.socket) -> str:
        """Read one NUL-terminated reply (or up to the close)."""
        buf = b""
        while b"\0" not in buf and len(buf) < REPLY_MAX:
            part = sock.recv(1024)
            if not part:
                break
            buf += part
        return buf.split(b"\0", 1)[0].decode("utf-8", "replace").strip()

    def ping(self, timeout: float | None = None) -> bool:
        """True when clamd answers PING with PONG."""
        try:
            with self._connect(timeout) as sock:
                sock.sendall(b"zPING\0")
                return self._reply(sock) == "PONG"
        except (ScanError, OSError) as e:
            log.info("clamd ping failed: %s", type(e).__name__)
            return False

    def scan_path(self, path: Path) -> ScanResult:
        """Stream the file to clamd. Raises ScanError when clamd is down or its answer is not OK or FOUND."""
        try:
            with self._connect() as sock, Path(path).open("rb") as f:
                sock.sendall(b"zINSTREAM\0")
                while True:
                    chunk = f.read(CHUNK)
                    if not chunk:
                        break
                    sock.sendall(struct.pack(">I", len(chunk)) + chunk)
                sock.sendall(struct.pack(">I", 0))
                reply = self._reply(sock)
        except ScanError:
            raise
        except OSError as e:
            raise ScanError(f"clamd scan failed: {type(e).__name__}") from None
        return parse_reply(reply)


def parse_reply(reply: str) -> ScanResult:
    """`stream: OK` is clean; `stream: <sig> FOUND` is a hit; anything else is a ScanError."""
    text = reply.strip()
    if text.startswith("stream:"):
        text = text[len("stream:") :].strip()
    if text == "OK":
        return ScanResult(clean=True, signature=None)
    if text.endswith(" FOUND"):
        sig = text[: -len(" FOUND")].strip()[:200]
        return ScanResult(clean=False, signature=sig or "unknown")
    raise ScanError(f"clamd answered: {text[:200]}")


def scanner_for(settings: Any) -> ClamdScanner | None:
    """The configured scanner, or None when CLAMAV_HOST is empty."""
    host = str(getattr(settings, "clamav_host", "") or "").strip()
    if not host:
        return None
    return ClamdScanner(
        host,
        int(getattr(settings, "clamav_port", 3310)),
        float(getattr(settings, "clamav_timeout_seconds", 30.0)),
    )


def check(settings: Any, path: Path) -> ScanResult | None:
    """Scan a complete file. None when no scanner is configured.

    Raises 422 infected (with the signature) on a hit and 503 scan_unavailable when the scanner fails.
    """
    scanner = scanner_for(settings)
    if scanner is None:
        return None
    try:
        result = scanner.scan_path(path)
    except ScanError as e:
        log.warning("virus scan unavailable: %s", e)
        raise ApiError(503, "scan_unavailable", MSG_UNAVAILABLE) from None
    if not result.clean:
        log.warning("upload refused by the virus scan: %s", result.signature)
        raise ApiError(422, "infected", MSG_INFECTED, extra={"signature": result.signature})
    return result