check_mail_loop

#!/usr/bin/env python
# -*- encoding: utf-8; py-indent-offset: 4 -*-
# +------------------------------------------------------------------+
# |             ____ _               _        __  __ _  __           |
# |            / ___| |__   ___  ___| | __   |  \/  | |/ /           |
# |           | |   | '_ \ / _ \/ __| |/ /   | |\/| | ' /            |
# |           | |___| | | |  __/ (__|   <    | |  | | . \            |
# |            \____|_| |_|\___|\___|_|\_\___|_|  |_|_|\_\           |
# |                                                                  |
# | Copyright Mathias Kettner 2014             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-
# tails. 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 email
import email.utils
import email.mime.text
import getopt
import imaplib
import os
import ast
import poplib
import random
import re
import smtplib
import socket
import sys
import time

import cmk.utils.password_store

cmk.utils.password_store.replace_passwords()


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


def bail_out(rc, s):
    stxt = ['OK', 'WARN', 'CRIT', 'UNKNOWN'][rc]
    sys.stdout.write('%s - %s\n' % (stxt, s))
    sys.exit(rc)


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

OPTIONS:
  --smtp-server ADDRESS   Host address of the SMTP server to send the mail to
  --smtp-port PORT        Port to use for SMTP (defaults to 25)
  --smtp-username USER    Username to use for SMTP communictation
                          (leave empty for anonymous SMTP)
  --smtp-password PW      Password to authenticate SMTP
  --smtp-tls              Use TLS over SMTP (disabled by default)
  --imap-tls              Use TLS for IMAP authentification (disabled by default)

  --fetch-protocol PROTO  Set to "IMAP" or "POP3", depending on your mailserver
                          (defaults to IMAP)
  --fetch-server ADDRESS  Host address of the IMAP/POP3 server hosting your mailbox
  --fetch-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)
  --fetch-username USER   Username to use for IMAP/POP3
  --fetch-password PW     Password to use for IMAP/POP3
  --fetch-ssl             Use SSL for feching the mailbox (disabled by default)

  --mail-from MAIL        Use this mail address as sender address
  --mail-to MAIL          Use this mail address as recipient address

  --warning AGE           Loop duration of the most recent mail in seconds or
                          the average of all received mails within a single
                          check to raise a WARNING state
  --critical AGE          Loop duration of the most recent mail in seconds or
                          the average of all received mails within a single
                          check to raise a CRITICAL state

  --connect-timeout       Timeout in seconds for network connects (defaults to 10)
  --status-dir PATH       This plugin needs a file to store information about
                          sent, received and expected mails. Defaults to either
                          /tmp/ or /omd/sites/<sitename>/var/check_mk when executed
                          from within an OMD site
  --status-suffix SUFFIX  Concantenated with "check_mail_loop.SUFFIX.status" to
                          generate the name of the status file. Empty by default
  --delete-messages       Delete all messages identified as being related to this
                          check plugin. This is disabled by default, which
                          might make your mailbox grow when you not clean it up
                          manually.
  --subject SUBJECT       You can specify the subject text. If choosen the subject
                          'Check_MK-Mail-Loop' will be replace by the stated text.

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

""")
    sys.exit(1)


short_options = 'dh'
long_options = [
    'smtp-server=',
    'smtp-port=',
    'smtp-username=',
    'smtp-password=',
    'smtp-tls',
    'imap-tls',
    'fetch-protocol=',
    'fetch-server=',
    'fetch-port=',
    'fetch-username=',
    'fetch-password=',
    'fetch-ssl',
    'mail-from=',
    'mail-to=',
    'warning=',
    'critical=',
    'connect-timeout=',
    'delete-messages',
    'help',
    'status-dir=',
    'status-suffix=',
    'subject=',
    "debug",
]

required_params = [
    'smtp-server',
    'fetch-server',
    'fetch-username',
    'fetch-password',
    'mail-from',
    'mail-to',
]

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

opt_debug = False
smtp_server = None
smtp_port = None
smtp_user = None
smtp_pass = None
smtp_tls = False
fetch_proto = 'IMAP'
fetch_server = None
fetch_port = None
fetch_user = None
fetch_pass = None
fetch_ssl = False
imap_tls = False
mail_from = None
mail_to = None
warning = None
critical = 3600
conn_timeout = 10
delete_messages = False
status_dir = None
status_suffix = None
subject = 'Check_MK-Mail-Loop'
for o, a in opts:
    if o in ['-h', '--help']:
        usage()
    elif o in ['-d', '--debug']:
        opt_debug = True
    elif o == '--smtp-server':
        smtp_server = a
    elif o == '--smtp-port':
        smtp_port = int(a)
    elif o == '--smtp-username':
        smtp_user = a
    elif o == '--smtp-password':
        smtp_pass = a
    elif o == '--smtp-tls':
        smtp_tls = True
    elif o == '--fetch-protocol':
        fetch_proto = a
    elif o == '--fetch-server':
        fetch_server = a
    elif o == '--fetch-port':
        fetch_port = int(a)
    elif o == '--fetch-username':
        fetch_user = a
    elif o == '--fetch-password':
        fetch_pass = a
    elif o == '--imap-tls':
        imap_tls = True
    elif o == '--fetch-ssl':
        fetch_ssl = True
    elif o == '--mail-from':
        mail_from = a
    elif o == '--mail-to':
        mail_to = a
    elif o == '--warning':
        warning = int(a)
    elif o == '--critical':
        critical = int(a)
    elif o == '--connect-timeout':
        conn_timeout = int(a)
    elif o == '--delete-messages':
        delete_messages = True
    elif o == '--status-dir':
        status_dir = a
    elif o == '--status-suffix':
        status_suffix = a
    elif o == '--subject':
        subject = a

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

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

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

if not status_dir:
    status_dir = os.environ.get('OMD_ROOT')
    if status_dir:
        status_dir += '/var/check_mk'
    else:
        status_dir = '/tmp'
if status_suffix:
    status_path = '%s/check_mail_loop.%s.status' % (status_dir, status_suffix)
else:
    status_path = '%s/check_mail_loop.status' % (status_dir)

socket.setdefaulttimeout(conn_timeout)

g_expected = {}
g_mails = {}
g_obsolete = {}
g_M = None


def load_expected_mails():
    try:
        for line in file(status_path):
            ts, key = line.rstrip().split(' ', 1)
            g_expected[ts + '-' + key] = (int(ts), int(key))
    except IOError:
        pass  # Skip errors on not existing file


def add_expected_msg(ts, key):
    g_expected[str(ts) + '-' + str(key)] = (int(ts), key)


def save_expected_mails():
    lines = []
    for ts, key in g_expected.values():
        lines.append('%d %s' % (ts, key))
    file(status_path, 'w').write('\n'.join(lines) + '\n')


def add_starttls_support(self, keyfile=None, certfile=None):
    import ssl
    name = "STARTTLS"
    typ, dat = self._simple_command(name)
    if typ != 'OK':
        raise self.error(dat[-1])

    self.sock = ssl.wrap_socket(self.sock)
    self.file = self.sock.makefile()

    cap = 'CAPABILITY'
    self._simple_command(cap)
    if not cap in self.untagged_responses:
        raise self.error('no CAPABILITY response from server')
    self.capabilities = tuple(self.untagged_responses[cap][-1].upper().split())


def fetch_mails():
    global g_M
    if not g_expected:
        return  # not expecting any mail, do not check for mails

    mails = {}
    try:
        if fetch_proto == 'POP3':
            # Get mails from POP3 mailbox
            fetch_class = poplib.POP3_SSL if fetch_ssl else poplib.POP3
            g_M = fetch_class(fetch_server, fetch_port)
            g_M.user(fetch_user)
            g_M.pass_(fetch_pass)

            num_messages = len(g_M.list()[1])

            for i in range(num_messages):
                index = i + 1
                lines = g_M.retr(index)[1]
                mails[i] = email.message_from_string("\n".join(lines))

        else:
            # Get mails from IMAP mailbox
            fetch_class = imaplib.IMAP4_SSL if fetch_ssl else imaplib.IMAP4

            if imap_tls:
                # starttls in imaplib is only available with python >= 3.2
                # we are going to implement our own version
                imaplib.Commands.update({"STARTTLS": ("NONAUTH")})
                g_M = imaplib.IMAP4(fetch_server, fetch_port)
                add_starttls_support(g_M)
            else:
                g_M = fetch_class(fetch_server, fetch_port)

            g_M.login(fetch_user, fetch_pass)
            g_M.select('INBOX', readonly=False)  # select INBOX
            retcode, messages = g_M.search(None, 'NOT', 'DELETED')
            if retcode == 'OK' and messages[0].strip():
                for num in messages[0].split(' '):
                    try:
                        data = g_M.fetch(num, '(RFC822)')[1]
                        mails[num] = email.message_from_string(data[0][1])
                    except Exception as e:
                        raise Exception('Failed to fetch mail %s (%s). Available messages: %r' %
                                        (num, parse_exception(e), messages))

        # Now filter out the messages for this check
        pattern = re.compile(r'(?:Re: |WG: )?%s ([^\s]+) ([^\s]+)' % subject)
        for index, msg in mails.items():
            matches = pattern.match(msg.get('Subject', ''))
            if matches:
                ts = matches.group(1).strip()
                key = matches.group(2).strip()

                # extract received time
                rx = msg.get('Received')
                if rx:
                    rx_ts = email.utils.mktime_tz(email.utils.parsedate_tz(rx.split(';')[-1]))
                else:
                    # use current time as fallback where no Received header could be found
                    rx_ts = int(time.time())

                if "%s-%s" % (ts, key) not in g_expected:
                    # Delete any "Check_MK-Mail-Loop" messages older than 24 hours, even if they are not in our list
                    if delete_messages and int(time.time()) - rx_ts > 24 * 3600:
                        g_obsolete[ts + '-' + key] = (index, rx_ts)
                    continue

                g_mails[ts + '-' + key] = (index, rx_ts)

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


def send_mail():
    now = time.time()
    key = random.randint(1, 1000)

    mail = email.mime.text.MIMEText("")
    mail['From'] = mail_from
    mail['To'] = mail_to
    mail['Subject'] = '%s %d %d' % (subject, now, key)
    mail['Date'] = email.utils.formatdate(localtime=True)

    try:
        S = smtplib.SMTP(smtp_server, smtp_port)
        if smtp_tls:
            S.starttls()
        if smtp_user:
            S.login(smtp_user, smtp_pass)
        S.sendmail(mail_from, mail_to, mail.as_string())
        S.quit()

        add_expected_msg(now, key)
    except Exception as e:
        if opt_debug:
            raise
        bail_out(3, 'Failed to send mail: %s' % parse_exception(e))


def check_mails():
    state = 0
    perfdata = []
    output = []

    num_received = 0
    num_pending = 0
    num_lost = 0
    duration = None
    now = time.time()

    # Loop all expected mails and check whether or not they have been received
    for ident, (send_ts, _unused_key) in g_expected.items():
        if ident in g_mails:
            recv_ts = g_mails[ident][1]

            if duration is None:
                duration = recv_ts - send_ts
            else:
                duration = (duration + (recv_ts - send_ts)) / 2  # average

            if critical is not None and duration >= critical:
                state = 2
                output.append(' (>= %d)' % critical)
            elif warning is not None and duration >= warning:
                state = max(state, 1)
                output.append(' (>= %d)' % warning)

            del g_expected[ident]  # remove message from expect list
            num_received += 1
            # FIXME: Also remove older mails which have not yet been seen?

        else:
            # drop expecting messages when older than critical threshold,
            # but keep waiting for other mails which have not yet reached it
            if now - send_ts >= critical:
                del g_expected[ident]
                num_lost += 1
                state = 2
            else:
                num_pending += 1

    if num_received == 1:
        output.insert(0, 'Mail received within %d seconds' % duration)
        perfdata.append(('duration', duration, warning or '', critical or ''))
    elif num_received > 1:
        output.insert(0,
                      'Received %d mails within average of %d seconds' % (num_received, duration))
        perfdata.append(('duration', duration, warning or '', critical or ''))
    else:
        output.insert(0, 'Did not receive any new mail')

    if num_lost:
        output.append('Lost: %d (Did not arrive within %d seconds)' % (num_lost, critical))

    if num_pending > 1:
        output.append('Currently waiting for %d mails' % num_pending)

    return state, output, perfdata


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 mail_index, _unused_recv_ts in g_mails.values() + g_obsolete.values():
            if fetch_proto == 'POP3':
                response = g_M.dele(mail_index + 1)
                if not response.startswith("+OK"):
                    raise Exception("Response from server: [%s]" % response)
            else:
                ty = g_M.store(mail_index, '+FLAGS', '\\Deleted')[0]
                if ty != 'OK':
                    raise Exception("Response from server: [%s]" % response)

        if fetch_proto == 'IMAP':
            g_M.expunge()
    except Exception as 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 main():
    # Enable showing protocol messages of imap for debugging
    if opt_debug:
        imaplib.Debug = 4

    load_expected_mails()

    fetch_mails()
    send_mail()

    state, output, perfdata = check_mails()

    if delete_messages:
        cleanup_mailbox()
    close_mailbox()

    save_expected_mails()

    sys.stdout.write(', '.join(output))
    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(state)


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