massive-passive-client

#!/usr/bin/python
# -*- coding: utf-8 -*-

__author__ = 'Robin Wittler'
__contact__ = 'r.wittler@mysportgroup.de'
__copyright__ = '(c) 2012 by mysportgroup GmbH'
__license__ = 'GPL3+'
__version__ = '0.3.9'

import os
import sys
import signal
import logging
from time import sleep
from Queue import Queue
from threading import Event
from threading import Thread
from pyinotify import IN_DELETE
from pyinotify import IN_CREATE
from pyinotify import IN_MOVED_TO
from mplib.config import ConfigDir
from mplib.daemon import daemonize
from mplib.config import ConfigFile
from pyinotify import IN_MOVED_FROM
from pyinotify import IN_CLOSE_WRITE
from setproctitle import setproctitle
from mplib.inotify import WatchManager
from mplib.getopts import get_gpl3_text
from mplib.log import BASE_FORMAT_STDOUT
from mplib.inotify import ThreadedNotifier
from apscheduler.threadpool import ThreadPool
from mplib.log import set_logfile_permissions
from mplib.inotify import ProcessConfigEvents
from mplib.getopts.client import client_getopt
from logging.handlers import WatchedFileHandler
from mplib.threads.client import SendNativeWorker
from mplib.scheduler import MassivePassiveScheduler
from mplib.net.client import PassiveCheckSubmitClient
from mplib.getopts.client import get_client_description


class SendNscaModeMain(object):
    def __init__(self, name, options, check_submit_client):
        super(SendNscaModeMain, self).__init__()
        self.name = name
        self.logger = logging.getLogger(
            '%s.%s' %(
                self.name,
                self.__class__.__name__
            )
        )
        self.options = options
        self.stop_event = Event()
        self.send_queue = Queue()
        self.check_submit_client = check_submit_client
        self.active_threads = list()

    def run(self):
        self.set_proc_title()
        self.logger.info('Startup!')
        message = '%s\n' %(self.read_message_from_stdin().rstrip())
        for ip in self.options.server:
            try:
                ip, port = ip.split(':')
            except ValueError:
                port = 5678

            check_submit_client = self.check_submit_client(
                message,
                ip,
                self.options.ssl_key,
                self.options.ssl_cert,
                self.options.ssl_ca_cert,
                port=int(port)
            )
            check_submit_client_thread = Thread(target=check_submit_client.start)
            check_submit_client_thread.start()
            self.active_threads.append(check_submit_client_thread)
        os.kill(os.getpid(), signal.SIGTERM)

    def read_message_from_stdin(self, timeout=3):
        def timeout_handler(signum, sigframe):
            raise RuntimeError('No message could read from stdin after %r seconds.' %(timeout))
        signal.signal(signal.SIGALRM, timeout_handler)
        signal.alarm(timeout)
        message = sys.stdin.read()
        signal.alarm(0)
        return message

    def set_proc_title(self):
        setproctitle(
            '%s %s' %(sys.argv[0], ' '.join(sys.argv[1:]))
        )

    def start(self):
        try:
            self.run()
        except Exception as error:
            self.logger.exception(error)
            self.logger.info('Emitting stop signal ...')
            os.kill(os.getpid(), signal.SIGTERM)

    def shutdown(self, signum, sigframe):
        self.logger.info('Received Signal %s.', signum)
        self.logger.info('Going down now ...')
        self.stop_event.set()
        for thread in self.active_threads:
            self.logger.debug('Joining %r ...', thread.name)
            thread.join()
            self.logger.debug('Thread %r alive? => %r', thread.name, thread.is_alive())
        self.logger.debug('Stop Event set.')
        self.logger.info('Exiting now!')

    def config_reload(self, signum, sigframe):
        pass


class Main(SendNscaModeMain):
    def __init__(self, name, options, send_worker, scheduler):
        self.name = name
        self.logger = logging.getLogger(
            '%s.%s' %(
                self.name,
                self.__class__.__name__
            )
        )
        self.options = options
        self.stop_event = Event()
        self.send_queue = Queue()
        self._send_worker = send_worker
        self.send_worker = self._send_worker(
            self.send_queue,
            self.stop_event,
            self.options.ssl_key,
            self.options.ssl_cert,
            self.options.ssl_ca_cert,
            max_wait=self.options.batch_wait_time,
            batch_mode=self.options.batch_mode,
            max_results=self.options.batch_max_items
        )
        self._scheduler = scheduler
        self.watch_manager = WatchManager()
        callbacks = {
            'IN_MOVED_FROM': self.remove_job_from_event,
            'IN_MOVED_TO': self.add_job_from_event,
            'IN_CLOSE_WRITE': self.add_job_from_event,
            'IN_CREATE': self.add_job_from_event,
            'IN_DELETE': self.remove_job_from_event,
        }
        self.notifier = ThreadedNotifier(
            self.watch_manager,
            default_proc_fun=ProcessConfigEvents(callbacks=callbacks)
        )
        self.watch_manager.add_watch(
            options.confdir,
            IN_CLOSE_WRITE|IN_DELETE|IN_CREATE|IN_MOVED_TO|IN_MOVED_FROM,
            rec=False,
            auto_add=True
        )

    def run(self):
        self.set_proc_title()
        self.drop_rights()
        if self.options.foreground is False:
            self.write_pidfile()
        self.logger.info('Startup!')
        self.start_send_worker()
        self.start_scheduler()
        self.add_jobs_to_scheduler()
        self.notifier.start()
        while not self.stop_event.is_set():
            sleep(0.1)

    def write_pidfile(self):
        pid = os.getpid()
        pidfile_fd = os.open(
            self.options.pidfile,
            os.O_CREAT | os.O_EXCL | os.O_RDWR,
            0640
        )
        pidfile = os.fdopen(pidfile_fd, 'w+')
        pidfile.write('%s\n' %(pid,))
        pidfile.flush()
        os.close(pidfile_fd)

    def drop_rights(self):
        current_uid = os.getuid()
        current_gid = os.getgid()

        if not options.group == current_gid:
            self.logger.debug('Switching group to %r', options.group)
            os.setgid(options.group)
            os.setegid(options.group)

        if not options.user == current_uid:
            self.logger.debug('Switching user to %r', options.user)
            os.setuid(options.user)
            os.seteuid(options.user)

    def start_scheduler(self):
        self.logger.debug('Calling scheduler start ...')
        jobs_counter = len(self.get_check_configs())
        self.scheduler = self._scheduler(
            self.send_queue,
            threadpool=ThreadPool(
                core_threads=(jobs_counter / 10),
                max_threads=jobs_counter,
                keepalive=30
            )
        )
        self.scheduler.start()
        self.logger.debug('Scheduler is running: %r.', self.scheduler.running)

    def add_jobs_to_scheduler(self):
        self.logger.debug('Getting passive check configs ...')
        configs = self.get_check_configs()
        self.scheduler.add_passive_checks(
            configs,
            wait_range_end=self.options.initial_random_wait_range
        )

    def add_job_from_event(self, event):
        self.logger.info('Adding job triggered by event: %r', event)
        self.scheduler.remove_job(event.pathname)
        self.scheduler.add_passive_checks(
            {1: ConfigFile(event.pathname)},
            wait_range_end=self.options.initial_random_wait_range
        )

    def remove_job_from_event(self, event):
        self.logger.info('Removing job triggered by event: %r', event)
        self.scheduler.remove_job(event.pathname)

    def start_send_worker(self):
        self.logger.debug('Calling send_worker thread start ...')
        self.send_worker.start()
        self.logger.debug(
            'send_worker is alive: %r', self.send_worker.is_alive()
        )

    def stop_scheduler(self):
        if self.scheduler.running:
            self.logger.debug('Calling scheduler shutdown ...')
            self.scheduler.shutdown()
            self.logger.debug('Scheduler shutdown done.')
        self.logger.debug('Scheduler is running: %r.', self.scheduler.running)

    def stop_send_worker(self):
        if self.send_worker.is_alive():
            self.logger.debug(
                'Calling %r thread join ...',
                self.send_worker.name
            )

            self.send_worker.join()

            self.logger.debug(
                '%r joined.',
                self.send_worker.name
            )

        self.logger.debug(
            '%r is alive: %r',
            self.send_worker.name,
            self.send_worker.is_alive()
        )

    def get_check_configs(self):
        if not os.access(self.options.confdir, os.R_OK|os.X_OK):
            raise RuntimeError(
                'Can not access %r with mode rx.',
                self.options.confdir
            )
        configs = ConfigDir(self.options.confdir)
        self.logger.debug(
            'Get this from confdir %r: %r', self.options.confdir, configs
        )
        return configs

    def config_reload(self, signum, sigframe):
        self.logger.info(
            'Received Signal %s ... reloading config now ...', signum
        )

        self.scheduler.remove_all_jobs()
        self.add_jobs_to_scheduler()
        self.logger.info('Config reload done.')

    def shutdown(self, signum, sigframe):
        self.logger.info('Received Signal %s.', signum)
        self.logger.info('Going down now ...')
        self.stop_event.set()
        self.logger.debug('Stop Event set.')
        self.logger.debug('Stopping notifier ...')
        self.notifier.stop()
        self.logger.debug('Notifier stopped.')
        self.stop_scheduler()
        self.stop_send_worker()

        if self.options.foreground is False:
            self.logger.debug('Removing pidfile at %r', self.options.pidfile)
            try:
                os.unlink(self.options.pidfile)
            except OSError as error:
                self.logger.debug('Ignoring error while unlinking pidfile.')
                self.logger.exception(error)
        self.logger.info('Exiting now!')

if __name__ == '__main__':
    options, args = client_getopt(
        version='%prog ' + __version__,
        description=get_client_description(),
        epilog=get_gpl3_text()
    )

    if options.silent is False:
        stream = None
    else:
        stream = open(os.devnull, 'w', 1)
    logging.basicConfig(stream=stream, format=BASE_FORMAT_STDOUT)
    root_logger = logging.getLogger('')
    logger = logging.getLogger(sys.argv[0])
    logger.setLevel(options.loglevel)

    if options.foreground is False and options.act_as_sender is False:
        daemonize(cwd='/')
        main = Main(sys.argv[0], options, SendNativeWorker, MassivePassiveScheduler)
        main.drop_rights()
    elif options.foreground is True:
        main = Main(sys.argv[0], options, SendNativeWorker, MassivePassiveScheduler)
    else:
        main = SendNscaModeMain(sys.argv[0], options, PassiveCheckSubmitClient)

    filehandler = WatchedFileHandler(options.logfile)
    filehandler.setLevel(options.loglevel)
    filehandler_formatter = logging.Formatter(BASE_FORMAT_STDOUT)
    filehandler.setFormatter(filehandler_formatter)
    root_logger.addHandler(filehandler)
    root_logger.setLevel(options.loglevel)
    set_logfile_permissions(options.logfile)

    signal.signal(signal.SIGTERM, main.shutdown)
    signal.signal(signal.SIGINT, main.shutdown)
    signal.signal(signal.SIGHUP, main.config_reload)
    main.start()
    sys.exit(0)


# vim: tabstop=8 expandtab shiftwidth=4 softtabstop=4