#!/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 the non-delivery SMTP command surface on port 25."""

from __future__ import annotations

import os
import re
import socket


HOST = os.environ.get("BONGO_TEST_HOST", "127.0.0.1")
PORT = int(os.environ.get("BONGO_TEST_SMTP_PORT", "25"))
TIMEOUT = float(os.environ.get("BONGO_TEST_TIMEOUT", "10"))
HELO_NAME = os.environ.get("BONGO_TEST_HELO", "client.bongo.test")
EXISTING_USER = os.environ.get("BONGO_TEST_USER", "test1")
MISSING_USER = os.environ.get(
    "BONGO_TEST_MISSING_USER", "smtp01-definitely-missing"
)
RESPONSE = re.compile(rb"^([0-9]{3})([- ])(.*)\r\n$")


class SMTPCheckError(RuntimeError):
    """Raised when the live SMTP peer violates the tested contract."""


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, expected: int) -> list[bytes]:
        lines: list[bytes] = []
        expected_code = f"{expected:03d}".encode("ascii")
        while True:
            line = self.reader.readline(8194)
            if not line:
                raise SMTPCheckError(
                    f"connection closed while waiting for SMTP {expected}"
                )
            if len(line) > 8192:
                raise SMTPCheckError("SMTP response line exceeded 8192 bytes")
            match = RESPONSE.match(line)
            if match is None:
                raise SMTPCheckError(f"malformed SMTP response: {line!r}")
            if match.group(1) != expected_code:
                raise SMTPCheckError(
                    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]:
        payload = command.encode("ascii") + b"\r\n"
        self.socket.sendall(payload)
        return self.response(expected)


def contains_capability(lines: list[bytes], capability: bytes) -> bool:
    wanted = capability.upper()
    return any(
        line.split(None, 1)[0].upper() == wanted
        for line in lines[1:]
        if line
    )


def check_session(extended: bool) -> tuple[int, int]:
    client = SMTPConnection()
    try:
        greeting = client.response(220)
        if not greeting or not greeting[0].strip():
            raise SMTPCheckError("SMTP greeting did not identify the server")

        if extended:
            hello = client.command(f"EHLO {HELO_NAME}", 250)
            if len(hello) < 2:
                raise SMTPCheckError("EHLO did not return a multiline response")
            if not contains_capability(hello, b"HELP"):
                raise SMTPCheckError("EHLO did not advertise HELP")
            if contains_capability(hello, b"VRFY"):
                raise SMTPCheckError(
                    "EHLO advertised VRFY while the policy disables it"
                )
        else:
            hello = client.command(f"HELO {HELO_NAME}", 250)
            if len(hello) != 1:
                raise SMTPCheckError("HELO unexpectedly returned extensions")

        client.command("NOOP smtp-01", 250)
        client.command("RSET", 250)
        help_lines = client.command("HELP", 214)
        if not any(b"HELO" in line and b"QUIT" in line for line in help_lines):
            raise SMTPCheckError("HELP omitted the basic SMTP commands")

        existing = client.command(f"VRFY {EXISTING_USER}", 252)
        missing = client.command(f"VRFY {MISSING_USER}", 252)
        if existing != missing:
            raise SMTPCheckError(
                "disabled VRFY disclosed whether a local user exists"
            )

        client.command("QUIT", 221)
        client.socket.settimeout(2)
        if client.reader.read(1) != b"":
            raise SMTPCheckError("server did not close the connection after QUIT")
        return len(greeting), len(hello)
    finally:
        client.close()


def main() -> int:
    helo_greeting, helo_lines = check_session(False)
    ehlo_greeting, ehlo_lines = check_session(True)
    print(
        "SMTP-01 PASS "
        f"host={HOST} port={PORT} greeting=yes "
        f"helo-lines={helo_lines} ehlo-lines={ehlo_lines} "
        "noop=yes rset=yes help=yes vrfy=252-no-disclosure quit=yes "
        f"sessions={helo_greeting + ehlo_greeting}"
    )
    return 0


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