From 0b8372c68cb4625459e73134c5ff80454810bd56 Mon Sep 17 00:00:00 2001
From: Balakrishnan Balasubramanian <git.user@balki.me>
Date: Mon, 12 Jun 2023 00:49:48 -0400
Subject: [PATCH] smtp refactor

---
 mail4one/config.py |  9 ++++--
 mail4one/smtp.py   | 80 +++++++++++++++++++++++++++++-----------------
 2 files changed, 58 insertions(+), 31 deletions(-)

diff --git a/mail4one/config.py b/mail4one/config.py
index 8e37aaf..87017a8 100644
--- a/mail4one/config.py
+++ b/mail4one/config.py
@@ -6,8 +6,8 @@ from jata import Jata, MutableDefault
 
 class Match(Jata):
     name: str
-    addrs: list[str] = MutableDefault(lambda: [])
-    addr_rexs: list[str] = MutableDefault(lambda: [])
+    addrs: list[str] = MutableDefault(lambda: []) # type: ignore
+    addr_rexs: list[str] = MutableDefault(lambda: []) # type: ignore
 
 
 DEFAULT_MATCH_ALL = "default_match_all"
@@ -100,3 +100,8 @@ def get_mboxes(addr: str, checks: list[Checker]) -> list[str]:
                     return
 
     return list(inner())
+
+
+def gen_addr_to_mboxes(cfg: Config) -> Callable[[str], [str]]:
+    checks = parse_checkers(cfg)
+    return lambda addr: get_mboxes(addr, checks)
diff --git a/mail4one/smtp.py b/mail4one/smtp.py
index 93a158b..a1e516b 100644
--- a/mail4one/smtp.py
+++ b/mail4one/smtp.py
@@ -3,41 +3,61 @@ import io
 import logging
 import mailbox
 import ssl
+import uuid
+import shutil
 from functools import partial
 from pathlib import Path
+from typing import Callable
 from . import config
+from email.message import Message
+import email.policy
+from email.generator import BytesGenerator
+import tempfile
 
-from aiosmtpd.handlers import Mailbox
+from aiosmtpd.handlers import Mailbox, AsyncMessage
 from aiosmtpd.smtp import SMTP, DATA_SIZE_DEFAULT
+from aiosmtpd.smtp import SMTP as SMTPServer
+from aiosmtpd.smtp import Envelope as SMTPEnvelope
+from aiosmtpd.smtp import Session as SMTPSession
 
 
-class MaildirCRLF(mailbox.Maildir):
-    _append_newline = True
+class MyHandler(AsyncMessage):
 
-    def _dump_message(self, message, target, mangle_from_=False):
-        temp_buffer = io.BytesIO()
-        super()._dump_message(message, temp_buffer, mangle_from_=mangle_from_)
-        temp_buffer.seek(0)
-        data = temp_buffer.read()
-        data = data.replace(b'\n', b'\r\n')
-        target.write(data)
+    def __init__(self, mails_path: Path, mbox_finder: Callable[[str], [str]]):
+        super().__init__()
+        self.mails_path = mails_path
+        self.mbox_finder = mbox_finder
+
+    async def handle_DATA(self, server: SMTPServer, session: SMTPSession,
+                          envelope: SMTPEnvelope) -> str:
+        self.rcpt_tos = envelope.rcpt_tos
+        return await super().handle_DATA(server, session, envelope)
+
+    async def handle_message(self, m: Message):  # type: ignore[override]
+        all_mboxes: set[str] = set()
+        for addr in self.rcpt_tos:
+            all_mboxes.union(self.mbox_finder(addr))
+        if not all_mboxes:
+            return
+        for mbox in all_mboxes:
+            for sub in ('new', 'tmp', 'cur'):
+                sub_path = self.mails_path / mbox / sub
+                sub_path.mkdir(mode=0o755, exist_ok=True, parents=True)
+        with tempfile.TemporaryDirectory() as tmpdir:
+            temp_email_path = Path(tmpdir) / f"{uuid.uuid4()}.eml"
+            with open(temp_email_path, "wb") as fp:
+                gen = BytesGenerator(fp, policy=email.policy.SMTP)
+                gen.flatten(m)
+            for mbox in all_mboxes:
+                shutil.copy(temp_email_path, self.mails_path / mbox / 'new')
 
 
-class MailboxCRLF(Mailbox):
-
-    def __init__(self, mail_dir: Path):
-        super().__init__(mail_dir)
-        for sub in ('new', 'tmp', 'cur'):
-            sub_path = mail_dir / sub
-            sub_path.mkdir(mode=0o755, exist_ok=True, parents=True)
-        self.mailbox = MaildirCRLF(mail_dir)
-
-
-def protocol_factory_starttls(dirpath: Path,
+def protocol_factory_starttls(mails_path: Path,
+                              mbox_finder: Callable[[str], [str]],
                               context: ssl.SSLContext | None = None):
-    logging.info("Got smtp client cb")
+    logging.info("Got smtp client cb starttls")
     try:
-        handler = MailboxCRLF(dirpath)
+        handler = MyHandler(mails_path, mbox_finder)
         smtp = SMTP(handler=handler,
                     require_starttls=True,
                     tls_context=context,
@@ -48,10 +68,10 @@ def protocol_factory_starttls(dirpath: Path,
     return smtp
 
 
-def protocol_factory(dirpath: Path):
+def protocol_factory(mails_path: Path, mbox_finder: Callable[[str], [str]]):
     logging.info("Got smtp client cb")
     try:
-        handler = MailboxCRLF(dirpath)
+        handler = MyHandler(mails_path, mbox_finder)
         smtp = SMTP(handler=handler, enable_SMTPUTF8=True)
     except Exception as e:
         logging.error("Something went wrong", e)
@@ -59,19 +79,21 @@ def protocol_factory(dirpath: Path):
     return smtp
 
 
-async def create_smtp_server_starttls(dirpath: Path,
+async def create_smtp_server_starttls(host: str,
                                       port: int,
-                                      host="",
+                                      mails_path: Path,
+                                      mbox_finder: Callable[[str], [str]],
                                       context: ssl.SSLContext | None = None):
     loop = asyncio.get_event_loop()
-    return await loop.create_server(partial(protocol_factory_starttls, dirpath,
-                                            context),
+    return await loop.create_server(partial(protocol_factory_starttls,
+                                            mails_path, mbox_finder, context),
                                     host=host,
                                     port=port,
                                     start_serving=False)
 
 
 async def create_smtp_server_tls(dirpath: Path,
+                                 mbox_finder: Callable[[str], [str]],
                                  port: int,
                                  host="",
                                  context: ssl.SSLContext | None = None):