minitor/minitor/main.py

import logging
import subprocess
import sys
from argparse import ArgumentParser
from datetime import datetime
from itertools import chain
from subprocess import CalledProcessError
from subprocess import check_output
from time import sleep

import yamlenv
from prometheus_client import Counter
from prometheus_client import Gauge
from prometheus_client import start_http_server


DEFAULT_METRICS_PORT = 8080
logging.basicConfig(
    level=logging.ERROR,
    format='%(asctime)s %(levelname)s %(name)s %(message)s'
)
logging.getLogger(__name__).addHandler(logging.NullHandler())


def read_yaml(path):
    """Loads config from a YAML file with env interpolation"""
    with open(path, 'r') as yaml:
        contents = yaml.read()
        return yamlenv.load(contents)


def validate_monitor_settings(settings):
    """Validates that settings for a Monitor are valid

    Note: Cannot yet validate the Alerts exist from within this class.
    That will be done by Minitor later
    """
    name = settings.get('name')
    if not name:
        raise InvalidMonitorException('Invalid name for monitor')
    if not settings.get('command'):
        raise InvalidMonitorException(
            'Invalid command for monitor {}'.format(name)
        )

    type_assertions = (
        ('check_interval', int),
        ('alert_after', int),
        ('alert_every', int),
    )

    for key, val_type in type_assertions:
        val = settings.get(key)
        if not isinstance(val, val_type):
            raise InvalidMonitorException(
                'Invalid type on {}: {}. Expected {} and found {}'.format(
                    name, key, val_type.__name__, type(val).__name__
                )
            )

    non_zero = (
        'check_interval',
        'alert_after',
    )

    for key in non_zero:
        if settings.get(key) == 0:
            raise InvalidMonitorException(
                'Invalid value for {}: {}. Value cannot be 0'.format(
                    name, key
                )
            )


def maybe_decode(bstr, encoding='utf-8'):
    try:
        return bstr.decode(encoding)
    except TypeError:
        return bstr


def call_output(*popenargs, **kwargs):
    """Similar to check_output, but instead returns output and exception"""
    # So we can capture complete output, redirect sderr to stdout
    kwargs.setdefault('stderr', subprocess.STDOUT)
    output, ex = None, None
    try:
        output = check_output(*popenargs, **kwargs)
    except CalledProcessError as e:
        output, ex = e.output, e

    output = output.rstrip(b'\n')
    return output, ex


class InvalidAlertException(Exception):
    pass


class InvalidMonitorException(Exception):
    pass


class MinitorAlert(Exception):
    def __init__(self, message, monitor):
        super().__init__(message)
        self.monitor = monitor


class Monitor(object):
    """Primary configuration item for Minitor"""

    def __init__(self, config, counter=None, logger=None):
        """Accepts a dictionary of configuration items to override defaults"""
        settings = {
            'alerts': ['log'],
            'check_interval': 30,
            'alert_after': 4,
            'alert_every': -1,
        }
        settings.update(config)
        validate_monitor_settings(settings)

        self.name = settings['name']
        self.command = settings['command']
        self.alert_down = settings.get('alert_down', [])
        if not self.alert_down:
            self.alert_down = settings.get('alerts', [])
        self.alert_up = settings.get('alert_up', [])
        self.check_interval = settings.get('check_interval')
        self.alert_after = settings.get('alert_after')
        self.alert_every = settings.get('alert_every')

        self.alert_count = 0
        self.last_check = None
        self.last_output = None
        self.last_success = None
        self.total_failure_count = 0

        self._counter = counter
        if logger is None:
            self._logger = logging.getLogger(
                '{}({})'.format(self.__class__.__name__, self.name)
            )
        else:
            self._logger = logger.getChild(
                '{}({})'.format(self.__class__.__name__, self.name)
            )

    def _count_check(self, is_success=True, is_alert=False):
        if self._counter is not None:
            self._counter.labels(
                monitor=self.name,
                status=('success' if is_success else 'failure'),
                is_alert=is_alert,
            ).inc()

    def should_check(self):
        """Determines if this Monitor should run it's check command"""
        if not self.last_check:
            return True
        since_last_check = (datetime.now() - self.last_check).total_seconds()
        return since_last_check >= self.check_interval

    def check(self):
        """Returns None if skipped, False if failed, and True if successful

        Will raise an exception if should alert
        """
        if not self.should_check():
            return None

        output, ex = call_output(
            self.command,
            shell=isinstance(self.command, str),
        )
        output = maybe_decode(output)
        self._logger.debug(output)
        self.last_check = datetime.now()
        self.last_output = output

        is_success = None
        try:
            if ex is None:
                is_success = True
                self.success()
            else:
                is_success = False
                self.failure()
        except MinitorAlert:
            self._count_check(is_success=is_success, is_alert=True)
            raise

        self._count_check(is_success=is_success)
        return is_success

    def success(self):
        """Handles success tasks"""
        back_up = None
        if not self.is_up():
            back_up = MinitorAlert(
                '{} check is up again!'.format(self.name),
                self,
            )
        self.total_failure_count = 0
        self.alert_count = 0
        self.last_success = datetime.now()
        if back_up:
            raise back_up

    def failure(self):
        """Handles failure tasks and possibly raises MinitorAlert"""
        self.total_failure_count += 1
        # Ensure we've hit the  minimum number of failures to alert
        if self.total_failure_count < self.alert_after:
            return

        failure_count = (self.total_failure_count - self.alert_after)
        if self.alert_every > 0:
            # Otherwise, we should check against our alert_every
            should_alert = (failure_count % self.alert_every) == 0
        elif self.alert_every == 0:
            # Only alert on the first failure
            should_alert = failure_count == 1
        else:
            should_alert = (failure_count >= (2 ** self.alert_count) - 1)

        if should_alert:
            self.alert_count += 1
            raise MinitorAlert(
                '{} check has failed {} times'.format(
                    self.name, self.total_failure_count
                ),
                self
            )

    def is_up(self):
        """Indicates if the monitor is already alerting failures"""
        return self.alert_count == 0


class Alert(object):
    def __init__(self, name, config, counter=None, logger=None):
        """An alert must be named and have a config dict"""
        self.name = name
        self.command = config.get('command')
        if not self.command:
            raise InvalidAlertException('Invalid alert {}'.format(self.name))

        self._counter = counter
        if logger is None:
            self._logger = logging.getLogger(
                '{}({})'.format(self.__class__.__name__, self.name)
            )
        else:
            self._logger = logger.getChild(
                '{}({})'.format(self.__class__.__name__, self.name)
            )

    def _count_alert(self, monitor):
        """Increments the alert counter"""
        if self._counter is not None:
            self._counter.labels(
                alert=self.name,
                monitor=monitor,
            ).inc()

    def _formated_command(self, **kwargs):
        """Formats command array or string with kwargs from Monitor"""
        if isinstance(self.command, str):
            return self.command.format(**kwargs)
        args = []
        for arg in self.command:
            args.append(arg.format(**kwargs))
        return args

    def _format_datetime(self, dt):
        """Formats a datetime for an alert"""
        if dt is None:
            return 'Never'
        return dt.isoformat()

    def alert(self, message, monitor):
        """Calls the alert command for the provided monitor"""
        self._count_alert(monitor.name)
        output, ex = call_output(
            self._formated_command(
                alert_count=monitor.alert_count,
                alert_message=message,
                failure_count=monitor.total_failure_count,
                last_output=monitor.last_output,
                last_success=self._format_datetime(monitor.last_success),
                monitor_name=monitor.name,
            ),
            shell=isinstance(self.command, str),
        )
        self._logger.error(maybe_decode(output))
        if ex is not None:
            raise ex


class Minitor(object):
    monitors = None
    alerts = None
    state = None
    check_interval = None

    def __init__(self):
        self._logger = logging.getLogger(self.__class__.__name__)
        self._alert_counter = None
        self._monitor_counter = None
        self._monitor_status_gauge = None

    def _parse_args(self, args=None):
        """Parses command line arguments and returns them"""
        parser = ArgumentParser(description='Minimal monitoring')
        parser.add_argument(
            '--config', '-c',
            dest='config_path',
            default='config.yml',
            help='Path to the config YAML file to use',
        )
        parser.add_argument(
            '--metrics', '-m',
            dest='metrics',
            action='store_true',
            help='Start webserver with metrics',
        )
        parser.add_argument(
            '--metrics-port', '-p',
            dest='metrics_port',
            type=int,
            default=DEFAULT_METRICS_PORT,
            help='Port to use when serving metrics',
        )
        parser.add_argument(
            '--verbose', '-v',
            action='count',
            help=('Adjust log verbosity by increasing arg count. Default log',
                  'level is ERROR. Level increases with each `v`'),
        )
        return parser.parse_args(args)

    def _setup(self, config_path):
        """Load all setup from YAML file at provided path"""
        config = read_yaml(config_path)
        self.check_interval = config.get('check_interval', 30)
        self.monitors = [
            Monitor(
                mon,
                counter=self._monitor_counter,
                logger=self._logger,
            )
            for mon in config.get('monitors', [])
        ]
        # Add default alert for logging
        self.alerts = {
            'log': Alert(
                'log',
                {'command': ['echo', '{alert_message}!']},
                counter=self._alert_counter,
                logger=self._logger,
            )
        }
        self.alerts.update({
            alert_name: Alert(
                alert_name,
                alert,
                counter=self._alert_counter,
                logger=self._logger,
            )
            for alert_name, alert in config.get('alerts', {}).items()
        })

    def _validate_monitors(self):
        """Validates monitors are valid against other config values"""
        for monitor in self.monitors:
            # Validate that the interval is valid
            if monitor.check_interval < self.check_interval:
                raise InvalidMonitorException(
                    'Monitor {} check interval is lower global value {}'.format(
                        monitor.name, self.check_interval
                    )
                )
            # Validate that the the alerts for the monitor exist
            for alert in chain(monitor.alert_down, monitor.alert_up):
                if alert not in self.alerts:
                    raise InvalidMonitorException(
                        'Monitor {} contains an unknown alert: {}'.format(
                            monitor.name, alert
                        )
                    )

    def _init_metrics(self):
        self._alert_counter = Counter(
            'minitor_alert_total',
            'Number of Minitor alerts',
            ['alert', 'monitor'],
        )
        self._monitor_counter = Counter(
            'minitor_check_total',
            'Number of Minitor checks',
            ['monitor', 'status', 'is_alert'],
        )
        self._monitor_status_gauge = Gauge(
            'minitor_monitor_up_count',
            'Currently responsive monitors',
            ['monitor'],
        )

    def _loop(self):
        while True:
            self._check()
            sleep(self.check_interval)

    def _check(self):
        """The main run loop"""
        for monitor in self.monitors:
            try:
                result = monitor.check()
                if result is not None:
                    self._logger.info(
                        '%s: %s',
                        monitor.name,
                        'SUCCESS' if result else 'FAILURE'
                    )
            except MinitorAlert as minitor_alert:
                self._logger.warning(minitor_alert)
                self._handle_minitor_alert(minitor_alert)

            # Track the status of the Monitor
            if self._monitor_status_gauge:
                self._monitor_status_gauge.labels(
                    monitor=monitor.name,
                ).set(int(monitor.is_up()))

    def _handle_minitor_alert(self, minitor_alert):
        """Issues all alerts for a provided monitor"""
        monitor = minitor_alert.monitor
        alerts = monitor.alert_up if monitor.is_up() else monitor.alert_down
        for alert in alerts:
            self.alerts[alert].alert(str(minitor_alert), monitor)

    def _set_log_level(self, verbose):
        """Sets the log level for the class using the provided verbose count"""
        if verbose == 1:
            self._logger.setLevel(logging.WARNING)
        elif verbose == 2:
            self._logger.setLevel(logging.INFO)
        elif verbose >= 3:
            self._logger.setLevel(logging.DEBUG)

    def run(self, args=None):
        """Runs Minitor in a loop"""
        args = self._parse_args(args)

        if args.verbose:
            self._set_log_level(args.verbose)

        if args.metrics:
            self._init_metrics()
            start_http_server(args.metrics_port)

        self._setup(args.config_path)
        self._validate_monitors()

        self._loop()


def main(args=None):
    try:
        Minitor().run(args)
    except KeyboardInterrupt:
        pass
    return 0


if __name__ == '__main__':
    sys.exit(main())