#!/usr/bin/env python3
"""Mock SMTP 服务器 —— 验证 SMTP 客户端发送流程
支持: EHLO(多行), AUTH PLAIN/LOGIN, MAIL FROM, RCPT TO, DATA(点填充), QUIT
用法: python3 mock_smtp_server.py <port> --auth-user USER --auth-pass PASS
"""
import socket
import threading
import argparse
import base64

received = []


def handle(conn, auth_user, auth_pass, out=''):
    conn.sendall(b'220 mock-smtp ESMTP ready\r\n')
    buf = b''
    authenticated = False
    try:
        while True:
            data = conn.recv(65536)
            if not data:
                break
            buf += data
            while b'\r\n' in buf:
                line, buf = buf.split(b'\r\n', 1)
                text = line.decode('utf-8', 'replace').strip()
                if not text:
                    continue
                upper = text.upper()
                if upper.startswith('EHLO'):
                    conn.sendall(b'250-mock-smtp\r\n')
                    conn.sendall(b'250-SIZE 52428800\r\n')
                    conn.sendall(b'250-AUTH PLAIN LOGIN\r\n')
                    conn.sendall(b'250 OK\r\n')
                elif upper.startswith('AUTH PLAIN'):
                    token = base64.b64decode(text.split(' ')[2]).decode('utf-8', 'replace')
                    _, user, pw = token.split('\0')
                    if user == auth_user and pw == auth_pass:
                        authenticated = True
                        conn.sendall(b'235 2.7.0 Authentication successful\r\n')
                    else:
                        conn.sendall(b'535 5.7.8 Authentication credentials invalid\r\n')
                elif upper.startswith('AUTH LOGIN'):
                    conn.sendall(b'334 VXNlcm5hbWU6\r\n')
                    # 简化: 直接读后续两行 base64
                    data = conn.recv(65536)
                    buf += data
                    while b'\r\n' in buf:
                        line, buf = buf.split(b'\r\n', 1)
                        user = base64.b64decode(line.strip()).decode()
                        conn.sendall(b'334 UGFzc3dvcmQ6\r\n')
                        break
                    data = conn.recv(65536)
                    buf += data
                    while b'\r\n' in buf:
                        line, buf = buf.split(b'\r\n', 1)
                        pw = base64.b64decode(line.strip()).decode()
                        break
                    if user == auth_user and pw == auth_pass:
                        authenticated = True
                        conn.sendall(b'235 2.7.0 Authentication successful\r\n')
                    else:
                        conn.sendall(b'535 5.7.8 Authentication credentials invalid\r\n')
                elif upper.startswith('MAIL FROM'):
                    conn.sendall(b'250 2.1.0 Ok\r\n')
                elif upper.startswith('RCPT TO'):
                    conn.sendall(b'250 2.1.5 Ok\r\n')
                elif upper.startswith('DATA'):
                    conn.sendall(b'354 End data with <CR><LF>.<CR><LF>\r\n')
                    # 读取邮件数据直到 \r\n.\r\n
                    data = b''
                    while True:
                        chunk = conn.recv(65536)
                        if not chunk:
                            break
                        data += chunk
                        if b'\r\n.\r\n' in data:
                            break
                    received.append(data)
                    if out:
                        with open(out, 'ab') as f:
                            f.write(data)
                            f.write(b'\n=====MOCK_MSG_END=====\n')
                    conn.sendall(b'250 2.0.0 Ok: queued as mock-1\r\n')
                elif upper.startswith('QUIT'):
                    conn.sendall(b'221 2.0.0 Bye\r\n')
                    return
                elif upper.startswith('RSET'):
                    conn.sendall(b'250 2.0.0 Ok\r\n')
                else:
                    conn.sendall(b'250 Ok\r\n')
    except Exception:
        pass
    finally:
        conn.close()


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('port', type=int)
    parser.add_argument('--auth-user', default='test@example.com')
    parser.add_argument('--auth-pass', default='secret')
    parser.add_argument('--out', default='')
    args = parser.parse_args()

    srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    srv.bind(('127.0.0.1', args.port))
    srv.listen(5)
    print(f'MOCK_SMTP_READY {args.port}', flush=True)
    while True:
        conn, _ = srv.accept()
        t = threading.Thread(
            target=handle,
            args=(conn, args.auth_user, args.auth_pass),
            kwargs={'out': args.out},
            daemon=True,
        )
        t.start()


if __name__ == '__main__':
    main()