bongo/contrib/testing/smtp-size-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

"""Exercise SMTP SIZE boundaries and low Queue-space responses."""

from __future__ import annotations

import argparse
import copy
import json
import os
import re
import socket
import subprocess
import sys
import time

from bongo.store.StoreClient import StoreClient


ADMIN = os.environ.get("BONGO_TEST_ADMIN", "/usr/bin/bongo-admin")
QUEUE_TOOL = os.environ.get(
    "BONGO_TEST_QUEUE_TOOL", "/usr/bin/bongo-queuetool"
)
SYSTEMCTL = os.environ.get("BONGO_TEST_SYSTEMCTL", "/usr/bin/systemctl")
HOST = os.environ.get("BONGO_TEST_HOST", "127.0.0.1")
PORT = int(os.environ.get("BONGO_TEST_SMTP_PORT", "25"))
STORE_PORT = int(os.environ.get("BONGO_TEST_STORE_PORT", "689"))
USER = os.environ.get("BONGO_TEST_USER", "test1")
DOMAIN = os.environ.get("BONGO_TEST_DOMAIN", "bongo.test")
PASSWORD = os.environ.get("BONGO_TEST_PASSWORD", "")
LIMIT = int(os.environ.get("BONGO_TEST_MESSAGE_LIMIT", "1024"))
TIMEOUT = float(os.environ.get("BONGO_TEST_TIMEOUT", "30"))
STARTUP_GRACE = float(os.environ.get("BONGO_TEST_STARTUP_GRACE", "2"))
ALLOW_LIVE = os.environ.get("BONGO_ALLOW_LIVE_SMTP_TEST") == "1"
TOKEN = f"smtp11-{os.getpid()}-{int(time.time())}"
TOKEN_PREFIX = b"smtp11-"
RESPONSE = re.compile(rb"^([0-9]{3})([- ])(.*)\r\n$")
QUEUE_ID = re.compile(r"^[0-9A-Fa-f]{3}-[0-9A-Fa-f]+$")


class SMTP11Error(RuntimeError):
    """Raised when SMTP size or free-space handling is incorrect."""


def run(
    arguments: list[str],
    *,
    input_text: str | None = None,
    binary: bool = False,
    check: bool = True,
) -> subprocess.CompletedProcess:
    completed = subprocess.run(
        arguments,
        input=input_text,
        capture_output=True,
        text=not binary,
        check=False,
    )
    if check and completed.returncode:
        stdout = completed.stdout
        stderr = completed.stderr
        if binary:
            stdout = stdout.decode("utf-8", "replace")
            stderr = stderr.decode("utf-8", "replace")
        raise SMTP11Error(
            f"{' '.join(arguments)} failed: "
            f"{stderr.strip() or stdout.strip()}"
        )
    return completed


def admin(*arguments: str, input_text: str | None = None) -> str:
    return run(
        ["sudo", "-n", ADMIN, *arguments], input_text=input_text
    ).stdout


def queue(
    *arguments: str, binary: bool = False, check: bool = True
) -> subprocess.CompletedProcess:
    return run(
        ["sudo", "-n", "-u", "bongo", QUEUE_TOOL, *arguments],
        binary=binary,
        check=check,
    )


def read_configuration() -> dict:
    try:
        value = json.loads(admin("__config-read", "smtp"))
    except json.JSONDecodeError as error:
        raise SMTP11Error(
            "bongo-admin returned invalid SMTP configuration"
        ) from error
    if not isinstance(value, dict):
        raise SMTP11Error("SMTP configuration is not an object")
    return value


def replace_configuration(configuration: dict) -> None:
    admin(
        "__config-replace",
        "smtp",
        input_text=json.dumps(configuration, separators=(",", ":")),
    )


class SMTPConnection:
    def __init__(self) -> None:
        self.socket = socket.create_connection((HOST, PORT), timeout=TIMEOUT)
        self.socket.settimeout(TIMEOUT)
        self.reader = self.socket.makefile("rb")

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

    def response(self) -> tuple[int, list[bytes]]:
        lines: list[bytes] = []
        response_code: int | None = None
        while True:
            line = self.reader.readline(8194)
            if not line:
                raise SMTP11Error("connection closed while awaiting response")
            if len(line) > 8192:
                raise SMTP11Error("SMTP response line exceeded 8192 bytes")
            match = RESPONSE.match(line)
            if match is None:
                raise SMTP11Error(f"malformed SMTP response: {line!r}")
            code = int(match.group(1))
            if response_code is None:
                response_code = code
            elif code != response_code:
                raise SMTP11Error(
                    f"multiline response changed status: {line!r}"
                )
            lines.append(match.group(3))
            if match.group(2) == b" ":
                return code, lines

    def expect(self, expected: int, context: str) -> list[bytes]:
        code, lines = self.response()
        if code != expected:
            raise SMTP11Error(
                f"{context}: expected SMTP {expected}, got {code}: {lines!r}"
            )
        return lines

    def command(self, command: str) -> tuple[int, list[bytes]]:
        self.socket.sendall(command.encode("ascii") + b"\r\n")
        return self.response()


def require(
    result: tuple[int, list[bytes]], expected: int, context: str
) -> list[bytes]:
    code, lines = result
    if code != expected:
        raise SMTP11Error(
            f"{context}: expected SMTP {expected}, got {code}: {lines!r}"
        )
    return lines


def wait_port() -> None:
    deadline = time.monotonic() + 30
    consecutive = 0
    while time.monotonic() < deadline:
        client: SMTPConnection | None = None
        try:
            client = SMTPConnection()
            client.expect(220, "readiness greeting")
            require(client.command("QUIT"), 221, "readiness QUIT")
            consecutive += 1
            if consecutive == 2:
                return
            time.sleep(0.25)
        except (OSError, SMTP11Error):
            consecutive = 0
            time.sleep(0.1)
        finally:
            if client is not None:
                client.close()
    raise SMTP11Error(f"SMTP did not become ready at {HOST}:{PORT}")


def restart_bongo() -> None:
    run(["sudo", "-n", SYSTEMCTL, "restart", "bongo.service"])
    time.sleep(STARTUP_GRACE)
    wait_port()
    active = run(
        ["sudo", "-n", SYSTEMCTL, "is-active", "bongo.service"]
    ).stdout.strip()
    if active != "active":
        raise SMTP11Error("bongo.service is not active after restart")


def make_message(length: int, token: str) -> bytes:
    prefix = (
        f"From: SMTP-11 <sender@example.test>\r\n"
        f"To: {USER}@{DOMAIN}\r\n"
        f"Subject: SMTP-11 {token}\r\n"
        f"Message-ID: <{token}@{DOMAIN}>\r\n"
        f"X-Bongo-Test: {token}\r\n"
        "\r\n"
        "size-boundary="
    ).encode("ascii")
    if length < len(prefix) + 2:
        raise SMTP11Error(
            f"message limit {length} is too small for the test fixture"
        )
    message = prefix + b"x" * (length - len(prefix) - 2) + b"\r\n"
    if len(message) != length:
        raise SMTP11Error("internal message-size construction error")
    return message


def ehlo(client: SMTPConnection) -> list[bytes]:
    client.expect(220, "greeting")
    return require(
        client.command("EHLO smtp11-client.example"), 250, "EHLO"
    )


def start_local_transaction(
    client: SMTPConnection, *, size: int | None
) -> None:
    suffix = "" if size is None else f" SIZE={size}"
    require(
        client.command(f"MAIL FROM:<sender@example.test>{suffix}"),
        250,
        "MAIL FROM",
    )
    require(
        client.command(f"RCPT TO:<{USER}@{DOMAIN}>"),
        250,
        "RCPT TO",
    )


def send_data(client: SMTPConnection, message: bytes) -> tuple[int, list[bytes]]:
    require(client.command("DATA"), 354, "DATA")
    client.socket.sendall(message + b".\r\n")
    return client.response()


def queue_token_ids() -> list[str]:
    identifiers: list[str] = []
    for line in queue("list").stdout.splitlines():
        fields = line.split()
        if not fields or QUEUE_ID.fullmatch(fields[0]) is None:
            continue
        message = queue("message", fields[0], binary=True, check=False)
        if message.returncode == 0 and TOKEN_PREFIX in message.stdout:
            identifiers.append(fields[0])
    return identifiers


def open_store() -> StoreClient:
    return StoreClient(
        USER,
        USER,
        authPassword=PASSWORD,
        host=HOST,
        port=STORE_PORT,
    )


def store_token_documents(store: StoreClient) -> list[str]:
    matches: list[str] = []
    # LIST is a streamed Store response. Consume it completely before issuing
    # READ commands on the same connection or payload bytes can be mistaken
    # for subsequent response lines.
    identifiers = [entry.uid for entry in store.List("/mail/INBOX")]
    for identifier in identifiers:
        try:
            data = store.Read(identifier)
        except Exception:
            continue
        if TOKEN_PREFIX in data:
            matches.append(identifier)
    return matches


def cleanup(store: StoreClient) -> None:
    deadline = time.monotonic() + 20
    quiet_since: float | None = None
    while time.monotonic() < deadline:
        found = False
        for document in store_token_documents(store):
            found = True
            store.Delete(document)
        for queue_id in queue_token_ids():
            found = True
            queue("delete", queue_id, check=False)
        if found:
            quiet_since = None
        elif quiet_since is None:
            quiet_since = time.monotonic()
        elif time.monotonic() - quiet_since >= 1:
            return
        time.sleep(0.25)


def wait_for_exact_delivery(store: StoreClient) -> str:
    deadline = time.monotonic() + TIMEOUT
    while time.monotonic() < deadline:
        documents = store_token_documents(store)
        if len(documents) > 1:
            raise SMTP11Error(
                f"exact-limit message delivered more than once: {documents!r}"
            )
        if documents:
            return documents[0]
        time.sleep(0.25)
    raise SMTP11Error("exact-limit message did not reach the Store")


def normal_checks() -> int:
    if not ALLOW_LIVE:
        raise SMTP11Error(
            "set BONGO_ALLOW_LIVE_SMTP_TEST=1 for the live configuration test"
        )
    if not PASSWORD:
        raise SMTP11Error("BONGO_TEST_PASSWORD is required")
    if LIMIT < 256:
        raise SMTP11Error("BONGO_TEST_MESSAGE_LIMIT must be at least 256")

    original = read_configuration()
    configured = copy.deepcopy(original)
    configured["message_size_limit"] = LIMIT
    restored = False
    store: StoreClient | None = None
    try:
        replace_configuration(configured)
        restart_bongo()
        store = open_store()
        cleanup(store)

        client = SMTPConnection()
        try:
            capabilities = ehlo(client)
            if f"SIZE {LIMIT}".encode("ascii") not in capabilities:
                raise SMTP11Error(
                    f"EHLO did not advertise SIZE {LIMIT}: {capabilities!r}"
                )

            lines = require(
                client.command(
                    f"MAIL FROM:<sender@example.test> SIZE={LIMIT + 1}"
                ),
                552,
                "declared over-limit MAIL FROM",
            )
            if b"size" not in b" ".join(lines).lower():
                raise SMTP11Error(
                    f"SMTP 552 did not explain the size failure: {lines!r}"
                )
            require(client.command("NOOP"), 250, "NOOP after MAIL rejection")

            start_local_transaction(client, size=LIMIT)
            require(
                send_data(client, make_message(LIMIT, TOKEN)),
                250,
                "exact-limit DATA",
            )
            require(client.command("NOOP"), 250, "NOOP after exact DATA")

            start_local_transaction(client, size=None)
            require(
                send_data(
                    client,
                    make_message(LIMIT + 1, f"{TOKEN}-over"),
                ),
                552,
                "undeclared over-limit DATA",
            )
            require(client.command("NOOP"), 250, "NOOP after DATA rejection")
            require(client.command("QUIT"), 221, "QUIT")
        finally:
            client.close()

        document = wait_for_exact_delivery(store)
        store.Delete(document)
        cleanup(store)
        if store_token_documents(store) or queue_token_ids():
            raise SMTP11Error("test content remains after cleanup")

        replace_configuration(original)
        restart_bongo()
        restored = True
    finally:
        if store is not None:
            try:
                if not restored:
                    cleanup(store)
            except Exception as error:
                print(
                    f"SMTP-11 cleanup warning: {error}", file=sys.stderr
                )
            finally:
                try:
                    store.Quit()
                except Exception:
                    pass
        if not restored:
            try:
                replace_configuration(original)
                restart_bongo()
            except Exception as error:
                print(
                    f"SMTP-11 restore warning: {error}", file=sys.stderr
                )

    print(
        "SMTP-11 SIZE PASS "
        f"limit={LIMIT} advertised=yes declared-over=552 "
        "exact-data=250 absent-size-over=552 session-recovery=yes "
        "queue-clean=yes store-clean=yes config-restored=yes"
    )
    return 0


def disk_checks() -> int:
    client = SMTPConnection()
    try:
        ehlo(client)
        lines = require(
            client.command("MAIL FROM:<sender@example.test> SIZE=1"),
            452,
            "declared-size full-spool MAIL FROM",
        )
        if b"storage" not in b" ".join(lines).lower():
            raise SMTP11Error(
                f"SMTP 452 did not explain the storage failure: {lines!r}"
            )
        require(
            client.command("NOOP"), 250, "NOOP after QDSPC rejection"
        )
        require(client.command("QUIT"), 221, "QUIT")
    finally:
        client.close()

    client = SMTPConnection()
    try:
        ehlo(client)
        lines = require(
            client.command("MAIL FROM:<sender@example.test>"),
            452,
            "undeclared-size full-spool MAIL FROM",
        )
        if b"storage" not in b" ".join(lines).lower():
            raise SMTP11Error(
                f"SMTP 452 did not explain the storage failure: {lines!r}"
            )
    finally:
        client.close()

    print(
        "SMTP-11 DISK PASS declared-size=452 no-size=452 "
        "enhanced-status=4.3.1"
    )
    return 0


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--disk-only",
        action="store_true",
        help="require a pre-existing full-spool fixture and test SMTP 452",
    )
    arguments = parser.parse_args()
    return disk_checks() if arguments.disk_only else normal_checks()


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except (OSError, SMTP11Error, ValueError) as error:
        print(f"SMTP-11 FAIL: {type(error).__name__}: {error}")
        raise SystemExit(1)