smtp-gee.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
'''
This lives in https://github.com/treibholz/smtp-gee

Yes, the code is bad, has mixed styles and I also want to burn it, rewrite
everything from scratch, but I don't have time for it, so I just add features
like to a car from the movie Mad Max.

It has to get the job done!

'''

import smtplib
import configparser
import time
import datetime
import hashlib
import socket
import imaplib
import argparse
import sys
import logging
import locale

from email.mime.text import MIMEText
import email

locale.setlocale(locale.LC_ALL, 'C')
log = logging.getLogger(__name__)
socket.setdefaulttimeout(30)


def parse_received_time(received_header):
    received_time = received_header.split(';')[-1]
    return email.utils.parsedate_to_datetime(received_time)


class Account(object):
    """docstring for Account"""
    def __init__(self, name, login=False, password=False, smtp_server="localhost", imap_server="localhost", imap_folder="INBOX", smtp_over_ssl=False, smtp_port=25):
        # super(Account, self).__init__()
        self.name           = name
        self.login          = login
        self.password       = password
        self.smtp_server    = smtp_server
        self.smtp_port      = smtp_port
        self.imap_server    = imap_server
        self.imap_folder    = imap_folder
        self.email          = login
        self.smtp_timeout   = 30
        self.imap_timeout   = 30
        self.imap_messages  = {}
        self.smtp_transfer_time = 0

        self.debug        =   False
        self.smtp_over_ssl  =   smtp_over_ssl
        self.error_string   =   ""
        self.current_folder =   ""

    def send(self, recipient, stopwatch=None):
        timestamp = time.time()
        payload =   f"""Hi,
                    this is a testmail, generated by SMTP-GEE.

                    sent on:   {socket.getfqdn()}
                    sent at:   {timestamp}
                    sent from: {self.email}
                    sent to:   {recipient.email}

                    Cheers.
                        SMTP-GEE"""

        # Stupid cleanup, to get nicer indenting...
        payload = payload.replace("                    ", "")

        test_id = hashlib.sha1(payload.encode('utf-8')).hexdigest()
        msg = MIMEText(payload)

        msg['From']     =   self.email
        msg['To']       =   recipient.email
        msg['Subject']  =   f"[SMTP-GEE] |{test_id}"
        msg['x-smtp-gee'] = socket.getfqdn() 

        try:
            if self.smtp_over_ssl:
                log.debug(f'SMTP: Connecting with TLS - {self.smtp_server}:{self.smtp_port}')
                s = smtplib.SMTP_SSL( self.smtp_server, port = self.smtp_port, timeout = self.smtp_timeout)
            else:
                log.debug(f'SMTP: Connecting with STARTTLS - {self.smtp_server}:{self.smtp_port}')
                s = smtplib.SMTP( self.smtp_server, port = self.smtp_port, timeout = self.smtp_timeout)
                s.starttls()
            if stopwatch != None:
                stopwatch.stage('CONNECT')

            log.debug('SMTP: Logging in')
            s.login(self.login, self.password )
            if stopwatch != None:
                stopwatch.stage('LOGIN')
            log.debug('SMTP: Login OK')

            s.sendmail( self.email, recipient.email, msg.as_string() )
            s.quit()

            return test_id
        except smtplib.SMTPAuthenticationError as smtp_error:
            self.error_string += f"SMTPAuthenticationError: {smtp_error}"
            return False
        except smtplib.SMTPConnectError as smtp_error:
            self.error_string += f"SMTPConnectError: {smtp_error}"
            return False
        except:
            self.error_string += f"Unexpected error: {sys.exc_info()}"
            return False


    def check(self, check_id, stopwatch=None):
        """Check via IMAP"""
        log.debug("Checking via IMAP")
        
        # get the check_date
        check_day = datetime.date.today()
        if datetime.datetime.now().hour == 0:
            check_day -= datetime.timedelta(days = 1)
        check_date = check_day.strftime('%d-%b-%Y')
        log.debug(f'IMAP: Check Mails since {check_date}') 

        initial_sleep_time = 1
        sleep_time = initial_sleep_time
        try:
            m = imaplib.IMAP4_SSL(self.imap_server)
            if self.debug:
                m.debug = 10
            if stopwatch != None:
                stopwatch.stage('CONNECT')

            log.debug("IMAP: LOGIN")
            m.login(self.login, self.password)
            if stopwatch != None:
                stopwatch.stage('LOGIN')
            log.debug("IMAP: LOGIN OK")

            data=[b'']
            found = False
            while data == [b'']:
                if stopwatch != None:
                    if stopwatch.gettime() > self.imap_timeout:
                        return False
                for current_folder in [f.strip() for f in self.imap_folder.split(',')]:
                    if not found:
                        typ, n = m.select(current_folder)

                        if typ != "OK":
                            self.error_string += "Can't select {current_folder}: {typ}"
                            return False
                        
                        msg_count = int(n[0])

                        log.debug(f"IMAP: {current_folder} {typ} ({msg_count})")
                        # Searching headers might be slow...
                        # typ, data = m.search(None, 'SUBJECT', f'"{check_id}"')
                    
                        log.debug(f"IMAP: {current_folder}  fetching all since {check_date}")
                        typ, _messages = m.search(None, f'(ALL) SINCE {check_date}')
                        if typ == 'OK':
                            messages = _messages[0].split()
                            messages.sort(reverse=True)
                            log.debug(f'IMAP: {_messages}')
                            self.imap_messages[current_folder] = len(messages)
                            for msg in messages:
                                if not found:
                                    msg_typ, subject = m.fetch(msg, 'BODY.PEEK[HEADER.FIELDS (SUBJECT)]')
                                    # check the Subject.
                                    if check_id in str(subject):
                                        log.debug(f"IMAP: {current_folder} Received!")
                                        found = True
                                        self.current_folder = current_folder
                                        data = subject
                                        num = subject[0][0].split()[0]
                        else:
                            log.debug(f'IMAP: {typ}, {messages}')

                if not found:
                    log.debug(f'IMAP: sleeping {sleep_time:0.1f} sec')
                    time.sleep(sleep_time)
                    if sleep_time <= 5:
                        sleep_time += 0.1

            # Fetch the whole mail, just for measuring
            message = m.fetch(num, '(RFC822)')[1][0][1]
            mail = email.message_from_bytes(message)
            received_headers = mail.get_all('received')
            smarthost_received = parse_received_time(received_headers[-1])
            mx_received = parse_received_time(received_headers[0])
            self.smtp_transfer_time = (mx_received-smarthost_received).total_seconds()

            # deleting should be more sophisticated, for debugging...
            m.store(num, '+FLAGS', r'\Deleted')
            m.expunge()
            self.imap_messages[self.current_folder] -= 1
            m.close()
            m.logout()

            return True
        except imaplib.IMAP4.error as imap_error:
            self.error_string += f"IMAP error: {imap_error}"
            return False
        except:
            self.error_string += f"Unexpected error: {sys.exc_info()}"
            return False


class Stopwatch(object):
    def __init__(self, debug=False):
        super(Stopwatch, self).__init__()
        self.__debug = debug
        self.__start   = -1
        self.counter = 0
        self.stage_counter = {}

    def gettime(self):
        return time.time() - self.__start

    def start(self):
        self.__start = time.time()

    def stop(self):
        self.counter += time.time() - self.__start
        self.__start  = -1

    def stage(self, stage_name):
        self.stage_counter[stage_name] = time.time() - self.__start



if __name__ == "__main__":

    # fallback returncode
    returncode = 3

    parser = argparse.ArgumentParser(
        description='Check how long it takes to send a mail (by SMTP) and how long it takes to find it in the IMAP-inbox',
        epilog = "Because e-mail is a realtime-medium and you know it!")


    main_parser_group = parser.add_argument_group('Main options')
    main_parser_group.add_argument('--from', dest='sender', action='store',
                    required=True,
                    metavar="<name>",
                    help='The account to send the message')

    main_parser_group.add_argument('--rcpt', dest='rcpt', action='store',
                    required=True,
                    metavar="<name>",
                    help='The account to receive the message')

    main_parser_group.add_argument('--nagios', dest='nagios', action='store_true',
                    required=False,
                    default=False,
                    help='output in Nagios mode')

    main_parser_group.add_argument('--prometheus', dest='prometheus', action='store_true',
                    required=False,
                    default=False,
                    help='output in Prometheus mode')

    main_parser_group.add_argument('--promextra', dest='promextra', action='store_true',
                    required=False,
                    default=False,
                    help='output addtional metrics')



    main_parser_group.add_argument('--except-means', dest='except_return', action='store',
                    metavar="<int>",
                    required=False,
                    default=2,
                    help='Map Exceptions to another returncode. Default: %(default)s')

    main_parser_group.add_argument('--debug', dest='debug', action='store_true',
                    required=False,
                    default=False,
                    help='Debug mode')

    main_parser_group.add_argument('--config',dest='config_file', action='store',
                    default='config.ini',
                    metavar="<file>",
                    required=False,
                    help='alternate config-file')


    smtp_parser_group = parser.add_argument_group('SMTP options')
    smtp_parser_group.add_argument('--smtp_warn', dest='smtp_warn', action='store',
                    required=False,
                    default=15,
                    metavar="<sec>",
                    type=int,
                    help='warning threshold to send the mail. Default: %(default)s')

    smtp_parser_group.add_argument('--smtp_crit', dest='smtp_crit', action='store',
                    required=False,
                    default=30,
                    metavar="<sec>",
                    type=int,
                    help='critical threshold to send the mail. Default: %(default)s')

    smtp_parser_group.add_argument('--smtp_timeout', dest='smtp_timeout', action='store',
                    required=False,
                    default=30,
                    metavar="<sec>",
                    type=int,
                    help='timeout to stop sending a mail. Default: %(default)s')


    imap_parser_group = parser.add_argument_group('IMAP options')
    imap_parser_group.add_argument('--imap_warn', dest='imap_warn', action='store',
                    required=False,
                    default=20,
                    metavar="<sec>",
                    type=int,
                    help='warning threshold until the mail appears in the INBOX. Default: %(default)s')

    imap_parser_group.add_argument('--imap_crit', dest='imap_crit', action='store',
                    required=False,
                    default=30,
                    metavar="<sec>",
                    type=int,
                    help='critical threshold until the mail appears in the INBOX. Default: %(default)s')

    imap_parser_group.add_argument('--imap_timeout', dest='imap_timeout', action='store',
                    required=False,
                    default=30,
                    metavar="<sec>",
                    type=int,
                    help='timeout to stop waiting for a mail to appear in the INBOX (not implemented yet). Default: %(default)s')


    args = parser.parse_args()

    # Read Config

    c = configparser.ConfigParser()
    c.read(args.config_file)

    if args.debug:
        logging.basicConfig(level=logging.DEBUG)
        log.setLevel(logging.DEBUG)
    
    a={}

    for s in c.sections():
        a[s] = Account(s)
        if args.debug:
            log.debug("Set to DEBUG")
            a[s].debug = True

        # This has to be more easy...
        a[s].smtp_server    = c.get(s, 'smtp_server')
        a[s].imap_server    = c.get(s, 'imap_server')
        a[s].password       = c.get(s, 'password')
        a[s].login          = c.get(s, 'login')
        a[s].email          = c.get(s, 'email')
        a[s].smtp_timeout   = args.smtp_timeout
        a[s].imap_timeout   = args.imap_timeout

        # FIXME: Really, really Ugly!
        a[s].smtp_port = 25
        try:
            a[s].smtp_over_ssl = c.get(s, 'smtp_over_ssl')
            a[s].smtp_port = 465
        except:
            pass

        try:
            a[s].smtp_port = c.get(s, 'smtp_port')
        except:
            pass

        a[s].imap_folder    = 'INBOX'
        try:
            a[s].imap_folder = c.get(s, 'imap_folder')
        except:
            pass

    ### Here the real work begins  ###

    # Create the stopwatches.
    smtp_time = Stopwatch()
    imap_time = Stopwatch()

    # send the mail by SMTP
    smtp_time.start()
    smtp_result = a[args.sender].send(a[args.rcpt], smtp_time)
    smtp_time.stop()

    log.debug(f'smtp_result: {smtp_result}')

    imap_result = False
    if smtp_result:

        # Receive the mail.
        imap_time.start()
        imap_result = a[args.rcpt].check(smtp_result, stopwatch=imap_time)
        imap_time.stop()

        log.debug(f'imap_result: {imap_result}')


    ### Present the results

    if args.prometheus:
        if smtp_result:
            smtp_state="success"
        elif smtp_time.counter >= args.smtp_timeout:
            smtp_state="timeout"
        else:
            smtp_state="error"

        if imap_result:
            imap_state="success"
        elif imap_time.counter >= args.imap_timeout:
            imap_state="timeout"
        elif smtp_state != "success":
            imap_state = "NoError"
        else:
            imap_state="error"


        print('# HELP smtp_gee exports metrics about SMTP and IMAP duration')
        print('# TYPE smtp_gee gauge')
        print(f'smtp_gee{{protocol="SMTP",from="{args.sender}",rcpt="{args.rcpt}",state="{smtp_state}",error_string="{a[args.sender].error_string}"}} {smtp_time.counter}')
        print(f'smtp_gee{{protocol="IMAP",from="{args.sender}",rcpt="{args.rcpt}",state="{imap_state}",error_string="{a[args.rcpt].error_string}",folder="{a[args.rcpt].current_folder}"}} {imap_time.counter}')
        if args.promextra:
            print('# HELP smtp_gee_stage exports metrics about different stages')
            print('# TYPE smtp_gee_stage gauge')
            print(f'smtp_gee_stage{{protocol="SMTP",stage="transmission",from="{args.sender}",rcpt="{args.rcpt}",state="{imap_state}"}} {a[args.rcpt].smtp_transfer_time}')
            for s in imap_time.stage_counter:
                print(f'smtp_gee_stage{{protocol="IMAP",stage="{s}",from="{args.sender}",rcpt="{args.rcpt}",state="{imap_state}"}} {imap_time.stage_counter[s]}')
            for s in smtp_time.stage_counter:
                print(f'smtp_gee_stage{{protocol="SMTP",stage="{s}",from="{args.sender}",rcpt="{args.rcpt}",state="{smtp_state}"}} {smtp_time.stage_counter[s]}')
            print('# HELP smtp_gee_messages exports metrics about the number of messages checked')
            print('# TYPE smtp_gee_messages gauge')
            for f in a[args.rcpt].imap_messages:
                print(f'smtp_gee_messages{{protocol="IMAP",folder="{f}",from="{args.sender}",rcpt="{args.rcpt}",state="{imap_state}"}} {a[args.rcpt].imap_messages[f]}')
        sys.exit(0)

    elif not args.nagios:

        # Default output
        print("SMTP, (%s) time to send the mail: %.3f sec." % (args.sender, smtp_time.counter, ))
        print("IMAP, (%s) time until the mail appeared in the destination mailbox/%s: %.3f sec." % (args.rcpt, a[args.rcpt].current_folder, imap_time.counter, ))

    else:

        # Nagios output
        # this could be beautified...
        # ...or removed.. Use Prometheus, all the cool kids are doing it!

        nagios_code = ('OK', 'WARNING', 'CRITICAL', 'UNKNOWN' )

        if   ((smtp_time.counter >= args.smtp_crit) or (imap_time.counter >= args.imap_crit)):
            returncode = 2
        elif ((smtp_time.counter >= args.smtp_warn) or (imap_time.counter >= args.imap_warn) or (a[args.rcpt].current_folder != 'INBOX')):
            returncode = 1
        else:
            returncode = 0

        if not smtp_result: # if it failed
            returncode = int(args.except_return)
            error_string = a[args.sender].error_string
            nagios_template="%s: (%s->%s) SMTP failed in %.3f sec, NOT received %s in %.3f sec (%s)|smtp=%.3f;%.3f;%.3f imap=%.3f;%.3f;%.3f"
        elif not imap_result: # if it failed
            error_string = a[args.rcpt].error_string
            returncode = int(args.except_return)
            nagios_template="%s: (%s->%s) sent in %.3f sec, IMAP failed, NOT received %s in %.3f sec (%s)|smtp=%.3f;%.3f;%.3f imap=%.3f;%.3f;%.3f"
        else:
            error_string=""
            nagios_template="%s: (%s->%s) sent in %.3f sec, received in %s in %.3f sec%s|smtp=%.3f;%.3f;%.3f imap=%.3f;%.3f;%.3f"

        print(nagios_template % (
            nagios_code[returncode],
            args.sender,
            args.rcpt,
            smtp_time.counter,
            a[args.rcpt].current_folder,
            imap_time.counter,
            error_string,
            smtp_time.counter,
            args.smtp_warn,
            args.smtp_crit,
            imap_time.counter,
            args.imap_warn,
            args.imap_crit,
        ))

        sys.exit(returncode)

## vim:ts=4:sw=4:sts=4:ai:sta:et