#!/usr/bin/env python3
# Standalone HTTP wrapper around ClamAV's clamd, reachable from browser JS (unlike clamd's raw
# TCP protocol). Used by the ThreadNet-Web client fork to scan files client-side, both before
# upload (plaintext, pre-encryption) and after download+decrypt of E2EE attachments - the two
# places Synapse's own check_media_file_for_spam module (Issue #19) can never see, since
# Synapse never holds the room's decryption key.
#
# Talks to clamd via its native 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 bytes
#   3. send a zero-length chunk to signal end of stream
#   4. read the reply: "stream: OK" (clean) or "stream: <name> FOUND" (infected)
#
# Stdlib only, synchronous/threaded (ThreadingHTTPServer) - no asyncio/Twisted constraints
# here since this runs as its own plain process, unlike the Synapse module.
#
# Auth: requires "Authorization: Bearer <matrix access token>", validated against Synapse's
# own /_matrix/client/v3/account/whoami - reuses Synapse's existing auth rather than building
# a new one, and stops this becoming an open "test your malware against our AV" oracle for
# anyone on the internet. This is a hard failure (401) - unlike scan errors below, this is an
# abuse-prevention concern, not a reliability one.
#
# Fails open on clamd connection errors (treats the file as clean, logs loudly) - matches the
# same fail-open design as the Synapse module, so a ClamAV hiccup doesn't block all uploads/
# downloads site-wide.

import json
import logging
import os
import socket
import sys
import urllib.error
import urllib.request
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

logging.basicConfig(level=logging.INFO, stream=sys.stdout)
logger = logging.getLogger("clamav-http-scanner")

CLAMD_HOST = os.environ["CLAMD_HOST"]
CLAMD_PORT = int(os.environ["CLAMD_PORT"])
SYNAPSE_WHOAMI_URL = os.environ["SYNAPSE_WHOAMI_URL"]
CLAMD_TIMEOUT_SECONDS = 30
MAX_BODY_BYTES = 100 * 1024 * 1024  # 100MB, matches typical Synapse upload size limits


def check_auth(authorization_header: "str | None") -> bool:
    if not authorization_header or not authorization_header.startswith("Bearer "):
        return False
    token = authorization_header.removeprefix("Bearer ").strip()
    request = urllib.request.Request(
        SYNAPSE_WHOAMI_URL, headers={"Authorization": f"Bearer {token}"}
    )
    try:
        with urllib.request.urlopen(request, timeout=10) as response:
            return response.status == 200
    except urllib.error.URLError:
        return False


def scan_bytes(data: bytes) -> "str | None":
    """Returns the detected signature name, or None if clean. Raises on connection errors."""
    with socket.create_connection(
        (CLAMD_HOST, CLAMD_PORT), timeout=CLAMD_TIMEOUT_SECONDS
    ) as sock:
        sock.sendall(b"zINSTREAM\0")
        chunk_size = 2**14
        for offset in range(0, len(data), chunk_size):
            chunk = data[offset : offset + chunk_size]
            sock.sendall(len(chunk).to_bytes(4, "big") + chunk)
        sock.sendall((0).to_bytes(4, "big"))

        response = b""
        while True:
            part = sock.recv(4096)
            if not part:
                break
            response += part

    text = response.decode("utf-8", errors="replace").strip("\x00 \n")
    if text.endswith("FOUND"):
        return text.removeprefix("stream:").removesuffix("FOUND").strip()
    return None


class Handler(BaseHTTPRequestHandler):
    def log_message(self, format: str, *args: object) -> None:
        logger.info("%s - %s", self.address_string(), format % args)

    def _send_json(self, status: int, payload: dict) -> None:
        body = json.dumps(payload).encode("utf-8")
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.send_header("Access-Control-Allow-Origin", "*")
        self.end_headers()
        self.wfile.write(body)

    def do_OPTIONS(self) -> None:
        self.send_response(204)
        self.send_header("Access-Control-Allow-Origin", "*")
        self.send_header("Access-Control-Allow-Headers", "Authorization, Content-Type")
        self.send_header("Access-Control-Allow-Methods", "POST, OPTIONS")
        self.end_headers()

    def do_POST(self) -> None:
        # matches the ingress PathPrefix(`/_scan`) route as-is (Traefik doesn't strip the
        # prefix by default) - keep client-facing and internal path identical.
        if self.path != "/_scan":
            self._send_json(404, {"error": "not found"})
            return

        if not check_auth(self.headers.get("Authorization")):
            self._send_json(401, {"error": "invalid or missing access token"})
            return

        length = int(self.headers.get("Content-Length", 0))
        if length <= 0 or length > MAX_BODY_BYTES:
            self._send_json(400, {"error": "missing or oversized body"})
            return
        data = self.rfile.read(length)

        try:
            signature = scan_bytes(data)
        except OSError:
            logger.exception(
                "ClamAV scan failed (clamd at %s:%s unreachable?) - "
                "treating file as clean (fail-open)",
                CLAMD_HOST,
                CLAMD_PORT,
            )
            self._send_json(200, {"clean": True, "scan_error": "scanner_unavailable"})
            return

        if signature is None:
            self._send_json(200, {"clean": True})
        else:
            logger.warning("ClamAV flagged an upload/download: %s", signature)
            self._send_json(200, {"clean": False, "signature": signature})


def main() -> None:
    server = ThreadingHTTPServer(("0.0.0.0", 8090), Handler)
    logger.info("Listening on :8090, clamd=%s:%s", CLAMD_HOST, CLAMD_PORT)
    server.serve_forever()


if __name__ == "__main__":
    main()
