remove Controller

mail4one/pop3.py

@@ -82,7 +82,6 @@ async def auth_stage():
 
 
 MAILS_PATH = ""
-WAIT_FOR_PRIVILEGES_TO_DROP = None
 
 
 async def transaction_stage(user: User):
@@ -138,9 +137,6 @@ def delete_messages(delete_ids):
 
 
 async def new_session(stream_reader: asyncio.StreamReader, stream_writer: asyncio.StreamWriter):
-    if WAIT_FOR_PRIVILEGES_TO_DROP:
-        logging.warning("Waiting for privileges to drop")
-        await WAIT_FOR_PRIVILEGES_TO_DROP
     reader.set(stream_reader)
     writer.set(stream_writer)
     logging.info(f"New session started with {stream_reader} and {stream_writer}")
@@ -153,23 +149,25 @@ async def new_session(stream_reader: asyncio.StreamReader, stream_writer: asynci
     except ClientError as c:
         write(err("Something went wrong"))
         logging.error(f"Unexpected client error", c)
-    except:
+    except Exception as e:
-        logging.error(f"Serious client error")
+        logging.error(f"Serious client error", e)
         raise
     finally:
         stream_writer.close()
 
 
-async def a_main(dirpath: Path, port: int, host="", context: ssl.SSLContext = None, waiter=None):
+async def create_pop_server(dirpath: Path, port: int, host="", context: ssl.SSLContext = None):
     logging.info(
-        f"Starting POP3 server Maildir={dirpath}, host={host}, port={port}, context={context}, waiter={waiter}")
+        f"Starting POP3 server Maildir={dirpath}, host={host}, port={port}, context={context}")
-    global MAILS_PATH, WAIT_FOR_PRIVILEGES_TO_DROP
+    global MAILS_PATH
     MAILS_PATH = dirpath / 'new'
-    WAIT_FOR_PRIVILEGES_TO_DROP = waiter
+    return await asyncio.start_server(new_session, host=host, port=port, ssl=context)
-    server = await asyncio.start_server(new_session, host=host, port=port, ssl=context)
+
+
+async def a_main(*args, **kwargs):
+    server = await create_pop_server(*args, **kwargs)
     await server.serve_forever()
 
 
 if __name__ == "__main__":
-    # noinspection PyTypeChecker
     asyncio.run(a_main(Path("/tmp/mails"), 9995))

mail4one/server.py

@@ -3,22 +3,15 @@ import asyncio
 # When running on privilege port after dropping privileges.
 # noinspection PyUnresolvedReferences
 import encodings.idna
-import io
 import logging
-import mailbox
 import os
 import ssl
 import sys
 from argparse import ArgumentParser
-from functools import partial
 from pathlib import Path
 
-from aiosmtpd.controller import Controller
+from .smtp import create_smtp_server
-from aiosmtpd.handlers import Mailbox
+from .pop3 import create_pop_server
-from aiosmtpd.main import DATA_SIZE_DEFAULT
-from aiosmtpd.smtp import SMTP
-
-from .pop3 import a_main as pop3_main
 
 
 def create_tls_context(certfile, keyfile):
@@ -27,49 +20,6 @@ def create_tls_context(certfile, keyfile):
     return context
 
 
-class STARTTLSController(Controller):
-    def __init__(self, *args, tls_context, smtp_args=None, **kwargs):
-        self.tls_context = tls_context
-        self.smtp_args = smtp_args or {}
-        self.has_privileges_dropped: asyncio.Future = None
-        if 'ssl_context' in kwargs:
-            raise Exception("ssl_context not allowed when using STARTTLS, set tls_context instead")
-        Controller.__init__(self, *args, **kwargs)
-
-    async def create_future(self):
-        self.has_privileges_dropped = asyncio.get_event_loop().create_future()
-
-    async def wait_for_privileges_to_drop(self):
-        await self.has_privileges_dropped
-
-    def factory(self):
-        if not self.has_privileges_dropped.done():
-            # Ideally  we should await here. But this is callback and not a coroutine
-            raise Exception("Client connected too fast before we could drop root privileges")
-        return SMTP(self.handler, require_starttls=True, tls_context=self.tls_context, **self.smtp_args)
-
-
-class MaildirCRLF(mailbox.Maildir):
-    _append_newline = True
-
-    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)
-
-
-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 parse_args():
     parser = ArgumentParser()
     parser.add_argument('--certfile')
@@ -83,8 +33,6 @@ def parse_args():
     args.host = '0.0.0.0'
     args.smtp_port = 25
     args.pop_port = 995
-    args.size = DATA_SIZE_DEFAULT
-    args.classpath = MailboxCRLF
     args.smtputf8 = True
     args.debug = True
     return args
@@ -97,7 +45,7 @@ def setup_logging(args):
         logging.basicConfig(level=logging.INFO)
 
 
-def drop_privileges(future_cb):
+def drop_privileges():
     try:
         import pwd
     except ImportError:
@@ -111,32 +59,24 @@ def drop_privileges(future_cb):
         logging.error("Cannot setuid nobody; run as root")
         sys.exit(1)
     logging.info("Dropped privileges")
-    future_cb().set_result("Go!")
     logging.debug("Signalled! Clients can come in")
 
 
+async def a_main(args, tls_context):
+    # pop_server = await create_pop_server(args.mail_dir_path, port=args.pop_port, host=args.host, context=tls_context)
+    smtp_server = await create_smtp_server(args.mail_dir_path, port=args.smtp_port, host=args.host, context=tls_context)
+    drop_privileges()
+    # await pop_server.start_serving()
+    await smtp_server.serve_forever()
+
+
 def main():
     args = parse_args()
     tls_context = create_tls_context(args.certfile, args.keyfile)
-    smtp_args = dict(data_size_limit=args.size, enable_SMTPUTF8=args.smtputf8)
     setup_logging(args)
-    handler = args.classpath(args.mail_dir_path)
     loop = asyncio.get_event_loop()
     loop.set_debug(args.debug)
-    controller = STARTTLSController(
+    asyncio.run(a_main(args, tls_context))
-        handler, tls_context=tls_context, smtp_args=smtp_args, hostname=args.host, port=args.smtp_port, loop=loop)
-
-    loop.create_task(controller.create_future())
-    loop.create_task(pop3_main(args.mail_dir_path, args.pop_port,
-                               host=args.host, context=tls_context, waiter=controller.wait_for_privileges_to_drop()))
-
-    controller.start()
-    loop.call_soon_threadsafe(partial(drop_privileges, lambda: controller.has_privileges_dropped))
-    logging.info("Server started. Press [ENTER] to stop")
-    input()
-    controller.stop()
-    # loop.create_task(a_main(controller))
-    # loop.run_forever()
 
 
 if __name__ == '__main__':

mail4one/smtp.py

@@ -0,0 +1,62 @@
+import asyncio
+import io
+import logging
+import mailbox
+import ssl
+from functools import partial
+from pathlib import Path
+
+from aiosmtpd.handlers import Mailbox
+from aiosmtpd.smtp import SMTP, DATA_SIZE_DEFAULT
+
+
+class MaildirCRLF(mailbox.Maildir):
+    _append_newline = True
+
+    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)
+
+
+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(dirpath: Path, context: ssl.SSLContext = None):
+    logging.info("Got smtp client cb")
+    try:
+        handler = MailboxCRLF(dirpath)
+        smtp = SMTP(handler=handler,
+                    require_starttls=True,
+                    tls_context=context,
+                    data_size_limit=DATA_SIZE_DEFAULT,
+                    enable_SMTPUTF8=True)
+    except Exception as e:
+        logging.error("Something went wrong", e)
+        raise
+    return smtp
+
+
+async def create_smtp_server(dirpath: Path, port: int, host="", context: ssl.SSLContext = None):
+    loop = asyncio.get_event_loop()
+    return await loop.create_server(partial(protocol_factory, dirpath, context),
+                                    host=host, port=port, start_serving=False)
+
+
+async def a_main(*args, **kwargs):
+    server = await create_smtp_server(*args, **kwargs)
+    await server.serve_forever()
+
+
+if __name__ == "__main__":
+    logging.basicConfig(level=logging.DEBUG)
+    asyncio.run(a_main(Path("/tmp/mails"), 9995))