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

"""Verify the trusted-network and rate-limit policy of internal SMTP."""

from __future__ import annotations

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


ADMIN = os.environ.get("BONGO_TEST_ADMIN", "/usr/bin/bongo-admin")
SYSTEMCTL = os.environ.get("BONGO_TEST_SYSTEMCTL", "/usr/bin/systemctl")
HOST = os.environ.get("BONGO_TEST_INTERNAL_HOST", "172.16.11.190")
PORT = int(os.environ.get("BONGO_TEST_INTERNAL_PORT", "26"))
ALLOWED_SOURCE = os.environ.get("BONGO_TEST_ALLOWED_SOURCE", HOST)
DENIED_SOURCE = os.environ.get("BONGO_TEST_DENIED_SOURCE", "127.0.0.1")
HELO_NAME = os.environ.get("BONGO_TEST_HELO", "device.bongo.test")
DOMAIN = os.environ.get("BONGO_TEST_DOMAIN", "bongo.test")
RECIPIENT1 = os.environ.get("BONGO_TEST_RECIPIENT1", f"test1@{DOMAIN}")
RECIPIENT2 = os.environ.get("BONGO_TEST_RECIPIENT2", f"test2@{DOMAIN}")
TIMEOUT = float(os.environ.get("BONGO_TEST_TIMEOUT", "10"))
STARTUP_GRACE = float(os.environ.get("BONGO_TEST_STARTUP_GRACE", "2"))
ALLOW_LIVE = os.environ.get("BONGO_ALLOW_LIVE_SMTP_TEST") == "1"
RESPONSE = re.compile(rb"^([0-9]{3})([- ])(.*)\r\n$")


class InternalRelayCheckError(RuntimeError):
    """Raised when internal SMTP violates the tested policy."""


def run(arguments: list[str], *, input_text: str | None = None) -> str:
    completed = subprocess.run(
        arguments,
        input=input_text,
        capture_output=True,
        text=True,
        check=False,
    )
    if completed.returncode:
        raise InternalRelayCheckError(
            f"{' '.join(arguments)} failed: "
            f"{completed.stderr.strip() or completed.stdout.strip()}"
        )
    return completed.stdout


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


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


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


def wait_port() -> None:
    deadline = time.monotonic() + 30
    consecutive = 0
    while time.monotonic() < deadline:
        sock = None
        try:
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.settimeout(0.5)
            sock.bind((ALLOWED_SOURCE, 0))
            sock.connect((HOST, PORT))
            if not sock.recv(8192).startswith(b"220 "):
                raise OSError("internal SMTP greeting was not ready")
            sock.sendall(b"QUIT\r\n")
            if not sock.recv(8192).startswith(b"221 "):
                raise OSError("internal SMTP did not complete readiness probe")
            consecutive += 1
            if consecutive == 2:
                return
            time.sleep(0.25)
        except OSError:
            consecutive = 0
            time.sleep(0.1)
        finally:
            if sock is not None:
                sock.close()
    raise InternalRelayCheckError(
        f"internal 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()
    if run(
        ["sudo", "-n", SYSTEMCTL, "is-active", "bongo.service"]
    ).strip() != "active":
        raise InternalRelayCheckError("bongo.service is not active")


def configuration_with_limits(
    original: dict,
    *,
    messages: int = 100,
    recipients: int = 100,
    bytes_per_minute: int = 1024 * 1024,
    connections: int = 5,
) -> dict:
    configuration = copy.deepcopy(original)
    configuration["internal_relay_messages_per_minute"] = messages
    configuration["internal_relay_recipients_per_minute"] = recipients
    configuration["internal_relay_bytes_per_minute"] = bytes_per_minute
    configuration["internal_relay_connections_per_ip"] = connections
    return configuration


class SMTPConnection:
    def __init__(self, source: str = ALLOWED_SOURCE) -> None:
        deadline = time.monotonic() + TIMEOUT
        while True:
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.settimeout(TIMEOUT)
            try:
                sock.bind((source, 0))
                sock.connect((HOST, PORT))
                break
            except ConnectionRefusedError:
                sock.close()
                if time.monotonic() >= deadline:
                    raise
                time.sleep(0.1)
            except BaseException:
                sock.close()
                raise
        self.socket = sock
        self.reader = sock.makefile("rb")

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

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

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

    def data(self, message: bytes, expected: int) -> list[bytes]:
        self.command("DATA", 354)
        if not message.endswith(b"\r\n"):
            message += b"\r\n"
        self.socket.sendall(message + b".\r\n")
        return self.response(expected)


def open_allowed() -> SMTPConnection:
    client = SMTPConnection()
    client.response(220)
    return client


def begin_transaction(client: SMTPConnection) -> None:
    client.command(f"EHLO {HELO_NAME}", 250)
    client.command("MAIL FROM:<root@localhost>", 250)
    client.command(f"RCPT TO:<{RECIPIENT1}>", 250)


def check_network_policy() -> None:
    allowed = open_allowed()
    try:
        allowed.command("QUIT", 221)
    finally:
        allowed.close()

    denied = SMTPConnection(DENIED_SOURCE)
    try:
        payload = denied.reader.read()
        expected = b"554 5.7.1 Source address not permitted\r\n"
        if payload != expected:
            raise InternalRelayCheckError(
                "denied network response was not exact "
                f"(expected {expected!r}, received {payload!r})"
            )
    finally:
        denied.close()


def check_connection_limit(original: dict) -> None:
    replace_configuration(
        configuration_with_limits(original, connections=1)
    )
    restart_bongo()
    first = open_allowed()
    try:
        second = SMTPConnection()
        try:
            lines = second.response(451)
            if lines != [b"4.7.1 Too many internal relay connections"]:
                raise InternalRelayCheckError(
                    f"unexpected connection-limit response {lines!r}"
                )
            if second.reader.read(1) != b"":
                raise InternalRelayCheckError(
                    "connection-limit rejection did not close the connection"
                )
        finally:
            second.close()
        first.command("QUIT", 221)
    finally:
        first.close()


def check_recipient_limit(original: dict) -> None:
    replace_configuration(
        configuration_with_limits(original, recipients=1)
    )
    restart_bongo()
    client = open_allowed()
    try:
        client.command(f"EHLO {HELO_NAME}", 250)
        client.command("MAIL FROM:<root@localhost>", 250)
        client.command(f"RCPT TO:<{RECIPIENT1}>", 250)
        lines = client.command(f"RCPT TO:<{RECIPIENT2}>", 451)
        if lines != [b"4.7.1 Internal relay recipient rate exceeded"]:
            raise InternalRelayCheckError(
                f"unexpected recipient-limit response {lines!r}"
            )
        client.command("RSET", 250)
        client.command("QUIT", 221)
    finally:
        client.close()


def check_byte_limit(original: dict) -> None:
    replace_configuration(
        configuration_with_limits(original, bytes_per_minute=64)
    )
    restart_bongo()
    client = open_allowed()
    try:
        begin_transaction(client)
        lines = client.data(
            b"From: root@localhost\r\n"
            + f"To: {RECIPIENT1}\r\n".encode("ascii")
            + b"Subject: SMTP-04 byte limit\r\n\r\n"
            + b"x" * 128
            + b"\r\n",
            451,
        )
        if lines != [b"4.7.1 Internal relay message rate exceeded"]:
            raise InternalRelayCheckError(
                f"unexpected byte-limit response {lines!r}"
            )
        client.command("QUIT", 221)
    finally:
        client.close()


def check_message_limit(original: dict) -> None:
    replace_configuration(
        configuration_with_limits(original, messages=1)
    )
    restart_bongo()
    client = open_allowed()
    try:
        begin_transaction(client)
        lines = client.data(
            b"From: root@localhost\r\n"
            + f"To: {RECIPIENT1}\r\n".encode("ascii")
            + b"Subject: Gr\xc3\xbc\xc3\x9fe\r\n\r\nnot committed\r\n",
            550,
        )
        if lines != [
            b"5.6.7 SMTPUTF8 required for valid internationalized headers"
        ]:
            raise InternalRelayCheckError(
                f"unexpected SMTPUTF8 rejection {lines!r}"
            )

        client.command("MAIL FROM:<root@localhost>", 250)
        client.command(f"RCPT TO:<{RECIPIENT1}>", 250)
        lines = client.data(
            b"From: root@localhost\r\n"
            + f"To: {RECIPIENT1}\r\n".encode("ascii")
            + b"Subject: SMTP-04 message limit\r\n\r\n"
            + b"not committed\r\n",
            451,
        )
        if lines != [b"4.7.1 Internal relay message rate exceeded"]:
            raise InternalRelayCheckError(
                f"unexpected message-limit response {lines!r}"
            )
        client.command("QUIT", 221)
    finally:
        client.close()


def phase(name: str, function, *arguments) -> None:
    print(f"SMTP-04: {name}", file=sys.stderr, flush=True)
    try:
        function(*arguments)
    except (InternalRelayCheckError, OSError, ValueError) as error:
        raise InternalRelayCheckError(f"{name}: {error}") from error


def main() -> int:
    if not ALLOW_LIVE:
        raise InternalRelayCheckError(
            "set BONGO_ALLOW_LIVE_SMTP_TEST=1 for this temporary "
            "live-configuration test"
        )
    if ALLOWED_SOURCE == DENIED_SOURCE:
        raise InternalRelayCheckError(
            "allowed and denied source addresses must differ"
        )
    if STARTUP_GRACE < 0:
        raise InternalRelayCheckError(
            "BONGO_TEST_STARTUP_GRACE cannot be negative"
        )

    original = read_configuration()
    if not original.get("internal_relay_enabled"):
        raise InternalRelayCheckError("internal SMTP is not enabled")
    if int(original.get("internal_relay_port", 0)) != PORT:
        raise InternalRelayCheckError(
            "BONGO_TEST_INTERNAL_PORT does not match the live configuration"
        )
    configured_host = str(original.get("internal_relay_bind_address", ""))
    if configured_host != HOST:
        raise InternalRelayCheckError(
            "BONGO_TEST_INTERNAL_HOST does not match the live configuration"
        )

    restored = False
    try:
        phase("network policy", check_network_policy)
        phase("connection limit", check_connection_limit, original)
        phase("recipient limit", check_recipient_limit, original)
        phase("byte limit", check_byte_limit, original)
        phase("message limit", check_message_limit, original)
        print("SMTP-04: restore configuration", file=sys.stderr, flush=True)
        replace_configuration(original)
        restart_bongo()
        restored = read_configuration() == original
        if not restored:
            raise InternalRelayCheckError(
                "SMTP configuration did not round-trip after restoration"
            )
    finally:
        if not restored:
            replace_configuration(original)
            restart_bongo()

    print(
        "SMTP-04 PASS "
        f"host={HOST}:{PORT} allowed={ALLOWED_SOURCE} denied={DENIED_SOURCE} "
        "network=554-exact connection-limit=451 recipient-limit=451 "
        "byte-limit=451 message-limit=451 queue-write=no restored=yes"
    )
    return 0


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except (
        InternalRelayCheckError,
        json.JSONDecodeError,
        OSError,
        ValueError,
    ) as error:
        print(f"SMTP-04 FAIL: {error}")
        raise SystemExit(1)