#!/usr/bin/env python3
import os
import sys
import logging
from logging import handlers
import asyncio
import signal
from pprint import pformat

from argparse import ArgumentParser, RawDescriptionHelpFormatter

logging.basicConfig(
    stream=sys.stdout,
    level=logging.ERROR,
    format="%(levelname)-7.7s [%(asctime)s] %(message)s",
    datefmt="%Y-%m-%d %H:%M",
)
log = logging.getLogger("spoofguard")


def main():
    args = parse_args()
    asyncio.run(run_server(args))


async def run_server(args):
    spoofchecker = SpoofChecker(args)
    HOST = spoofchecker.listen_host
    PORT = spoofchecker.listen_port
    log.debug(f"Starting server on {HOST}:{PORT}...")
    pol_server = await asyncio.start_server(spoofchecker.handle_socket, HOST, PORT)

    await asyncio.gather(
        asyncio.ensure_future(pol_server.serve_forever()), return_exceptions=True
    )


class SpoofChecker:
    def __init__(self, args):
        self._read_config_and_cli_args(args)

        self.stats = {}
        self.cache = {
            "sender_login_map": {},
        }
        self.superusers = set()
        self._reload_mapping_file()
        signal.signal(signal.SIGHUP, self._reload_mapping_file)

    def _read_config_and_cli_args(self, args):
        cfg = args.config
        self.listen_host = cfg.get("DEFAULT", "interface", fallback="127.0.0.1")
        self.listen_port = cfg.getint("DEFAULT", "port", fallback=4242)

        # Get mapping file from config, maybe overriden by cli args
        self.mapping_file = args.mapping_file
        if not self.mapping_file:
            self.mapping_file = os.path.join(
                os.path.dirname(cfg.get("DEFAULT", "config_path")),
                cfg.get("DEFAULT", "mapping_file", fallback="sender_login_map"),
            )

        # No need to go further if we don't have a mapping file
        if not os.path.isfile(self.mapping_file):
            print(f"Cannot read mapping file {self.mapping_file}. Stopping...")
            sys.exit(1)

        # Configure logging output stream
        logstream = args.log_stream or cfg.get(
            "DEFAULT", "logstream", fallback="console"
        )
        if logstream == "syslog":
            logging.getLogger().removeHandler(logging.getLogger().handlers[0])
            mail_facility = logging.handlers.SysLogHandler.LOG_MAIL
            maillog_handler = logging.handlers.SysLogHandler(
                facility=mail_facility, address="/dev/log"
            )
            maillog_handler.setFormatter(
                logging.Formatter("%(name)s: %(levelname)s %(message)s")
            )
            log.addHandler(maillog_handler)
        else:
            log.addHandler(logging.getLogger().handlers[0])
            logging.getLogger().removeHandler(logging.getLogger().handlers[0])

        # Set log level
        log.setLevel(logging.WARNING)  # Default log level
        if args.verbose > 0:
            if args.verbose >= 2:
                log.setLevel(logging.DEBUG)
            elif args.verbose == 1:
                log.setLevel(logging.INFO)
        else:
            log.setLevel(
                logging.getLevelName(
                    cfg.get("DEFAULT", "loglevel", fallback=logging.WARNING)
                )
            )

        # Set dry-run mode
        self.dry_run = args.dry_run or cfg.getboolean(
            "DEFAULT", "dry-run", fallback=False
        )

    def _reload_mapping_file(self, *args, **kwargs):
        if len(args) and args[0] == signal.SIGHUP:
            log.info(f"Received SIGHUP. Reloading {self.mapping_file}...")
        else:
            log.info(f"Loading {self.mapping_file}...")
        superusers_before = self.superusers
        self.superusers = set()
        _res = {}
        with open(self.mapping_file, "r") as fin:
            for line in fin.readlines():
                line = line.strip()
                if not line:
                    continue
                if line.startswith("# *"):
                    self.superusers.update(line.strip("# *").split())
                    continue
                if line.startswith("#"):
                    continue
                items = line.split()
                key = items[0]
                vals = items[1:]
                if key in _res:
                    log.info(
                        f"Mapping file contains duplicate key: '{key}'. Merging..."
                    )
                else:
                    _res[key] = set()
                _res[key].update(vals)
        self.cache["sender_login_map"] = _res
        if self.superusers != superusers_before:
            log.info(f"superusers: {self.superusers}")

    async def handle_socket(self, reader, writer):
        while True:
            try:
                req = (await reader.readuntil(separator=b"\n\n")).decode("utf8").strip()
                log.debug(f"Received request:\n{req}")
                attrs = self.parse_request(req)
                if "request" not in attrs:
                    log.error("Missing key 'request'.")
                if attrs.get("request") == "monitoring":
                    res = "status OK"
                else:
                    res = "action=%s\n\n" % self.smtpd_access_policy(attrs)
                    log.info(res.strip())
                writer.write(res.encode("utf-8") + b"\n\n")
                await writer.drain()
            except asyncio.exceptions.IncompleteReadError as e:
                last_data_received = e.partial.decode("utf-8").strip()
                if last_data_received:
                    log.error(
                        f"Remote side closed the connection. Last data received: '{last_data_received}'"
                    )
                writer.close()
                break
            except Exception as e:
                log.exception("Unexpected error while handling request", exc_info=True)
                break

    def parse_request(self, data):
        attrs = {}
        for line in data.split("\n"):
            line = line.strip()
            if "=" not in line:
                log.warning(f"Ignoring invalid line: '{line}'")
                continue
            key, val = line.split("=", 1)
            attrs[key.strip()] = val.strip()
        return attrs

    def smtpd_access_policy(self, attrs):
        """
        This is where the decision is taken.

        :attrs: dict received from postfix
        """
        log.debug("Handle mail policy")

        if "sasl_username" not in attrs:
            log.warning("Attr 'sasl_username' is missing. Ignoring.")
            return "DUNNO"

        action = "DUNNO"
        message = ""
        user = attrs.get("sasl_username")
        sender = attrs.get("sender")
        queue_id = attrs.get("queue_id", "???")

        # Special use case. FIXME: add config parameter instead of hard-coded values
        if user.endswith("@smtp-02.laforet.com"):  # XXX HARD-CODED values
            user = user.replace("@smtp-02.laforet.com", "@laforet.com")

        # Build a readable string for client ip/port for logging
        client = "unknown client"
        if "client_address" in attrs and "client_port" in attrs:
            client = "{client_address}:{client_port}".format(**attrs)

        if not user:
            action = "DUNNO"
            message = f"client '{client}' sent an email as '{sender}'"
        elif user in self.superusers:
            action = "DUNNO"
            message = f"superuser '{user}' sent an email as '{sender}'"
        elif sender.lower() == user.lower():
            action = "DUNNO"
            message = f"sasl_username matches sender: '{sender}'"
        elif not sender:
            # Weird use case, but apparently very common: empty sender.
            # Let's allow this, for now...
            rcpt = attrs.get("recipient", "???")
            action = "DUNNO"
            message = f"user '{user}' sent an email with empty sender address from client '{client}' to '{rcpt}'"
        elif sender not in self.cache["sender_login_map"]:
            action = "REJECT"
            message = f"user '{user}' using an unknown sender address: '{sender}'"
        elif user not in self.cache["sender_login_map"].get(sender, []):
            action = "REJECT"
            message = f"user '{user}' is not allowed to send emails as '{sender}'"
        else:
            action = "DUNNO"
            message = f"user '{user}' sent an email as '{sender}'"

        if self.dry_run:
            if action != "DUNNO":
                message = f"(would be {action} without dry-run mode) {message}"
            action = "DUNNO"
        if message:
            message = f"-- {queue_id} {message}"
        return f"{action} {message}"


def parse_args():
    parser = ArgumentParser(
        prog="spoofguard", formatter_class=RawDescriptionHelpFormatter
    )
    parser.add_argument(
        "-v", "--verbose", action="count", default=0, help="Print debug messages"
    )
    parser.add_argument(
        "-n",
        "--dry-run",
        action="store_true",
        help="Don't really apply modifications, just simulate.",
    )
    parser.add_argument(
        "-c",
        "--config",
        type=read_config,
        default="",
        help="Read configuration from the specified file.",
    )
    parser.add_argument(
        "-f",
        "--mapping-file",
        help="Read sender login map from this file",
    )
    parser.add_argument(
        "--log-stream",
        choices=["console", "syslog"],
        help="Output log on 'console' (stdout) or to 'syslog' (mail facility). Default is 'console'",
    )

    args = parser.parse_args()
    if args.verbose > 1:
        log.setLevel(logging.DEBUG)
        logging.getLogger().handlers[0].setFormatter(
            logging.Formatter(
                "%(levelname)-8.8s %(asctime)-19.19s [%(module)s.py l.%(lineno)s] %(message)s"
            )
        )
    elif args.verbose > 0:
        log.setLevel(logging.INFO)

    return args


def read_config(configfile=None):
    """Returns a ``ConfigParser`` object."""

    from configparser import ConfigParser

    conf = ConfigParser()

    possible_configfiles = [
        "./spoofguard.ini",
        "~/.config/spoofguard.ini",
        "/etc/spoofguard.ini",
        "/etc/spoofguard/spoofguard.ini",
    ]

    if configfile is None or not os.path.isfile(configfile):
        for candidate in possible_configfiles:
            candidate = os.path.expanduser(candidate)
            if os.path.isfile(candidate):
                configfile = candidate
                log.debug("Found configuration file %s." % configfile)
                break

    if configfile:
        conf.read(configfile)
        conf.set("DEFAULT", "config_path", os.path.realpath(configfile))
    else:
        conf.read(possible_configfiles)
        conf.set("DEFAULT", "config_path", "/etc/spoofguard/missing_config.ini")

    return conf


if __name__ == "__main__":
    main()
