#!/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 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 envelope import Envelope as EnvelopeEnvelope


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

    """

    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_and_sign(self, message: Message) -> Message:
        # Currently does NOT work
        new_message = (
            EnvelopeEnvelope()
            .smime()
            .load(message)
            .sign(
                key=Path(self._config.smime_cert_private),
                cert=Path(self._config.smime_cert),
            )
            .as_message()
        )
        new_message["Original-Sender"] = message["Original-Sender"]
        new_message["Original-Recipient"] = message["Original-Recipient"]
        new_message["From"] = message["From"]
        new_message["To"] = message["To"]
        return new_message

    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 including the From and To Header
        """
        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])
        )
        # Add correct headers
        new = b"From: " + self._config.sender.encode() + b"\r\n" + output
        new = b"To: " + ", ".join(self._config.recipients).encode() + b"\r\n" + new
        new = b"Subject: " + message.get("Subject", "").encode() + b"\r\n" + new
        return new

    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._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 main() -> int:
    """Main routine

    Returns:
        int: exit code
    """
    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
    controller = SMTPServer(
        config=config.server, handler=Handler(SMTPClient(config.client))
    )
    return 0 if controller.run() else 1


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