PacketStation/tests/server/fakesmtp.py

"""An in-process SMTP server with STARTTLS, and a throwaway CA made at test time. Never sends mail anywhere.

`make_ca(dir)` writes a CA and a server certificate for 127.0.0.1 and localhost (with `cryptography` when it is
installed, else the `openssl` command). `FakeSmtp(certs)` listens on 127.0.0.1 and keeps every message it takes.
"""

from __future__ import annotations

import contextlib
import datetime as dt
import shutil
import socket
import ssl
import subprocess
import threading
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

CRLF = b"\r\n"


@dataclass
class Certs:
    ca: Path
    cert: Path
    key: Path


def can_make_ca() -> bool:
    try:
        import cryptography  # noqa: F401
    except ImportError:
        return shutil.which("openssl") is not None
    return True


def make_ca(where: Path, name: str = "ca") -> Certs:
    where.mkdir(parents=True, exist_ok=True)
    try:
        return _with_cryptography(where, name)
    except ImportError:
        return _with_openssl(where, name)


def _with_cryptography(where: Path, name: str) -> Certs:
    import ipaddress

    from cryptography import x509
    from cryptography.hazmat.primitives import hashes, serialization
    from cryptography.hazmat.primitives.asymmetric import rsa
    from cryptography.x509.oid import ExtendedKeyUsageOID, NameOID

    now = dt.datetime.now(dt.UTC)
    ca_key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
    ca_name = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, f"Packet Station test {name}")])
    ca_cert = (
        x509.CertificateBuilder()
        .subject_name(ca_name)
        .issuer_name(ca_name)
        .public_key(ca_key.public_key())
        .serial_number(x509.random_serial_number())
        .not_valid_before(now - dt.timedelta(minutes=5))
        .not_valid_after(now + dt.timedelta(days=1))
        .add_extension(x509.BasicConstraints(ca=True, path_length=None), critical=True)
        .add_extension(x509.KeyUsage(True, False, False, False, False, True, True, False, False), critical=True)
        .add_extension(x509.SubjectKeyIdentifier.from_public_key(ca_key.public_key()), critical=False)
        .sign(ca_key, hashes.SHA256())
    )
    key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
    cert = (
        x509.CertificateBuilder()
        .subject_name(x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "localhost")]))
        .issuer_name(ca_name)
        .public_key(key.public_key())
        .serial_number(x509.random_serial_number())
        .not_valid_before(now - dt.timedelta(minutes=5))
        .not_valid_after(now + dt.timedelta(days=1))
        .add_extension(
            x509.SubjectAlternativeName([x509.DNSName("localhost"), x509.IPAddress(ipaddress.ip_address("127.0.0.1"))]),
            critical=False,
        )
        .add_extension(x509.BasicConstraints(ca=False, path_length=None), critical=True)
        .add_extension(x509.ExtendedKeyUsage([ExtendedKeyUsageOID.SERVER_AUTH]), critical=False)
        .add_extension(x509.AuthorityKeyIdentifier.from_issuer_public_key(ca_key.public_key()), critical=False)
        .sign(ca_key, hashes.SHA256())
    )
    pem = serialization.Encoding.PEM
    certs = Certs(where / f"{name}.pem", where / f"{name}-server.pem", where / f"{name}-server.key")
    certs.ca.write_bytes(ca_cert.public_bytes(pem))
    certs.cert.write_bytes(cert.public_bytes(pem))
    certs.key.write_bytes(
        key.private_bytes(pem, serialization.PrivateFormat.TraditionalOpenSSL, serialization.NoEncryption())
    )
    return certs


def _with_openssl(where: Path, name: str) -> Certs:
    exe = shutil.which("openssl")
    if exe is None:
        raise RuntimeError("neither cryptography nor openssl is available")
    certs = Certs(where / f"{name}.pem", where / f"{name}-server.pem", where / f"{name}-server.key")
    ca_key = where / f"{name}.key"
    csr = where / f"{name}-server.csr"
    ext = where / f"{name}-ext.cnf"
    ext.write_text(
        "basicConstraints=critical,CA:FALSE\nextendedKeyUsage=serverAuth\n"
        "subjectAltName=DNS:localhost,IP:127.0.0.1\nauthorityKeyIdentifier=keyid\n"
    )
    ca_cfg = where / f"{name}-ca.cnf"
    ca_cfg.write_text(
        "[req]\ndistinguished_name=dn\nx509_extensions=v3_ca\nprompt=no\n[dn]\nCN=Packet Station test "
        f"{name}\n[v3_ca]\nbasicConstraints=critical,CA:TRUE\nkeyUsage=critical,keyCertSign,cRLSign\n"
        "subjectKeyIdentifier=hash\n"
    )

    def run(*args: str) -> None:
        subprocess.run([exe, *args], check=True, capture_output=True, timeout=60)

    run("req", "-x509", "-newkey", "rsa:2048", "-nodes", "-sha256", "-days", "1", "-keyout", str(ca_key),
        "-out", str(certs.ca), "-config", str(ca_cfg))  # fmt: skip
    run("req", "-newkey", "rsa:2048", "-nodes", "-sha256", "-keyout", str(certs.key), "-out", str(csr),
        "-subj", "/CN=localhost")  # fmt: skip
    run("x509", "-req", "-in", str(csr), "-CA", str(certs.ca), "-CAkey", str(ca_key), "-CAcreateserial",
        "-days", "1", "-sha256", "-extfile", str(ext), "-out", str(certs.cert))  # fmt: skip
    return certs


@dataclass
class Message:
    mail_from: str
    rcpts: list[str]
    data: bytes
    tls: bool

    @property
    def text(self) -> str:
        return self.data.decode("utf-8", "replace")

    def header(self, name: str) -> str | None:
        import email

        msg = email.message_from_bytes(self.data)
        value: Any = msg.get(name)
        return None if value is None else str(value)


@dataclass
class FakeSmtp:
    """A threaded SMTP server. offer_tls=False leaves STARTTLS out of EHLO. refuse: addresses RCPT refuses."""

    certs: Certs | None = None
    offer_tls: bool = True
    refuse: set[str] = field(default_factory=set)
    messages: list[Message] = field(default_factory=list)
    connections: int = 0
    tls_failures: int = 0

    def __post_init__(self) -> None:
        self._sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self._sock.bind(("127.0.0.1", 0))
        self._sock.listen(8)
        self._sock.settimeout(0.2)
        self.host, self.port = self._sock.getsockname()
        self._stop = threading.Event()
        self._thread = threading.Thread(target=self._serve, daemon=True)
        self._lock = threading.Lock()

    def __enter__(self) -> FakeSmtp:
        self._thread.start()
        return self

    def __exit__(self, *exc: object) -> None:
        self._stop.set()
        self._thread.join(5)
        self._sock.close()

    def _serve(self) -> None:
        while not self._stop.is_set():
            try:
                conn, _ = self._sock.accept()
            except (TimeoutError, OSError):
                continue
            threading.Thread(target=self._session, args=(conn,), daemon=True).start()

    def _session(self, conn: socket.socket) -> None:
        with self._lock:
            self.connections += 1
        conn.settimeout(10)
        try:
            self._talk(conn)
        finally:
            with contextlib.suppress(OSError):
                conn.close()

    def _talk(self, conn: socket.socket) -> None:
        sock: Any = conn
        buf = b""
        tls = False
        mail_from, rcpts = "", []  # type: str, list[str]

        def send(line: str) -> None:
            sock.sendall(line.encode() + CRLF)

        def readline() -> bytes | None:
            nonlocal buf
            while CRLF not in buf:
                chunk = sock.recv(65536)
                if not chunk:
                    return None
                buf += chunk
            line, buf = buf.split(CRLF, 1)
            return line

        with contextlib.suppress(OSError, ssl.SSLError, ValueError):
            send("220 fake.test ESMTP")
            while True:
                raw = readline()
                if raw is None:
                    return
                cmd = raw.decode("utf-8", "replace")
                verb = cmd.split(" ", 1)[0].upper()
                if verb == "EHLO":
                    caps = ["250-fake.test", "250-8BITMIME"]
                    if self.offer_tls and not tls and self.certs is not None:
                        caps.append("250-STARTTLS")
                    caps.append("250 SMTPUTF8")
                    sock.sendall("\r\n".join(caps).encode() + b"\r\n")
                elif verb == "HELO":
                    send("250 fake.test")
                elif verb == "STARTTLS":
                    if not self.offer_tls or self.certs is None or tls:
                        send("454 TLS not available")
                        continue
                    send("220 Ready to start TLS")
                    ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
                    ctx.load_cert_chain(str(self.certs.cert), str(self.certs.key))
                    try:
                        sock = ctx.wrap_socket(conn, server_side=True)
                    except (ssl.SSLError, OSError):
                        with self._lock:
                            self.tls_failures += 1
                        return
                    tls, buf, mail_from, rcpts = True, b"", "", []
                elif verb == "MAIL":
                    mail_from = cmd.split(":", 1)[1].split()[0].strip("<>")
                    rcpts = []
                    send("250 OK")
                elif verb == "RCPT":
                    addr = cmd.split(":", 1)[1].split()[0].strip("<>")
                    if addr.lower() in self.refuse:
                        send("550 No such user")
                    else:
                        rcpts.append(addr)
                        send("250 OK")
                elif verb == "DATA":
                    if not rcpts:
                        send("554 No valid recipients")
                        continue
                    send("354 End data with <CR><LF>.<CR><LF>")
                    while b"\r\n.\r\n" not in b"\r\n" + buf:
                        chunk = sock.recv(65536)
                        if not chunk:
                            return
                        buf += chunk
                    full = b"\r\n" + buf
                    data, rest = full.split(b"\r\n.\r\n", 1)
                    buf = rest
                    with self._lock:
                        self.messages.append(Message(mail_from, list(rcpts), data[2:], tls))
                    send("250 OK queued")
                elif verb in ("RSET", "NOOP"):
                    send("250 OK")
                elif verb == "QUIT":
                    send("221 Bye")
                    return
                else:
                    send("502 Command not implemented")