#!/usr/bin/env python3

#     SMTP forwarding relay daemon with signing and encryption
#
#     Copyright (C) 2024  F. Schröder (See https://www.hackitu.de/smtprd/smtprd.py.html)
#     Copyright (C) 2024  F. Brandes (additions to original code)
#
#     This program is free software: you can redistribute it and/or modify
#     it under the terms of the GNU Affero General Public License as
#     published by the Free Software Foundation, either version 3 of the
#     License, or (at your option) any later version.
#
#     This program is distributed in the hope that it will be useful,
#     but WITHOUT ANY WARRANTY; without even the implied warranty of
#     MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#     GNU Affero General Public License for more details.
#
#     You should have received a copy of the GNU Affero General Public License
#     along with this program.  If not, see <http://www.gnu.org/licenses/>.

"""
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 dataclasses import dataclass
from email import message_from_bytes, message_from_string
from email.message import Message
from email.policy import default as EmailMessagePolicy
from locale import LC_ALL, setlocale
from pathlib import Path
from typing import List, Optional, Set, Tuple, Union

from aiosmtpd.controller import SMTP as Server
from aiosmtpd.controller import Controller
from aiosmtpd.smtp import Envelope, Session
from aiosmtplib import SMTP, SMTPException, SMTPRecipientRefused, SMTPRecipientsRefused
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import padding
from cryptography.hazmat.primitives.serialization import (
    Encoding,
    load_pem_private_key,
    pkcs7,
)
from cryptography.x509 import load_pem_x509_certificate
from M2Crypto import BIO, SMIME, X509


@dataclass(frozen=True)
class ServerConfig:
    """Serverconfig"""

    hostname: str
    port: int


# pylint: disable=too-many-instance-attributes
@dataclass(frozen=True)
class ClientConfig:
    """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
    smime_cert: str
    smime_cert_private: str
    smime_to_cert: str


@dataclass(frozen=True)
class Config:
    """Config main module

    Raises:
        RuntimeError: When config file cannot be parsed

    """

    server: ServerConfig
    client: ClientConfig

    @classmethod
    def _from_config(cls, config: configparser.RawConfigParser) -> "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),
                smime_cert=config.get("client", "smime_cert", fallback=""),
                smime_cert_private=config.get(
                    "client", "smime_cert_private", fallback=""
                ),
                smime_to_cert=config.get("client", "smime_to_cert", fallback=""),
            ),
        )

    @classmethod
    def from_ini(cls, filename: str) -> "Config":
        """Opens and reads .INI config file

        Args:
            filename (str): config file path

        Raises:
            RuntimeError: on error in config file

        Returns:
            Config: ConfigParser object
        """
        try:
            config_parser: configparser.ConfigParser = configparser.ConfigParser()
            with open(filename, "r", encoding="utf8") 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):
    """Client part of library

    This will connect to an upstream SMTP server to deliver the mail

    """

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

    def _encrypt(self, message: bytes, subject: str) -> bytes:
        """Encrypt the message

        Args:
            message (bytes): message in bytes format (can/should be signed)

        Returns:
            bytes: Encrypted message
        """
        buf = BIO.MemoryBuffer(message)
        s = SMIME.SMIME()

        # Load target cert to encrypt to.
        x509 = X509.load_cert(self._config.smime_to_cert)
        sk = X509.X509_Stack()
        sk.push(x509)
        s.set_x509_stack(sk)

        s.set_cipher(SMIME.Cipher("aes_256_cbc"))

        # Encrypt the buffer.
        p7 = s.encrypt(buf)

        # Output p7 in mail-friendly format.
        out = BIO.MemoryBuffer()

        # Add header
        out.write("From: " + self._config.sender + "\r\n")
        out.write("To: " + ", ".join(self._config.recipients) + "\r\n")
        out.write("Subject: " + subject + "\r\n")
        s.write(out, p7)

        return out.read()

    def _sign(self, message: Message) -> bytes:
        """Sign the message

        This function signs the message. All existing headers will be moved
        inside the signature boundary.
        Args:
            message (Message): message object

        Returns:
            bytes: The signed message
        """
        with open(Path(self._config.smime_cert_private), "rb") as key_data:
            key = load_pem_private_key(key_data.read(), password=None)
        with open(self._config.smime_cert, "rb") as cert_data:
            cert = load_pem_x509_certificate(cert_data.read())

        # sign
        output = (
            pkcs7.PKCS7SignatureBuilder()
            .set_data(message.as_bytes())
            .add_signer(cert, key, hashes.SHA512(), rsa_padding=padding.PKCS1v15())
            .sign(Encoding.SMIME, [pkcs7.PKCS7Options.DetachedSignature])
        )
        return output

    def _encrypt_and_sign(self, message: Message) -> bytes:
        """Sign and encrypt the message

        Args:
            message (Message): message object

        Returns:
            bytes: The signed and encrypted message including the From and To Header
        """
        signed = self._sign(message)
        encrypt = self._encrypt(signed, message.get("Subject", ""))
        return encrypt

    async def _send_message(
        self, message: bytes, sender: str, recipients: List[str]
    ) -> None:
        # if we use send_message, it expects a message object
        # This will add a newline after the signing process. The message object is fine
        # but as soon as _as_bytes or _as_string is called, it will add a new line
        # between the boundary of the signed message, which will ruin the signature
        failed_recipients: Set[str] = set(
            (
                await self.sendmail(
                    sender,
                    recipients,
                    message,
                )
            )[0].keys()
        )
        if len(failed_recipients) > 0:  # 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:
        """Modify message to be passed on.

        Remove original From, To and Reply-To and replace with "Original-X" header.
        Call _sign function, if there is a cert saved in the config.
        Pass it on to the actual send function

        Args:
            message (Message): message
            sender (Optional[str]): Original sender
            recipients (List[str]): Original recipients

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

        print(
            f"{self._config.sender} ({message['Original-Sender']}) -> "
            f"{', '.join(self._config.recipients)} ({message['Original-Recipient']}): "
            f"'{message.get('Subject', '')}'"
        )
        if self._config.smime_cert:
            message = self._encrypt_and_sign(message)
        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()


# pylint: disable=too-few-public-methods
class Handler:
    """Reimplement Handler"""

    def __init__(self, client: SMTPClient) -> None:
        """Impelent our SMTPClient"""
        self._client: SMTPClient = client

    async def handle_DATA(  # pylint: disable=invalid-name
        self,
        server: Server,
        session: Session,  # pylint: disable=unused-argument
        envelope: Envelope,
    ) -> str:
        """Forward the message from our server to our client function

        Args:
            server (Server): Our server
            session (Session): Current session (Not used here)
            envelope (Envelope): message envelope

        Returns:
            str: Status to client
        """
        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__}"
        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)
        if isinstance(data, str):
            return message_from_string(data, policy=EmailMessagePolicy)
        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):
    """Implementation of local SMTP Server 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:
        """Run routine

        Returns:
            bool: True on finished routine
        """
        shutdown_requested: threading.Event = threading.Event()

        def _handler(signum: int, frame) -> None:  # pylint: disable=unused-argument
            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 parse_args(args=None) -> argparse.ArgumentParser.parse_args:
    """Parse arguments

    Parameters
    ----------
    args : List of strings

    Returns
    -------
    argparse.ArgumentParser.parse_args : Namespace of arguments
    """
    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",
    )
    return parser.parse_args(args)


def main(args=None) -> int:
    """Main routine

    Returns:
        int: exit code
    """
    args = parse_args(args)
    setlocale(LC_ALL, "C")  # for strftime
    if len(args.config) == 0:
        raise OSError("No config file supplied")
    if not Path(args.config).is_file():
        raise OSError("Config file not found: " + str(args.config))
    try:
        config: Config = Config.from_ini(args.config)
    except RuntimeError as e:
        print(str(e), file=sys.stderr)
        return 1
    controller = SMTPServer(
        config=config.server, handler=Handler(SMTPClient(config.client))
    )
    return 0 if controller.run() else 1


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