# Synapse spam-checker module (Issue #19): scans locally-stored and remote/federated media
# through ClamAV's clamd daemon via its native INSTREAM protocol, before Synapse serves it.
#
# Talks to clamd using only the standard library (asyncio/socket) - no third-party
# dependency, since Synapse's container runs with a read-only root filesystem and this
# module is mounted in via a ConfigMap, not baked into a custom image.
#
# clamd INSTREAM protocol (docs.clamav.net/manual/Usage/ClamdProtocol.html):
#   1. send b"zINSTREAM\0"
#   2. send one or more chunks, each framed as a 4-byte big-endian length + that many
#      data bytes
#   3. send a zero-length chunk (b"\x00\x00\x00\x00") to signal end of stream
#   4. read the reply: "stream: OK\0" (clean) or "stream: <name> FOUND\0" (infected)
#
# Fails open (allows the file through) on any connection/timeout error against clamd,
# so a scanner outage can't take down media uploads for the whole homeserver - logged
# loudly so an outage is still visible in the logs.

import asyncio
import logging
from typing import Any, Union

from synapse.module_api import ModuleApi, NOT_SPAM
from synapse.module_api.errors import Codes

logger = logging.getLogger(__name__)

CHUNK_SIZE = 2**14  # matches ReadableFileWrapper.CHUNK_SIZE
CLAMD_TIMEOUT_SECONDS = 30


class ClamAVSpamChecker:
    def __init__(self, config: dict, api: ModuleApi):
        self.api = api
        self.clamd_host = config["clamd_host"]
        self.clamd_port = config["clamd_port"]
        self.api.register_spam_checker_callbacks(
            check_media_file_for_spam=self.check_media_file_for_spam,
        )

    @staticmethod
    def parse_config(config: dict) -> dict:
        if "clamd_host" not in config or "clamd_port" not in config:
            raise ValueError(
                "clamav_spam_checker config requires 'clamd_host' and 'clamd_port'"
            )
        return config

    async def check_media_file_for_spam(
        self, file_wrapper: Any, file_info: Any
    ) -> Union[Any, Codes, bool]:
        buffer = bytearray()
        await file_wrapper.write_chunks_to(buffer.extend)

        try:
            verdict = await asyncio.wait_for(
                self._scan(bytes(buffer)), timeout=CLAMD_TIMEOUT_SECONDS
            )
        except Exception:
            logger.exception(
                "ClamAV scan failed (clamd at %s:%s unreachable?) - "
                "allowing file through (fail-open)",
                self.clamd_host,
                self.clamd_port,
            )
            return NOT_SPAM

        if verdict is None:
            return NOT_SPAM

        logger.warning("ClamAV rejected an upload: %s", verdict)
        return Codes.FORBIDDEN

    async def _scan(self, data: bytes) -> "str | None":
        """Returns the detected signature name, or None if the file is clean."""
        reader, writer = await asyncio.open_connection(self.clamd_host, self.clamd_port)
        try:
            writer.write(b"zINSTREAM\0")
            for offset in range(0, len(data), CHUNK_SIZE):
                chunk = data[offset : offset + CHUNK_SIZE]
                writer.write(len(chunk).to_bytes(4, "big") + chunk)
            writer.write((0).to_bytes(4, "big"))
            await writer.drain()

            response = await reader.read()
        finally:
            writer.close()
            await writer.wait_closed()

        text = response.decode("utf-8", errors="replace").strip("\x00 \n")
        # "stream: OK" or "stream: <signature name> FOUND"
        if text.endswith("FOUND"):
            return text.removeprefix("stream:").removesuffix("FOUND").strip()
        return None