bongo/contrib/testing/store-protocol-safety-check.py

#!/usr/bin/env python3
# This program is free software, licensed under the terms of the GNU GPL.
# See the Bongo COPYING file for full details.
# Copyright (c) 2026 Bongo Project contributors

"""Check that the live Store protocol rejects unsafe input without damage."""

from __future__ import annotations

import os
import socket
import sys
from dataclasses import dataclass


HOST = os.environ.get("BONGO_TEST_STORE_HOST", "127.0.0.1")
PORT = int(os.environ.get("BONGO_TEST_STORE_PORT", "689"))
USER = os.environ.get("BONGO_TEST_USER", "test1")
OTHER_USER = os.environ.get("BONGO_TEST_OTHER_USER", "test2")
PASSWORD = os.environ.get("BONGO_TEST_PASSWORD", "")
TIMEOUT = float(os.environ.get("BONGO_TEST_TIMEOUT", "5"))


class ProtocolError(RuntimeError):
    """Raised when the Store protocol does not produce the expected result."""


@dataclass
class StoreConnection:
    socket: socket.socket
    stream: object

    @classmethod
    def connect(cls) -> "StoreConnection":
        connection = socket.create_connection((HOST, PORT), timeout=TIMEOUT)
        connection.settimeout(TIMEOUT)
        result = cls(connection, connection.makefile("rwb"))
        result.expect((4242,), "Store greeting")
        return result

    def close(self) -> None:
        try:
            self.stream.close()
        finally:
            self.socket.close()

    def read_line(self) -> bytes:
        line = self.stream.readline(16384)
        if not line:
            raise ProtocolError("Store closed the connection unexpectedly")
        if len(line) >= 16384:
            raise ProtocolError("Store returned an oversized response")
        return line.rstrip(b"\r\n")

    def expect(self, codes: tuple[int, ...], label: str) -> bytes:
        response = self.read_line()
        try:
            code = int(response.split(maxsplit=1)[0])
        except (ValueError, IndexError) as error:
            raise ProtocolError(f"{label}: invalid response {response!r}") from error
        if code not in codes:
            raise ProtocolError(
                f"{label}: expected {codes}, received {response.decode(errors='replace')}"
            )
        return response

    def command(
        self, command: str, codes: tuple[int, ...], label: str
    ) -> bytes:
        self.stream.write(command.encode("utf-8") + b"\r\n")
        self.stream.flush()
        return self.expect(codes, label)

    def authenticate(self) -> None:
        self.command(
            f"AUTH USER {USER} {PASSWORD}",
            (1000,),
            f"authentication for {USER}",
        )


def require_password() -> None:
    if not PASSWORD:
        raise ProtocolError("BONGO_TEST_PASSWORD must be set")
    if any(character.isspace() for character in PASSWORD):
        raise ProtocolError(
            "BONGO_TEST_PASSWORD must not contain whitespace for AUTH USER"
        )


def unauthenticated_checks() -> int:
    client = StoreConnection.connect()
    try:
        client.command("NOT-A-COMMAND", (3000,), "unknown command")
        client.command("NOOP", (1000,), "connection after unknown command")
        client.command("STORE " + USER, (3241,), "store before authentication")
        client.command("LIST /mail", (3243,), "list before Store selection")
        client.command("NOOP", (1000,), "connection after unauthorised access")
    finally:
        client.close()
    return 4


def malformed_checks() -> int:
    client = StoreConnection.connect()
    try:
        client.authenticate()
        client.command("STORE " + USER, (1000,), "select own Store")
        client.command("READ", (3010,), "missing READ argument")
        client.command("READ not-a-guid", (3011,), "invalid document GUID")
        client.command(
            "READ 1 0 1 extra",
            (3010,),
            "excess READ arguments",
        )
        client.command("NOOP", (1000,), "connection after malformed commands")
    finally:
        client.close()
    return 3


def privileged_command_checks() -> int:
    client = StoreConnection.connect()
    try:
        client.authenticate()
        client.command("STORE " + OTHER_USER, (1000,), "select foreign Store")
        client.command("LIST /mail", (4240,), "foreign Store contents")
        for command in ("REINDEX", "REPAIR", "RESET", "SHUTDOWN"):
            client.command(
                command,
                (3243,),
                f"unprivileged {command}",
            )
            client.command(
                "NOOP",
                (1000,),
                f"connection after rejected {command}",
            )
    finally:
        client.close()
    return 5


def oversized_line_check() -> int:
    client = StoreConnection.connect()
    try:
        client.socket.sendall(b"X" * 4096 + b"\r\n")
        client.socket.shutdown(socket.SHUT_WR)
        response = client.socket.recv(1)
        if response:
            raise ProtocolError(
                f"oversized command produced unexpected response {response!r}"
            )
    except (ConnectionResetError, BrokenPipeError):
        pass
    finally:
        client.close()

    recovery = StoreConnection.connect()
    try:
        recovery.command("NOOP", (1000,), "Store after oversized command")
    finally:
        recovery.close()
    return 1


def main() -> int:
    require_password()
    checks = 0
    checks += unauthenticated_checks()
    checks += malformed_checks()
    checks += privileged_command_checks()
    checks += oversized_line_check()
    print(
        "STQ-08 PASS "
        f"checks={checks} oversized-bytes=4096 "
        "manager-only=REINDEX,REPAIR,RESET,SHUTDOWN "
        f"foreign-store={OTHER_USER}"
    )
    return 0


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except (OSError, ProtocolError) as error:
        print(f"STQ-08 FAIL: {error}", file=sys.stderr)
        raise SystemExit(1)