remove Controller

This commit is contained in:
Balakrishnan Balasubramanian 2018-12-18 22:58:52 -05:00
parent dd321ab8b7
commit bcfaeb1569
3 changed files with 84 additions and 84 deletions

View File

@ -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:
logging.error(f"Serious client error")
except Exception as e:
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}")
global MAILS_PATH, WAIT_FOR_PRIVILEGES_TO_DROP
f"Starting POP3 server Maildir={dirpath}, host={host}, port={port}, context={context}")
global MAILS_PATH
MAILS_PATH = dirpath / 'new'
WAIT_FOR_PRIVILEGES_TO_DROP = waiter
server = await asyncio.start_server(new_session, host=host, port=port, ssl=context)
return 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))

View File

@ -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 aiosmtpd.handlers import Mailbox
from aiosmtpd.main import DATA_SIZE_DEFAULT
from aiosmtpd.smtp import SMTP
from .pop3 import a_main as pop3_main
from .smtp import create_smtp_server
from .pop3 import create_pop_server
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(
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()
asyncio.run(a_main(args, tls_context))
if __name__ == '__main__':

62
mail4one/smtp.py Normal file
View File

@ -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))