#!/usr/bin/python
# encoding: utf-8
# +------------------------------------------------------------------+
# |             ____ _               _        __  __ _  __           |
# |            / ___| |__   ___  ___| | __   |  \/  | |/ /           |
# |           | |   | '_ \ / _ \/ __| |/ /   | |\/| | ' /            |
# |           | |___| | | |  __/ (__|   <    | |  | | . \            |
# |            \____|_| |_|\___|\___|_|\_\___|_|  |_|_|\_\           |
# |                                                                  |
# | Copyright Mathias Kettner 2013             mk@mathias-kettner.de |
# +------------------------------------------------------------------+
#
# This file is part of Check_MK.
# The official homepage is at http://mathias-kettner.de/check_mk.
#
# check_mk is free software;  you can redistribute it and/or modify it
# under the  terms of the  GNU General Public License  as published by
# the Free Software Foundation in version 2.  check_mk is  distributed
# in the hope that it will be useful, but WITHOUT ANY WARRANTY;  with-
# out even the implied warranty of  MERCHANTABILITY  or  FITNESS FOR A
# PARTICULAR PURPOSE. See the  GNU General Public License for more de-
# ails.  You should have  received  a copy of the  GNU  General Public
# License along with GNU Make; see the file  COPYING.  If  not,  write
# to the Free Software Foundation, Inc., 51 Franklin St,  Fifth Floor,
# Boston, MA 02110-1301 USA.

import sys, getopt, os, time, random, poplib, imaplib, email
import traceback, socket, re

def parse_exception(e):
    e = str(e)
    if e[0] == '{':
        e = "%d - %s" % eval(e).values()[0]
    return str(e)

def bail_out(rc, s, perfdata=[]):
    stxt = ['OK', 'WARN', 'CRIT', 'UNKNOWN'][rc]
    sys.stdout.write('%s - %s' % (stxt, parse_exception(s)))
    if perfdata:
        sys.stdout.write(' | %s' % (' '.join(['%s=%s' % (p[0], ';'.join(map(str, p[1:]))) for p in perfdata ])))
    sys.stdout.write('\n')
    sys.exit(rc)

def usage(msg=None):
    if msg:
        sys.stderr.write('ERROR: %s\n' % msg)
    sys.stderr.write("""
USAGE: check_mail [OPTIONS]

OPTIONS:
  --protocol PROTO      Set to "IMAP" or "POP3", depending on your mailserver
                        (defaults to IMAP)
  --server ADDRESS      Host address of the IMAP/POP3 server hosting your mailbox
  --port PORT           IMAP or POP3 port
                        (defaults to 110 for POP3 and 995 for POP3 with SSL and
                         143 for IMAP and 993 for IMAP with SSL)
  --username USER       Username to use for IMAP/POP3
  --password PW         Password to use for IMAP/POP3
  --ssl                 Use SSL for feching the mailbox (disabled by default)
  --connect-timeout     Timeout in seconds for network connects (defaults to 10)

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

  -d, --debug           Enable debug mode
  -h, --help            Show this help message and exit

""")
    sys.exit(1)

short_options = 'dh'
long_options  = ['protocol=', 'server=', 'port=', 'username=', 'password=', 'ssl',
    'connect-timeout=', 'cleanup=', 'forward-ec', 'match-subject=',
    'forward-facility=', 'forward-app=', 'forward-method=', 'forward-host=', 'body-limit=',
    'help', 'debug', ]

required_params = [
    'server', 'username', 'password',
]

try:
    opts, args = getopt.getopt(sys.argv[1:], short_options, long_options)
except getopt.GetoptError, err:
    sys.stderr.write("%s\n" % err)
    sys.exit(1)

opt_debug        = False
fetch_proto      = 'IMAP'
fetch_server     = None
fetch_port       = None
fetch_user       = None
fetch_pass       = None
fetch_ssl        = False
conn_timeout     = 10
cleanup_messages = False
forward_ec       = False
forward_facility = 16 # default to "mail" (2 << 3)
forward_app      = None
forward_method   = None # local event console
forward_host     = ''
match_subject    = None
body_limit       = 1000

g_M         = None
g_forwarded = []


for o,a in opts:
    if o in [ '-h', '--help' ]:
        usage()
    elif o in [ '-d', '--debug' ]:
        opt_debug = True
    elif o == '--protocol':
        fetch_proto = a
    elif o == '--server':
        fetch_server = a
    elif o == '--port':
        fetch_port = int(a)
    elif o == '--username':
        fetch_user = a
    elif o == '--password':
        fetch_pass = a
    elif o == '--ssl':
        fetch_ssl = True
    elif o == '--connect-timeout':
        conn_timeout = int(a)
    elif o == '--cleanup':
        cleanup_messages = a
    elif o == '--forward-ec':
        forward_ec = True
    elif o == '--match-subject':
        match_subject = re.compile(a)
    elif o == '--forward-facility':
        forward_facility = int(a) << 3
    elif o == '--forward-app':
        forward_app = a
    elif o == '--forward-method':
        if ',' in a:
            forward_method = tuple(a.split(','))
        else:
            forward_method = a
    elif o == '--forward-host':
        forward_host = a
    elif o == '--body-limit':
        body_limit = int(a)

param_names = dict(opts).keys()
for p in required_params:
    if '--'+p not in param_names:
        usage('The needed parameter --%s is missing' % p)

if fetch_proto not in [ 'IMAP', 'POP3' ]:
    usage('The given protocol is not supported.')

if fetch_port == None:
    if fetch_proto == 'POP3':
        fetch_port = fetch_ssl and 995 or 110
    else:
        fetch_port = fetch_ssl and 993 or 143

def connect():
    global g_M
    try:
        if fetch_proto == 'POP3':
            fetch_class = fetch_ssl and poplib.POP3_SSL or poplib.POP3
            g_M = fetch_class(fetch_server, fetch_port)
            g_M.user(fetch_user)
            g_M.pass_(fetch_pass)
        else:
            fetch_class = fetch_ssl and imaplib.IMAP4_SSL or imaplib.IMAP4
            g_M = fetch_class(fetch_server, fetch_port)
            g_M.login(fetch_user, fetch_pass)
            g_M.select('INBOX', readonly=False) # select INBOX
    except Exception, e:
        if opt_debug:
            raise
        bail_out(3, 'Failed connect to %s:%d: %s' % (fetch_server, fetch_port, parse_exception(e)))

def fetch_mails():
    mails = {}
    try:
        # Get mails from mailbox
        if fetch_proto == 'POP3':
            num_messages = len(g_M.list()[1])
            for i in range(num_messages):
                index = i+1
                lines = g_M.retr(index, 0)[1]
                mails[i] = email.message_from_string("".join(lines))
        else:
            retcode, messages = g_M.search(None, 'NOT', 'DELETED')
            if retcode == 'OK' and messages[0].strip():
                for num in messages[0].split(' '):
                    try:
                        ty, data = g_M.fetch(num, '(RFC822)')
                        mails[num] = email.message_from_string(data[0][1])
                    except Exception, e:
                        raise Exception('Failed to fetch mail %s (%s). Available messages: %r' % (num, parse_exception(e), messages))

        if match_subject:
            # Now filter out the messages not wanted to be handled by this check
            for index, msg in mails.items():
                matches = match_subject.match(msg.get('Subject', ''))
                if not matches:
                    del mails[index]

        return mails
    except Exception, e:
        if opt_debug:
            raise
        bail_out(3, 'Failed to check for mails: %s' % parse_exception(e))

def cleanup_mailbox():
    if not g_M:
        return # do not deal with mailbox when none sent yet
    try:
        # Do not delete all messages in the inbox. Only the ones which were
        # processed before! In the meantime there might be occured new ones.
        for index in g_forwarded:
            if fetch_proto == 'POP3':
                if cleanup_messages == 'delete':
                    response = g_M.dele(index)
                if response != "+OK":
                    raise Exception("Response from server: [%s]" % response)
            else:
                if cleanup_messages != 'delete':
                    # The user wants the message to be moved to the folder
                    # refered by the string stored in "cleanup_messages"
                    folder = cleanup_messages.strip('/')

                    # Create maybe missing folder hierarchy
                    target = ''
                    for level in folder.split('/'):
                        target += "%s/" % level
                        g_M.create(target)

                    # Copy the mail
                    ty, data= g_M.copy(str(index), folder)
                    if ty != 'OK':
                        raise Exception("Response from server: [%s]" % data)

                # Now delete the mail
                ty, data = g_M.store(index, '+FLAGS', '\\Deleted')
                if ty != 'OK':
                    raise Exception("Response from server: [%s]" % data)

        if fetch_proto == 'IMAP':
            g_M.expunge()
    except Exception, e:
        if opt_debug:
            raise
        bail_out(2, 'Failed to delete mail: %s' % parse_exception(e))

def close_mailbox():
    if not g_M:
        return # do not deal with mailbox when none sent yet
    if fetch_proto == 'POP3':
        g_M.quit()
    else:
        g_M.close()
        g_M.logout()

def syslog_time():
    localtime = time.localtime()
    day = int(time.strftime("%d", localtime)) # strip leading 0
    value = time.strftime("%b %%d %H:%M:%S", localtime)
    return value % day

def forward_to_ec(mails):
    # 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 = []
    cur_time = syslog_time()
    priority = 5 # OK
    for index, msg in mails.items():
        subject = msg.get('Subject', 'None')
        log_line = subject
        # Now add the body to the event
        if msg.is_multipart():
            for payload in msg.get_payload():
                log_line += '|' + payload.get_payload()[:body_limit]
        else:
            log_line += '|' + msg.get_payload()[:body_limit]
        log_line = log_line.replace('\r\n', '\0')
        log_line = log_line.replace('\n', '\0')

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

        # Construct the final syslog message
        log = '<%d>%s' % (forward_facility + priority, cur_time)
        log += ' %s %s: %s' % (forward_host or fetch_server, application, log_line)
        messages.append(log)
        g_forwarded.append(index)

    # send lines to event console
    # a) local in same omd site
    # b) local pipe
    # c) remote via udp
    # d) remote via tcp
    global forward_method
    if not forward_method:
        forward_method = os.getenv('OMD_ROOT') + "/tmp/run/mkeventd/eventsocket"
    elif forward_method == 'spool:':
        forward_method += os.getenv('OMD_ROOT') + "/var/mkeventd/spool"

    try:
        if messages:
            if isinstance(forward_method, tuple):
                # connect either via tcp or udp
                if forward_method[0] == 'udp':
                    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
                else:
                    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                sock.connect((forward_method[1], forward_method[2]))
                for message in messages:
                    sock.send(message + "\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)
                sock.connect(forward_method)
                sock.send('\n'.join(messages) + '\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
                spool_path = forward_method[6:]
                file_name  = '.%s_%d_%d' % (forward_host, os.getpid(), time.time())
                if not os.path.exists(spool_path):
                    os.makedirs(spool_path)
                file('%s/%s' % (spool_path, file_name), 'w').write('\n'.join(messages) + '\n')
                os.rename('%s/%s' % (spool_path, file_name), '%s/%s' % (spool_path, file_name[1:]))

        if cleanup_messages:
            cleanup_mailbox()

        bail_out(0, 'Forwarded %d messages to event console' % len(messages), [('messages', len(messages))])
    except Exception, e:
        bail_out(3, 'Unable to forward messages to event console (%s). Left %d messages untouched.' % (e, len(messages)))

def main():
    # Enable showing protocol messages of imap for debugging
    if opt_debug:
        imaplib.Debug = 4

    connect()
    if forward_ec:
        forward_to_ec(fetch_mails())
    else:
        bail_out(0, 'Successfully logged in to mailbox')

socket.setdefaulttimeout(conn_timeout)

try:
    main()
except Exception, e:
    if opt_debug:
        raise
    bail_out(2, 'Unhandled exception: %s' % parse_exception(e))
