#!/usr/bin/python3

import aiosmtpd.controller
import email
import email.policy
from aiosmtpd.smtp import AuthResult, LoginPassword

USER = "test"
PASS = "test123"

COUNTER_FAILED_LOGIN = 0

def authenticator_func(server, session, envelope, mechanism, auth_data):

    assert isinstance(auth_data, LoginPassword)

    username = auth_data.login.decode("utf-8")
    password = auth_data.password.decode("utf-8")

    if USER == username and PASS == password:
        return AuthResult(success=True)
    else:
        return AuthResult(success=False, handled=False)

class CustomSMTPHandler:

    async def handle_DATA(self, server, session, envelope):
        global COUNTER_FAILED_LOGIN

        mail = email.message_from_bytes(envelope.content, policy=email.policy.default)
        if "Login error" in mail.get("subject"):
            COUNTER_FAILED_LOGIN += 1
        print(COUNTER_FAILED_LOGIN)
        if COUNTER_FAILED_LOGIN > 5:
            print("Brute Force")
        return '250 OK'

if __name__ == "__main__":
    handler = CustomSMTPHandler()
    server = aiosmtpd.controller.Controller(handler, hostname="0.0.0.0", port=8025,
                                            authenticator=authenticator_func,
                                            auth_required=True,
                                            auth_require_tls=False)
    server.start()
    input("Server started. Press Return to quit.")
    server.stop()