#!/usr/bin/env python3
import socketserver
import sys


class Handler(socketserver.StreamRequestHandler):
    def write(self, line: bytes) -> None:
        try:
            self.wfile.write(line + b"\r\n")
            self.wfile.flush()
        except BrokenPipeError:
            return

    def handle(self) -> None:
        self.write(b"220 localhost ESMTP test")
        auth_step = 0
        data_mode = False
        message = bytearray()
        while True:
            line = self.rfile.readline()
            if not line:
                break
            stripped = line.rstrip(b"\r\n")
            if data_mode:
                if stripped == b".":
                    with open(sys.argv[2], "wb") as output:
                        output.write(message)
                    self.write(b"250 2.0.0 queued")
                    data_mode = False
                else:
                    message.extend(line)
                continue
            upper = stripped.upper()
            if auth_step == 1:
                auth_step = 2
                self.write(b"334 UGFzc3dvcmQ6")
            elif auth_step == 2:
                auth_step = 0
                self.write(b"235 2.7.0 authenticated")
            elif upper.startswith(b"EHLO "):
                self.write(b"250-localhost")
                self.write(b"250-AUTH LOGIN")
                self.write(b"250 OK")
            elif upper == b"AUTH LOGIN":
                auth_step = 1
                self.write(b"334 VXNlcm5hbWU6")
            elif upper.startswith(b"MAIL FROM:") or upper.startswith(b"RCPT TO:"):
                self.write(b"250 2.1.0 accepted")
            elif upper == b"DATA":
                data_mode = True
                self.write(b"354 End data with <CR><LF>.<CR><LF>")
            elif upper == b"QUIT":
                self.write(b"221 2.0.0 bye")
                break
            else:
                self.write(b"500 unexpected command")


if len(sys.argv) != 3:
    raise SystemExit("usage: fake_smtp_server.py PORT OUTPUT")

with socketserver.TCPServer(("127.0.0.1", int(sys.argv[1])), Handler) as server:
    server.handle_request()