#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# Copyright (C) 2019 tribe29 GmbH - License: GNU General Public License v2
# This file is part of Checkmk (https://checkmk.com). It is subject to the terms and
# conditions defined in the file COPYING, which is part of this source code package.
"""Tries to connect to an IMAP or POP3 server and returns OK if successfull.

"""

import argparse
import base64
import logging
import os
import re
import socket
import time
from pathlib import Path
from typing import List, Tuple, Union

from cmk.utils.mailbox import (
    active_check_main,
    Args,
    CheckResult,
    ConnectError,
    ForwardToECError,
    Mailbox,
    MailMessages,
)


def create_argument_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--forward-ec", action="store_true", help="Forward matched mails to the event console (EC)"
    )
    parser.add_argument(
        "--forward-method",
        type=str,
        metavar="METHOD",
        help=(
            "Configure how to connect to the event console to forward"
            "\nthe messages to. Can be configured to:"
            "\n    udp,<ADDR>,<PORT>        - Connect to remove EC via UDP"
            "\n    tcp,<ADDR>,<PORT>        - Connect to remove EC via TCP"
            "\n    spool:                   - Write to site local spool directory"
            "\n    spool:/path/to/spooldir  - Spool to given directory"
            "\n    /path/to/pipe            - Write to given EC event pipe"
            "\nDefaults to use the event console of the local OMD sites."
        ),
    )
    parser.add_argument(
        "--forward-facility",
        type=int,
        default=2,
        metavar="F",
        help="Syslog facility to use for forwarding (Defaults to '2' -> mail)",
    )
    parser.add_argument(
        "--forward-app",
        type=str,
        metavar="APP",
        help=(
            "Specify which string to use for the syslog application field "
            "when forwarding to the event console. "
            "\nYou can specify macros like \\1 or \\2 when you specified "
            '"--match-subject" with regex groups.'
            "\nDefaults to use the whole subject of the e mail"
        ),
    )
    parser.add_argument(
        "--forward-host", type=str, metavar="HOST", help="Hostname to use for the generated events"
    )
    parser.add_argument(
        "--body-limit",
        type=int,
        metavar="SIZE",
        default=1000,
        help=("Limit the number of characters of the body to forward " "(default=1000)"),
    )
    parser.add_argument(
        "--match-subject",
        type=str,
        metavar="REGEX",
        help="Use this option to not process all messages found in the inbox, "
        "but only the whones whose subject matches the given regular expression.",
    )
    parser.add_argument(
        "--cleanup",
        type=str,
        metavar="METHOD",
        help="Delete processed messages (see --match-subject) or move to subfolder "
        "a matching the given path. This is configured with the following "
        "METHOD: "
        "\n   delete             - Simply delete mails"
        "\n   path/to/subfolder  - Move to this folder (Only supported with IMAP)"
        "\nBy default the mails are not cleaned up, which might make your "
        "mailbox grow when you not clean it up manually.",
    )
    return parser


def syslog_time(localtime: time.struct_time = time.localtime()) -> str:
    """
    >>> syslog_time(time.strptime("Sep 22 21:20:19", "%b %d %H:%M:%S"))
    'Sep 22 21:20:19'
    """
    return time.strftime("%b %d %H:%M:%S", localtime)


def prepare_messages_for_ec(args: Args, mails: MailMessages) -> Tuple[List[str], List[int]]:
    # create syslog message from each mail
    # <128> Oct 24 10:44:27 Klappspaten /var/log/syslog: Oct 24 10:44:27 Klappspaten logger: asdasdad as
    # <facility+priority> timestamp hostname application: message
    messages: List[str] = []
    forwarded: List[int] = []
    cur_time = syslog_time()
    priority = 5  # OK

    for index, msg in sorted(mails.items()):
        subject = msg.get("Subject", "None")
        encoding = msg.get("Content-Transfer-Encoding", "None")
        log_line = subject
        # Now add the body to the event
        if msg.is_multipart():
            # only care for the first text/plain element
            for part in msg.walk():
                content_type = part.get_content_type()
                disposition = str(part.get("Content-Disposition"))
                encoding = part.get("Content-Transfer-Encoding", "None")
                if content_type == "text/plain" and "attachment" not in disposition:
                    payload = part.get_payload()
                    if encoding == "base64":
                        payload = str(base64.b64decode(payload))
                    log_line += "|" + payload[: args.body_limit]
                    break
        else:
            payload = msg.get_payload()
            if encoding == "base64":
                payload = str(base64.b64decode(payload))
            log_line += "|" + payload[: args.body_limit]

        log_line = log_line.replace("\r\n", "\0")
        log_line = log_line.replace("\n", "\0")

        # replace match groups in "forward_app"
        if args.forward_app:
            application = args.forward_app
            if args.match_subject:
                matches = re.match(args.match_subject, subject)
                for num, match in enumerate(matches.groups() if matches else []):
                    application = application.replace("\\%d" % (num + 1,), match)
        else:
            application = subject.replace("\n", "")

        # Construct the final syslog message
        log = "<%d>%s" % ((int(args.forward_facility) << 3) + priority, cur_time)
        log += " %s %s: %s" % (args.forward_host or args.fetch_server, application, log_line)
        messages.append(log)
        forwarded.append(index)
    return messages, forwarded


def forward_to_ec(args: Args, messages: List[str]) -> CheckResult:
    # send lines to event console
    # a) local in same omd site
    # b) local pipe
    # c) remote via udp
    # d) remote via tcp
    def evaluate_forward_method(description: str) -> Union[str, Tuple[str, str, str]]:
        if not description:
            return "%s/tmp/run/mkeventd/eventsocket" % os.getenv("OMD_ROOT", "")
        if description == "spool:":
            return "spool: %s/var/mkeventd/spool" % os.getenv("OMD_ROOT", "")
        if "," in description:
            prot_addr_port = description.split(",")
            return prot_addr_port[0], prot_addr_port[1], prot_addr_port[3]
        return description

    forward_method = evaluate_forward_method(args.forward_method)
    socket.setdefaulttimeout(args.connect_timeout)

    try:
        if messages:
            if isinstance(forward_method, tuple):
                # connect either via tcp or udp
                prot, addr, port = forward_method
                sock = (
                    socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
                    if prot == "udp"
                    else socket.socket(socket.AF_INET, socket.SOCK_STREAM)  #
                )
                sock.connect((addr, port))
                for message in messages:
                    sock.send(message.encode())
                    sock.send(b"\n")
                sock.close()

            elif not forward_method.startswith("spool:"):
                # write into local event pipe
                # Important: When the event daemon is stopped, then the pipe
                # is *not* existing! This prevents us from hanging in such
                # situations. So we must make sure that we do not create a file
                # instead of the pipe!
                sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
                try:
                    sock.connect(forward_method)
                except FileNotFoundError:
                    raise RuntimeError("%r does not exist" % forward_method)
                sock.send(("\n".join(messages)).encode())
                sock.send(b"\n")
                sock.close()

            else:
                # Spool the log messages to given spool directory.
                # First write a file which is not read into ec, then
                # perform the move to make the file visible for ec
                if not (spool_path := Path(forward_method[6:])).exists():
                    os.makedirs(spool_path)

                file_name = ".%s_%d_%d" % (args.forward_host, os.getpid(), time.time())
                path_spool_file = spool_path / file_name
                path_spool_file.write_text("\n".join(messages) + "\n")
                path_spool_file.rename(spool_path / file_name[1:])

        return (
            0,
            "Forwarded %d messages to event console" % len(messages),
            [("messages", len(messages))],
        )
    except Exception as exc:
        raise ForwardToECError(
            "Unable to forward messages to event console (%s). Left %d messages untouched."
            % (exc, len(messages))
        ) from exc


def check_mail(args: Args) -> CheckResult:
    """Create a mailbox an try to connect. In case of success our check will happily return OK
    If desired mails can also be fetched and forwarded to the event console (EC)
    """
    with Mailbox(args) as mailbox:
        try:
            mailbox.connect()
        except ConnectError as exc:
            # Handle this distinct error directly since it's the thing this check is all about
            return 2, str(exc), None

        if args.forward_ec:
            messages, forwarded = prepare_messages_for_ec(
                args,
                mailbox.fetch_mails(args.match_subject),
            )
            result = forward_to_ec(args, messages)
            if args.cleanup:
                if args.cleanup != "delete":
                    mailbox.copy_mails(forwarded, args.cleanup)
                mailbox.delete_mails(forwarded)

            return result

    return 0, "Successfully logged in to mailbox", None


if __name__ == "__main__":
    logging.getLogger().name = "check_mail"
    active_check_main(create_argument_parser(), check_mail)
