import asyncio
import email
import logging
from smtplib import SMTP as SMTPClient
from typing import Dict

from aiosmtpd.proxy_protocol import ProxyData
from python_socks.sync import Proxy
import os
import ssl
from aiosmtpd.controller import Controller
from aiosmtpd.smtp import Envelope, Session, SMTP
from email.message import Message
import dkim
import spf
from functools import lru_cache
import dns.resolver

PROXY = Proxy.from_url('socks5://5.78.86.156:2025')

END_HOST = "localhost"
END_PORT = 25

PROXY_SMTP_HOST = "*"
PROXY_SMTP_PORT = 2025
END_SMTP_HOST = "*"
END_SMTP_PORT = 3025
GENERIC_SMTP_HOST = "*"
GENERIC_SMTP_PORT = 25

log = logging.getLogger("smtphandler")
VALID_PROXY_ADDRS = {"5.78.86.156", "2a01:4ff:1f0:83de::1", "m-labs-intl.com", "mail.m-labs-intl.com"}
VALID_END_ADDRS = {"localhost", "127.0.0.1"}
VALID_RECEPIENTS = {"m-labs.hk", "m-labs.ph", "m-labs-intl.com", "193thz.com", "malloctech.fr"}


@lru_cache(maxsize=256)
def get_mx(domain):
    records = dns.resolver.resolve(domain, "MX")
    if not records:
        return None
    result = max(records, key=lambda r: r.preference)
    return str(result.exchange)


class ProxiedSMTP(SMTPClient):
    def _get_socket(self, host, port, timeout):
        return PROXY.connect(dest_host=host, dest_port=port, timeout=timeout)


class GenericMailServer:
    """
    Accepts mail from everywhere, and passes it to the real SMTP server, after proper SPF and DKIM checks.
    """

    async def handle_MAIL(self,
                          server: SMTP,
                          session: Session,
                          envelope: Envelope,
                          address: str,
                          mail_options: list):
        ip = session.peer[0]
        result, description = spf.check2(ip, address, session.host_name)
        valid_spf = result == 'pass'
        envelope.spf = valid_spf

        log.info("SPF: %s, %s", result, description)

        if not valid_spf:
            return '550 SPF validation failed'

        envelope.mail_from = address
        envelope.mail_options.extend(mail_options)

        return '250 OK'

    async def handle_RCPT(self,
                          server: SMTP,
                          session: Session,
                          envelope: Envelope,
                          address: str,
                          rcpt_options: list):
        if address.split("@")[1] not in VALID_RECEPIENTS:
            return '550 not relaying to that domain'

        log.debug("Handle RCPT for %s", address)
        envelope.rcpt_tos.append(address)
        return '250 OK'

    async def handle_DATA(self,
                          server: SMTP,
                          session: Session,
                          envelope: Envelope):
        valid_dkim = dkim.verify(envelope.content)
        envelope.dkim = valid_dkim
        log.info("DKIM: %s", valid_dkim)

        message: Message = email.message_from_bytes(envelope.content)

        if not valid_dkim:
            return '550 DKIM validation failed'

        log.info('Message: %s', message)
        try:
            with SMTPClient(END_HOST, END_PORT) as client:
                client.sendmail(
                    from_addr=envelope.mail_from,
                    to_addrs=envelope.rcpt_tos,
                    msg=envelope.original_content
                )
        except BaseException as e:
            print(e)
            return '500 Could not process your message'

        return '250 Message accepted for delivery'


class EndMailServer:
    """
    Accepts mail only from end server. Checks if mail is signed to be from proxy, and sends to proxy if needed.
    SPF and DKIM needs to be already included.
    """

    async def handle_DATA(self,
                          server: SMTP,
                          session: Session,
                          envelope: Envelope):
        ip = session.peer[0]
        if ip not in VALID_END_ADDRS:
            return "521 Server doesn't accept mail"
        message: Message = email.message_from_bytes(envelope.content)
        log.info('Message: %s', message)
        mx_rcpt: Dict[str, list[str]] = {}
        for rcpt in envelope.rcpt_tos:
            _, _, domain = rcpt.partition("@")
            mx = get_mx(domain)
            if mx is None:
                continue
            mx_rcpt.setdefault(mx, []).append(rcpt)

        try:
            for mx, rcpts in mx_rcpt.items():
                if envelope.mail_from in VALID_PROXY_ADDRS:
                    with ProxiedSMTP(mx, 25) as client:
                        client.sendmail(
                            from_addr=envelope.mail_from,
                            to_addrs=rcpts,
                            msg=envelope.original_content
                        )
                else:
                    with SMTPClient(mx, 25) as client:
                        client.sendmail(
                            from_addr=envelope.mail_from,
                            to_addrs=rcpts,
                            msg=envelope.original_content
                        )
        except BaseException as e:
            print(e)
            return '500 Could not process your message'

        return '250 Message accepted for delivery'


class ProxyMailServer:
    """
    Accepts mail only from proxy server, and passes it to the real SMTP server, after proper SPF and DKIM checks.
    """

    async def handle_PROXY(self, server: SMTP, session: Session, envelope: Envelope, proxy_data: ProxyData):
        ip = session.peer[0]
        envelope.proxy_data = proxy_data
        return ip in VALID_PROXY_ADDRS

    async def handle_MAIL(self,
                          server: SMTP,
                          session: Session,
                          envelope: Envelope,
                          address: str,
                          mail_options: list):
        ip = envelope.proxy_data.src_addr
        result, description = spf.check2(ip, address, session.host_name)
        valid_spf = result == 'pass'
        envelope.spf = valid_spf

        log.info("SPF: %s, %s", result, description)

        if not valid_spf:
            return '550 SPF validation failed'

        envelope.mail_from = address
        envelope.mail_options.extend(mail_options)

        return '250 OK'

    async def handle_RCPT(self,
                          server: SMTP,
                          session: Session,
                          envelope: Envelope,
                          address: str,
                          rcpt_options: list):
        if address.split("@")[1] not in VALID_RECEPIENTS:
            return '550 not relaying to that domain'

        log.debug("Handle RCPT for %s", address)
        envelope.rcpt_tos.append(address)
        return '250 OK'

    async def handle_DATA(self,
                          server: SMTP,
                          session: Session,
                          envelope: Envelope):
        valid_dkim = dkim.verify(envelope.content)
        envelope.dkim = valid_dkim
        log.info("DKIM: %s", valid_dkim)

        message: Message = email.message_from_bytes(envelope.content)

        if not valid_dkim:
            return '550 DKIM validation failed'

        log.info('Message: %s', message)
        try:
            with SMTPClient(END_HOST, END_PORT) as client:
                client.sendmail(
                    from_addr=envelope.mail_from,
                    to_addrs=envelope.rcpt_tos,
                    msg=envelope.original_content
                )
        except BaseException as e:
            print(e)
            return '500 Could not process your message'

        return '250 Message accepted for delivery'


if __name__ == '__main__':
    host = os.getenv('SMTP_HOST', '*')
    port = int(os.getenv('SMTP_PORT', '25'))
    accept_host = os.getenv('ACCEPT_HOST')
    ssl_keys = os.getenv('SSL_KEYS')

    loop = asyncio.get_event_loop()

    end_handler = EndMailServer()
    proxy_handler = ProxyMailServer()
    generic_handler = GenericMailServer()

    ssl_context = None

    if ssl_keys:
        ssl_context = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
        ssl_context.load_cert_chain(ssl_keys + '.crt', ssl_keys + '.key')

    generic_controller = Controller(GenericMailServer, hostname=GENERIC_SMTP_HOST, port=GENERIC_SMTP_PORT)
    generic_controller.factory = lambda: SMTP(GenericMailServer, enable_SMTPUTF8=True, tls_context=ssl_context)
    generic_controller.start()
    log.info("Generic SMTP server started on %s:%s", GENERIC_SMTP_HOST, GENERIC_SMTP_PORT)

    end_controller = Controller(EndMailServer, hostname=END_SMTP_HOST, port=END_SMTP_PORT)
    end_controller.factory = lambda: SMTP(EndMailServer, enable_SMTPUTF8=True)
    end_controller.start()
    log.info("End SMTP server started on %s:%s", END_SMTP_HOST, END_SMTP_PORT)

    proxy_controller = Controller(ProxyMailServer, hostname=PROXY_SMTP_HOST, port=PROXY_SMTP_PORT)
    proxy_controller.factory = lambda: SMTP(ProxyMailServer, enable_SMTPUTF8=True, tls_context=ssl_context)
    proxy_controller.start()
    log.info("Proxy SMTP server started on %s:%s", PROXY_SMTP_HOST, PROXY_SMTP_PORT)

    try:
        loop.run_forever()
    except KeyboardInterrupt:
        print("Shutting down")
    finally:
        generic_controller.stop()
        end_controller.stop()
        proxy_controller.stop()
        loop.stop()
        loop.close()