Source code for mentat.module.detmngr

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
#-------------------------------------------------------------------------------
# This file is part of Mentat system (https://mentat.cesnet.cz/).
#
# Copyright (C) since 2011 CESNET, z.s.p.o (http://www.ces.net/)
# Use of this source is governed by the MIT license, see LICENSE file.
#-------------------------------------------------------------------------------


"""
This Mentat module is a script providing functions for detectors management
for Mentat system database.

This script is implemented using the :py:mod:`pyzenkit.zenscript` framework and
so it provides all of its core features. See the documentation for more in-depth
details.

.. note::

    Still work in progress, use with caution.


Usage examples
--------------

.. code-block:: shell

    # Display help message and exit.
    mentat-detmngr.py --help

    # Run in debug mode (enable output of debugging information to terminal).
    mentat-detmngr.py --debug

    # Run with increased logging level.
    mentat-detmngr.py --log-level debug


Available script commands
-------------------------

``status`` (*default*)
    Detect and display the state of internal whois database contents according
    to the data in given reference whois file.

``update``
    Attempt to update the state of internal whois database contents according
    to the data in given reference whois file.


Custom configuration
--------------------

Custom command line options
^^^^^^^^^^^^^^^^^^^^^^^^^^^

``--detectors-file file-path``
    Path to reference detectors file containing data.

    *Type:* ``string``, *default:* ``None``

``--source``
    Origin of the whois file.

    *Type:* ``string``, *default:* ``detectors-file``
"""


__author__  = "Rajmund Hruška <rajmund.hruska@cesnet.cz>"
__credits__ = "Jan Mach <jan.mach@cesnet.cz>, Pavel Kácha <pavel.kacha@cesnet.cz>, Andrea Kropáčová <andrea.kropacova@cesnet.cz>"


import json
import collections

#
# Custom libraries.
#
import pyzenkit.jsonconf
import mentat.script.fetcher
import mentat.const
import mentat.datatype.internal
from mentat.datatype.sqldb import DetectorModel, detectormodel_from_typeddict

DETECTORS_FILE_GENERIC = 'detectors-file'
DETECTORS_FILE_WARDEN = 'warden'

[docs]class MentatDetmngrScript(mentat.script.fetcher.FetcherScript): """ Implementation of Mentat module (script) providing functions for detectors management for Mentat database. """ # # Class constants. # # List of configuration keys. CONFIG_DETECTORS_FILE = 'detectors_file' CONFIG_DETECTORS_SOURCE = 'source' def __init__(self): """ Initialize detmngr script object. This method overrides the base implementation in :py:func:`pyzenkit.zenscript.ZenScript.__init__` and it aims to even more simplify the script object creation by providing configuration values for parent contructor. """ self.eventservice = None self.sqlservice = None super().__init__( description = 'mentat-detmngr.py - Detector management script for Mentat database', ) def _init_argparser(self, **kwargs): """ Initialize script command line argument parser. This method overrides the base implementation in :py:func:`pyzenkit.zenscript.ZenScript._init_argparser` and it must return valid :py:class:`argparse.ArgumentParser` object. It appends additional command line options custom for this script object. This method is called from the main constructor in :py:func:`pyzenkit.baseapp.BaseApp.__init__` as a part of the **__init__** stage of application`s life cycle. :param kwargs: Various additional parameters passed down from object constructor. :return: Valid argument parser object. :rtype: argparse.ArgumentParser """ argparser = super()._init_argparser(**kwargs) # # Create and populate options group for custom script arguments. # arggroup_script = argparser.add_argument_group('custom script arguments') arggroup_script.add_argument( '--detectors-file', type = str, default = None, help = 'path to reference detectors file containing data' ) arggroup_script.add_argument( '--source', type = str, default = DETECTORS_FILE_GENERIC, help = 'origin of the detectors file' ) return argparser def _init_config(self, cfgs, **kwargs): """ Initialize default script configurations. This method overrides the base implementation in :py:func:`pyzenkit.zenscript.ZenScript._init_config` and it appends additional configurations via ``cfgs`` parameter. This method is called from the main constructor in :py:func:`pyzenkit.baseapp.BaseApp.__init__` as a part of the **__init__** stage of application`s life cycle. :param list cfgs: Additional set of configurations. :param kwargs: Various additional parameters passed down from constructor. :return: Default configuration structure. :rtype: dict """ cfgs = ( (self.CONFIG_DETECTORS_FILE, None), (self.CONFIG_DETECTORS_SOURCE, DETECTORS_FILE_GENERIC) ) + cfgs return super()._init_config(cfgs, **kwargs) #---------------------------------------------------------------------------
[docs] def get_default_command(self): """ Return the name of the default script command. This command will be executed in case it is not explicitly selected either by command line option, or by configuration file directive. :return: Name of the default command. :rtype: str """ return 'status'
[docs] def cbk_command_status(self): """ Implementation of the **status** command (*default*). Detect and display the status of detectors collection. """ result = self._process_detectors(True) return result
[docs] def cbk_command_update(self): """ Implementation of the **update** command. Attempt to update the state of internal detectors database contents according to the data in given reference detectors file. """ result = self._process_detectors(False) return result
#--------------------------------------------------------------------------- def _process_detectors(self, status_only): """ The actual worker method for processing detectors records. :param bool status_only: Do not actually perform any database operations, just report status. :return: Structure containing information about changes. :rtype: dict """ result = {'create': [], 'delete': [], 'update': []} det_db = {} det_file = self.c(self.CONFIG_DETECTORS_FILE) det_file_type, det_file_data_raw = self._load_detectors_file(det_file) self.logger.debug("Raw data: %s", str(det_file_data_raw)) det_file_data = self._process_detectors_data(det_file_data_raw, det_file_type) self.logger.info("Number of detectors in reference detectors file: %d", len(det_file_data.keys())) detectors = self.sqlservice.session.query(DetectorModel).order_by(DetectorModel.name).all() self.sqlservice.session.commit() self.logger.info("Number of detectors in the database: %d", len(detectors)) for detector in detectors: det_db[detector.name] = detector self._detectors_create_missing(det_db, det_file_data, det_file_type, result, status_only) self._detectors_report_extra(det_db, det_file_data, det_file_type, result, status_only) self._detectors_update_existing(det_db, det_file_data, det_file_type, result, status_only) return result def _load_detectors_file(self, detectors_file): """ Load reference detectors file. :param str detectors_file: Name of the reference detectors file. :return: Data content of detectors file. :rtype: dict """ try: with open(detectors_file, 'r') as jsf: json_data = jsf.read() detectors_file_data = json.loads(json_data) except: raise pyzenkit.zenscript.ZenScriptException("Invalid detectors file '{}', expected JSON formated file".format(detectors_file)) detectors_file_type = self.c(self.CONFIG_DETECTORS_SOURCE) self.logger.info("Loaded reference detectors file '%s :: %s'", detectors_file, detectors_file_type) return (detectors_file_type, detectors_file_data) def _process_detectors_data(self, detectors_file_data, detectors_file_type): """ Process reference detectors file data into format more appropriate for searching and comparisons. :param dict whois_file_data: Whois data as loaded by :py:func:`_load_whois_file`. :param str whois_file_type: Type of the whois file (value of ``__whois_type__`` meta attribute). :return: Processed whois file data into format more appropriate for searching. :rtype: dict """ if 'clients' not in detectors_file_data: raise pyzenkit.zenscript.ZenScriptException("Invalid detectors file format, expected 'clients' key.") processed_data = collections.defaultdict(dict) for client in detectors_file_data['clients']: det = mentat.datatype.internal.t_detector_record(client, source=detectors_file_type) self.logger.debug(det) processed_data[det['name']] = det return processed_data def _detectors_create_missing(self, det_db, det_file_data, det_file_type, result, status_only): """ Create missing detector records in the database. :param dict det_db: Detectors loaded from the database. :param dict det_file_data: Detectors loaded from the reference detectors file. :param str det_file_type: Source of the detectors in the reference detectors file. :param dict result: Structure containing processing log, will be appended to script runlog. :param bool status_only: Do not actually perform any database operations, just report status. """ for detector_name in sorted(det_file_data.keys()): # Try finding the detector from the file in the database by the name. if not detector_name in det_db: gkey = '{}::{}'.format(detector_name, det_file_type) result['create'].append(gkey) if status_only: self.logger.warning("'%s' Found new detector.", gkey) continue sqldet = detectormodel_from_typeddict( det_file_data[detector_name], {'description' : 'Detector created automatically by mentat-detmngr.py utility.'} ) self.logger.warning("'%s' Creating new detector.", gkey) self.sqlservice.session.add(sqldet) self.sqlservice.session.commit() def _detectors_report_extra(self, det_db, det_file_data, det_file_type, result, status_only): """ Report extra detectors from database. :param dict det_db: Detectors loaded from the database. :param dict det_file_data: Detectors loaded from the reference detectors file. :param str det_file_type: Source of the detectors in the reference detectors file. :param dict result: Structure containing processing log, will be appended to script runlog. :param bool status_only: Do not actually perform any database operations, just report status. """ for detector_name in sorted(det_db.keys()): det = det_db[detector_name] # For deletion consider only detectors with the same origin (source) as # the loaded detectors file. if det.source == det_file_type and detector_name not in det_file_data: detkey = '{}::{}'.format(det.name, det.source) result['delete'].append(detkey) self.logger.warning("'%s' Detector was not found in the loaded detectors file, consider deletion.", detkey) def _detectors_update_existing(self, det_db, det_file_data, det_file_type, result, status_only): """ Update existing detectors within the database. :param dict det_db: Detectors loaded from the database. :param dict det_file_data: Detectors loaded from the reference detectors file. :param str det_file_type: Source of the detectors in the reference detectors file. :param dict result: Structure containing processing log, will be appended to script runlog. :param bool status_only: Do not actually perform any database operations, just report status. """ for detector_name in sorted(det_db.keys()): if det_db[detector_name].source != det_file_type or detector_name not in det_file_data: continue detector_db = det_db[detector_name] detector_file = det_file_data[detector_name] if self._detectors_differ(detector_db, detector_file): detkey = '{}::{}'.format(detector_db.name, detector_db.source) result['update'].append(detkey) if status_only: self.logger.warning("Detector '%s' has changed.", detkey) continue self.logger.warning("Updating existing detector '%s'.", detkey) detector_db.credibility = detector_file['credibility'] detector_db.registered = detector_file['registered'] self.sqlservice.session.commit() @staticmethod def _detectors_differ(det_db, det_file): """ Check if the given detectors differ. It is assumed that detectors have the same name and source. :param det_db: Instance of :py:class:`mentat.datatype.sqldb.DetectorModel` :param det_file: Instance of :py:class:`mentat.datatype.internal.Detector` :return: True or False :rtype: bool """ return det_db.credibility != det_file['credibility'] or det_db.registered != det_file['registered']