Skip to content
Snippets Groups Projects
cli.py 6.95 KiB
"""
This script queries Inventory Provider for changes
and configures Sensu with the snmp polling checks
required by BRIAN.

.. code-block:: console

    % python brian_polling_manager/cli.py --help
    Usage: cli.py [OPTIONS]

      Update BRIAN snmp checks based on Inventory Provider data.

    Options:
      --config TEXT         configuration filename
      --force / --no-force  update even if inventory hasn't been updated
      --help                Show this message and exit.


The required configuration file must be
formatted according to the following schema:

.. asjson::
    brian_polling_manager.cli.CONFIG_SCHEMA

"""
import json
import logging
import os
from typing import Union

import click
import jsonschema

from brian_polling_manager import inventory, interfaces

logger = logging.getLogger(__name__)

_DEFAULT_CONFIG = {
    'inventory': [
        'http://test-inventory-provider01.geant.org:8080',
        'http://test-inventory-provider02.geant.org:8080'
    ],
    'sensu': {
        'api-base': [
            'https://test-poller-sensu-agent01.geant.org:8080',
            'https://test-poller-sensu-agent02.geant.org:8080',
            'https://test-poller-sensu-agent03.geant.org:8080'
        ],
        'token': '696a815c-607e-4090-81de-58988c83033e',
        'interface-check': {
            'script': '/var/lib/sensu/bin/counter2influx.sh',
            'measurement': 'counters',
            'interval': 300,
            'subscriptions': ['interfacecounters'],
            'output_metric_handlers': ['influx-db-handler'],
            'namespace': 'default',
            'round_robin': True,
            'command': '{script} {measurement} {community} {hostname} {interface} {ifIndex}',
        }
    },
    'statedir': '/tmp/'
}

CONFIG_SCHEMA = {
    '$schema': 'http://json-schema.org/draft-07/schema#',
    'definitions': {
        'influx-check': {
            'type': 'object',
            'properties': {
                'script': {'type': 'string'},
                'measurement': {'type': 'string'},
                'interval': {'type': 'integer'},
                'subscriptions': {
                    'type': 'array',
                    'items': {'type': 'string'}
                },
                'output_metric_handlers': {
                    'type': 'array',
                    'items': {'type': 'string'}
                },
                'namespace': {'type': 'string'},
                'round_robin': {'type': 'boolean'},
                'command': {'type': 'string'},
            },
            'required': ['script', 'measurement', 'interval',
                         'subscriptions', 'output_metric_handlers',
                         'namespace', 'round_robin', 'command'],
            'additionalProperties': False
        },
        'sensu': {
            'type': 'object',
            'properties': {
                'api-base': {
                    'type': 'array',
                    'items': {'type': 'string'},
                    'minItems': 1
                },
                'token': {'type': 'string'},
                'interface-check': {'$ref': '#/definitions/influx-check'}
            },
            'required': ['api-base', 'token', 'interface-check'],
            'additionalProperties': False
        }
    },
    'type': 'object',
    'properties': {
        'inventory': {
            'type': 'array',
            'items': {'type': 'string'},
            'minItems': 1
        },
        'sensu': {'$ref': '#/definitions/sensu'},
        'statedir': {'type': 'string'}
    },
    'required': ['inventory', 'sensu', 'statedir'],
    'additionalProperties': False
}


class State(object):

    INTERFACES = 'interfaces.json'
    STATE = 'state.json'

    STATE_SCHEMA = {
        '$schema': 'http://json-schema.org/draft-07/schema#',
        'type': 'object',
        'properties': {
            'last': {'type': 'number'}
        },
        'required': ['last'],
        'additionalProperties': False
    }

    def __init__(self, state_dir: str):
        assert os.path.isdir(state_dir)
        self.filenames = {
            'state': os.path.join(state_dir, State.STATE),
            'cache': os.path.join(state_dir, State.INTERFACES)
        }

    @staticmethod
    def _load_json(filename, schema):
        try:
            with open(filename) as f:
                state = json.loads(f.read())
                jsonschema.validate(state, schema)
                return state
        except (json.JSONDecodeError, jsonschema.ValidationError, OSError):
            logger.exception(
                f'unable to open state file {filename}')
            return None

    @property
    def last(self) -> int:
        state = State._load_json(self.filenames['state'], State.STATE_SCHEMA)
        return state['last'] if state else -1

    @last.setter
    def last(self, new_last: Union[float, None]):
        if not new_last or new_last < 0:
            os.unlink(self.filenames['state'])
        else:
            state = {'last': new_last}
            with open(self.filenames['state'], 'w') as f:
                f.write(json.dumps(state))

    @property
    def interfaces(self) -> list:
        return State._load_json(self.filenames['cache'], inventory.INVENTORY_INTERFACES_SCHEMA)

    @interfaces.setter
    def interfaces(self, new_interfaces):
        try:
            jsonschema.validate(new_interfaces, inventory.INVENTORY_INTERFACES_SCHEMA)
        except jsonschema.ValidationError:
            logger.exception('invalid interface state data')
            return

        with open(self.filenames['cache'], 'w') as f:
            f.write(json.dumps(new_interfaces))


def _validate_config(ctx, param, value):
    """
    loads, validates and returns configuration parameters

    :param ctx:
    :param param:
    :param value: filename (string)
    :return: a dict containing configuration parameters
    """
    if value is None:
        config = _DEFAULT_CONFIG
    else:
        try:
            with open(value) as f:
                config = json.loads(f.read())
        except (json.JSONDecodeError, jsonschema.ValidationError) as e:
            raise click.BadParameter(str(e))

    try:
        jsonschema.validate(config, CONFIG_SCHEMA)
    except jsonschema.ValidationError as e:
        raise click.BadParameter(str(e))

    return config


@click.command()
@click.option(
    '--config',
    default=None,
    type=click.STRING,
    callback=_validate_config,
    help='configuration filename')
@click.option(
    '--force/--no-force',
    default=False,
    help="update even if inventory hasn't been updated")
def main(config, force):
    """
    Update BRIAN snmp checks based on Inventory Provider data.
    """

    state = State(config['statedir'])
    last = inventory.last_update_timestamp(config['inventory'])
    if force or not last or last != state.last:
        state.last = last
        state.interfaces = inventory.load_interfaces(config['inventory'])

    interfaces.refresh(config['sensu'], state)


if __name__ == '__main__':
    logging.basicConfig(level=logging.DEBUG)
    main()