PacketStation/tests/server/fakeclam.py

"""A fake clamd for tests: a threading TCP server that speaks the part of the clamd protocol Packet Station uses.

PING answers PONG. INSTREAM reads 4-byte big-endian chunk lengths and the chunks until a zero length, then answers
`stream: OK`, or `stream: Eicar-Test-Signature FOUND` when the bytes hold the EICAR test string. A command sent as
`z<CMD>\0` gets its answer ended by NUL, `n<CMD>\n` or a bare `<CMD>\n` by a newline.

`FakeClam(refuse=True)` reserves a port and closes it, so a connection is refused. `mode="garbage"` answers
INSTREAM with a line that is not a clamd answer.
"""

from __future__ import annotations

import socket
import socketserver
import struct
import threading
from typing import Any

# Built from parts so the source file itself does not hold the test string (a desktop scanner would flag it).
EICAR = rb"X5O!P%@AP[4\PZX54(P^)7CC)7}$" + b"EICAR-STANDARD-" + b"ANTIVIRUS-TEST-FILE!$H+H*"
SIGNATURE = "Eicar-Test-Signature"


def _read_exact(sock: socket.socket, n: int) -> bytes:
    buf = b""
    while len(buf) < n:
        chunk = sock.recv(n - len(buf))
        if not chunk:
            raise ConnectionError("client closed the stream early")
        buf += chunk
    return buf


def _read_command(sock: socket.socket) -> tuple[str, bytes]:
    """The command name and the terminator its answer uses."""
    first = _read_exact(sock, 1)
    if first == b"z":
        end = b"\0"
    elif first == b"n":
        end = b"\n"
    else:
        end = b"\n"
    buf = b"" if first in (b"z", b"n") else first
    while True:
        c = _read_exact(sock, 1)
        if c in (b"\0", b"\n"):
            break
        buf += c
        if len(buf) > 64:
            break
    return buf.decode("ascii", "replace"), end


class _Handler(socketserver.BaseRequestHandler):
    server: _Server

    def handle(self) -> None:
        sock: socket.socket = self.request
        sock.settimeout(10)
        try:
            cmd, end = _read_command(sock)
        except (ConnectionError, OSError):
            return
        fake = self.server.fake
        with fake.lock:
            fake.commands.append(cmd)
        if cmd == "PING":
            sock.sendall(b"PONG" + end)
            return
        if cmd == "INSTREAM":
            data = b""
            try:
                while True:
                    (n,) = struct.unpack(">I", _read_exact(sock, 4))
                    if n == 0:
                        break
                    data += _read_exact(sock, n)
            except (ConnectionError, OSError):
                return
            with fake.lock:
                fake.scanned.append(data)
            if fake.mode == "garbage":
                sock.sendall(b"what is this" + end)
            elif EICAR in data:
                sock.sendall(f"stream: {SIGNATURE} FOUND".encode() + end)
            else:
                sock.sendall(b"stream: OK" + end)
            return
        sock.sendall(b"UNKNOWN COMMAND" + end)


class _Server(socketserver.ThreadingTCPServer):
    daemon_threads = True
    allow_reuse_address = True
    fake: FakeClam


class FakeClam:
    """Start with `with FakeClam() as clam:`; point settings at `clam.host` and `clam.port`."""

    def __init__(self, *, refuse: bool = False, mode: str = "normal") -> None:
        self.host = "127.0.0.1"
        self.refuse = refuse
        self.mode = mode
        self.lock = threading.Lock()
        self.commands: list[str] = []
        self.scanned: list[bytes] = []
        self._server: _Server | None = None
        self._thread: threading.Thread | None = None
        self.port = 0

    @property
    def calls(self) -> int:
        with self.lock:
            return len(self.commands)

    def start(self) -> FakeClam:
        if self.refuse:
            s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            s.bind((self.host, 0))
            self.port = int(s.getsockname()[1])
            s.close()  # nothing listens there now: a connect is refused
            return self
        self._server = _Server((self.host, 0), _Handler)
        self._server.fake = self
        self.port = int(self._server.server_address[1])
        self._thread = threading.Thread(target=self._server.serve_forever, name="fakeclam", daemon=True)
        self._thread.start()
        return self

    def stop(self) -> None:
        if self._server is not None:
            self._server.shutdown()
            self._server.server_close()
            self._server = None
        if self._thread is not None:
            self._thread.join(5)
            self._thread = None

    def __enter__(self) -> FakeClam:
        return self.start()

    def __exit__(self, *exc: Any) -> None:
        self.stop()


def ping(host: str, port: int) -> bytes:
    """A raw zPING, for checking the fake itself."""
    with socket.create_connection((host, port), timeout=5) as s:
        s.sendall(b"zPING\0")
        return s.recv(64)


def instream(host: str, port: int, data: bytes, chunk: int = 7) -> bytes:
    """A raw zINSTREAM in small chunks, for checking the fake itself."""
    with socket.create_connection((host, port), timeout=5) as s:
        s.sendall(b"zINSTREAM\0")
        for i in range(0, len(data), chunk):
            part = data[i : i + chunk]
            s.sendall(struct.pack(">I", len(part)) + part)
        s.sendall(struct.pack(">I", 0))
        out = b""
        while not out.endswith(b"\0"):
            got = s.recv(256)
            if not got:
                break
            out += got
        return out