smtprd-ng/smtprd.py

#!/usr/bin/env python3

"""
Minimal SMTP forwarding relay daemon.
Accepts unsolicited mails via SMTP and sends them using a remote SMTP server as configured.
"""

import argparse
import asyncio
import configparser
import signal
import sys
import threading
import time
import uuid

from locale import setlocale, LC_ALL

from email import message_from_bytes, message_from_string
from email.message import Message
from email.policy import default as EmailMessagePolicy

from aiosmtpd.controller import Controller, SMTP as Server
from aiosmtpd.smtp import Session, Envelope
from aiosmtplib import SMTP, SMTPException, SMTPRecipientsRefused, SMTPRecipientRefused

from dataclasses import dataclass
from typing import Optional, Union, List, Tuple, Set


@dataclass(frozen=True)
class ServerConfig:
    hostname: str
    port: int


@dataclass(frozen=True)
class ClientConfig:
    hostname: str
    port: int
    sender: str
    recipients: List[str]
    username: Optional[str]
    password: Optional[str]
    set_reply_to: bool
    use_tls: bool
    start_tls: bool


@dataclass(frozen=True)
class Config:
    server: ServerConfig
    client: ClientConfig

    @classmethod
    def _from_config(cls, config: configparser.ConfigParser) -> 'Config':
        return cls(
            server=ServerConfig(
                hostname=config.get("server", "hostname", fallback="localhost"),
                port=config.getint("server", "port", fallback=8025),
            ),
            client=ClientConfig(
                hostname=config.get("client", "hostname"),
                port=config.getint("client", "port"),
                sender=config.get("client", "sender"),
                recipients=[_.strip() for _ in config.get("client", "recipients").split(",")],
                username=config.get("client", "username"),
                password=config.get("client", "password"),
                set_reply_to=config.getboolean("client", "set_reply_to", fallback=False),
                use_tls=config.getboolean("client", "use_tls", fallback=True),
                start_tls=config.getboolean("client", "start_tls", fallback=False),
            ),
        )

    @classmethod
    def from_ini(cls, filename: str) -> 'Config':
        try:
            config_parser: configparser.ConfigParser = configparser.ConfigParser()
            with open(filename, "r") as fp:
                config_parser.read_file(fp)
            return cls._from_config(config_parser)
        except (OSError, configparser.Error, UnicodeDecodeError) as e:
            raise RuntimeError(f"Cannot parse config '{filename}': {str(e)}") from None


class SMTPClient(SMTP):
    def __init__(self, config: ClientConfig) -> None:
        self._config: ClientConfig = config
        self._lock: asyncio.Lock = asyncio.Lock()
        super().__init__(
            hostname=self._config.hostname, port=self._config.port,
            username=self._config.username, password=self._config.password,
            use_tls=self._config.use_tls, start_tls=self._config.start_tls,
        )

    async def _send_message(self, message: Message, sender: str, recipients: List[str]) -> None:
        failed_recipients: Set[str] = set((await self.send_message(
            message, sender=sender, recipients=recipients
        ))[0].keys())
        if len(failed_recipients):  # raise also when not all have been refused
            raise SMTPRecipientsRefused([SMTPRecipientRefused(0, "", _) for _ in failed_recipients])

    async def forward_message(self, message: Message, sender: Optional[str], recipients: List[str]) -> None:
        del message["From"]
        del message["To"]
        del message["Reply-To"]
        message["Original-Sender"] = sender if sender is not None else ""
        message["Original-Recipient"] = ", ".join(recipients)
        message["From"] = self._config.sender
        message["To"] = ", ".join(self._config.recipients)
        if self._config.set_reply_to and sender is not None and sender != self._config.sender:
            message["Reply-To"] = sender

        print(f"{message['From']} ({message['Original-Sender']}) -> "
              f"{message['To']} ({message['Original-Recipient']}): "
              f"'{message.get('Subject', '')}'")

        async with self._lock:  # TODO: consumer task from spool queue, reusing connections
            try:
                await self.connect()
                await self._send_message(message, self._config.sender, self._config.recipients)
            except SMTPRecipientsRefused as e:
                raise RuntimeError(f"Recipients refused: {', '.join(_.recipient for _ in e.recipients)}") from e
            except (SMTPException, OSError) as e:
                raise RuntimeError(str(e)) from e
            finally:
                self.close()


class Handler:
    def __init__(self, client: SMTPClient) -> None:
        self._client: SMTPClient = client

    async def handle_DATA(self, server: Server, session: Session, envelope: Envelope) -> str:
        try:
            await self._client.forward_message(self._prepare_message(server, envelope),
                                               sender=envelope.mail_from,
                                               recipients=envelope.rcpt_tos)
        except (RuntimeError, ValueError) as e:
            print(f"Cannot forward: {str(e)}", file=sys.stderr)
            return f"550 {e.__cause__.__class__.__name__ if e.__cause__ is not None else e.__class__.__name__}"
        else:
            return "250 OK"

    @classmethod
    def _parse_message(cls, data: Optional[Union[bytes, str]]) -> Message:
        if isinstance(data, bytes):
            return message_from_bytes(data, policy=EmailMessagePolicy)
        elif isinstance(data, str):
            return message_from_string(data, policy=EmailMessagePolicy)
        else:
            raise ValueError(str(type(data)))

    def _prepare_message(self, server: Server, envelope: Envelope) -> Message:
        message: Message = self._parse_message(envelope.content)

        if server.transport is None:
            raise RuntimeError("Cannot get transport socket")
        peer: Tuple[str, int] = server.transport.get_extra_info("peername")
        sock: Tuple[str, int] = server.transport.get_extra_info("sockname")
        host: str = server.hostname  # getfqdn
        now: str = time.strftime("%a, %d %b %Y %H:%M:%S %z")

        message["Received"] = f"from {peer[0]} ([{peer[0]}]) by {host} ([{sock[0]}]); {now}"
        message["X-Peer"] = f"[{peer[0]}]:{peer[1]}"
        if "Date" not in message:
            message["Date"] = now
        if "Message-ID" not in message:
            message["Message-ID"] = f"<{str(uuid.uuid4())}@{host}>"

        return message


class SMTPServer(Controller):
    def __init__(self, config: ServerConfig, handler: Handler) -> None:
        super().__init__(handler=handler, hostname=config.hostname, port=config.port, ready_timeout=10.0)

    def run(self) -> bool:
        shutdown_requested: threading.Event = threading.Event()

        def _handler(signum: int, frame) -> None:
            shutdown_requested.set()

        signal.signal(signal.SIGINT, _handler)
        signal.signal(signal.SIGTERM, _handler)

        print(f"Starting on {self.hostname}:{self.port}", file=sys.stderr)
        self.start()
        shutdown_requested.wait()
        print("Stopping", file=sys.stderr)
        self.stop()
        return True


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__.strip(),
                                     formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument("--config", metavar="CONFIG.INI", type=str, default="./config.ini",
                        help="configuration file")
    args = parser.parse_args()
    setlocale(LC_ALL, "C")  # for strftime

    try:
        config: Config = Config.from_ini(args.config)
    except RuntimeError as e:
        print(str(e), file=sys.stderr)
        return 1
    else:
        controller = SMTPServer(config=config.server, handler=Handler(SMTPClient(config.client)))
        return 0 if controller.run() else 1


if __name__ == "__main__":
    sys.exit(main())