#! /usr/bin/env python3
"""
 discovery-wrapper A small tool which wraps around discovery and tries to
                guide the discovery process with a more modern approach with a
                Queue and workers.

 Based on the original version of poller-wrapper.py by Job Snijders

 Author:        Neil Lathwood <neil@librenms.org>
                Orsiris de Jong <contact@netpower.fr>
 Date:          Oct 2019

 Usage:         This program accepts one command line argument: the number of threads
                that should run simultaneously. If no argument is given it will assume
                a default of 1 thread.

 Ubuntu Linux:  apt-get install python-mysqldb
 FreeBSD:       cd /usr/ports/*/py-MySQLdb && make install clean
 RHEL 7:        yum install MySQL-python
 RHEL 8:        dnf install mariadb-connector-c-devel gcc && python -m pip install mysqlclient

 Tested on:     Python 3.6.8 / PHP 7.2.11 / CentOS 8

 License:       This program is free software: you can redistribute it and/or modify it
                under the terms of the GNU General Public License as published by the
                Free Software Foundation, either version 3 of the License, or (at your
                option) any later version.

                This program is distributed in the hope that it will be useful, but
                WITHOUT ANY WARRANTY; without even the implied warranty of
                MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General
                Public License for more details.

                You should have received a copy of the GNU General Public License along
                with this program. If not, see http://www.gnu.org/licenses/.

                LICENSE.txt contains a copy of the full GPLv3 licensing conditions.
"""

import LibreNMS.library as LNMS

try:

    import json
    import os
    import queue
    import subprocess
    import sys
    import threading
    import time
    from optparse import OptionParser

except ImportError as exc:
    print('ERROR: missing one or more of the following python modules:')
    print('threading, queue, sys, subprocess, time, os, json')
    print('ERROR: %s' % exc)
    sys.exit(2)

APP_NAME = "discovery_wrapper"
LOG_FILE = "logs/" + APP_NAME + ".log"
_DEBUG = False
distdisco = False
real_duration = 0
discovered_devices = 0

# (c) 2015, GPLv3, Daniel Preussker <f0o@devilcode.org> <<<EOC0
def memc_alive():
    try:
        global memc
        key = str(uuid.uuid4())
        memc.set('discovery.ping.' + key, key, 60)
        if memc.get('discovery.ping.' + key) == key:
            memc.delete('discovery.ping.' + key)
            return True
        else:
            return False
    except:
        return False


def memc_touch(key, time):
    try:
        global memc
        val = memc.get(key)
        memc.set(key, val, time)
    except:
        pass


# EOC0

"""
    A seperate queue and a single worker for printing information to the screen prevents
    the good old joke:

        Some people, when confronted with a problem, think,
        "I know, I'll use threads," and then they two they hav erpoblesms.
"""


def printworker():
    nodeso = 0
    while True:
        # (c) 2015, GPLv3, Daniel Preussker <f0o@devilcode.org> <<<EOC4
        global IsNode
        global distdisco
        if distdisco:
            if not IsNode:
                memc_touch('discovery.master', 30)
                nodes = memc.get('discovery.nodes')
                if nodes is None and not memc_alive():
                    print("WARNING: Lost Memcached. Taking over all devices. Nodes will quit shortly.")
                    distdisco = False
                    nodes = nodeso
                if nodes is not nodeso:
                    print("INFO: %s Node(s) Total" % (nodes))
                    nodeso = nodes
            else:
                memc_touch('discovery.nodes', 30)
            try:
                worker_id, device_id, elapsed_time = print_queue.get(False)
            except:
                pass
                try:
                    time.sleep(1)
                except:
                    pass
                continue
        else:
            worker_id, device_id, elapsed_time = print_queue.get()
        # EOC4
        global real_duration
        global per_device_duration
        global discovered_devices
        real_duration += elapsed_time
        per_device_duration[device_id] = elapsed_time
        discovered_devices += 1
        if elapsed_time < 300:
            print("INFO: worker %s finished device %s in %s seconds" % (worker_id, device_id, elapsed_time))
        else:
            print("WARNING: worker %s finished device %s in %s seconds" % (worker_id, device_id, elapsed_time))
        print_queue.task_done()


"""
    This class will fork off single instances of the discovery.php process, record
    how long it takes, and push the resulting reports to the printer queue
"""


def poll_worker():
    while True:
        device_id = poll_queue.get()
        # (c) 2015, GPLv3, Daniel Preussker <f0o@devilcode.org> <<<EOC5
        if not distdisco or memc.get('discovery.device.' + str(device_id)) is None:
            if distdisco:
                result = memc.add('discovery.device.' + str(device_id), config['distributed_poller_name'], 300)
                if not result:
                    print("This device (%s) appears to be being discovered by another discovery node" % (device_id))
                    poll_queue.task_done()
                    continue
                if not memc_alive() and IsNode:
                    print("Lost Memcached, Not discovering Device %s as Node. Master will discover it." % device_id)
                    poll_queue.task_done()
                    continue
            # EOC5
            try:
                start_time = time.time()

                output = "-d >> %s/discover_device_%s.log" % (log_dir, device_id) if debug else ">> /dev/null"
                command = "/usr/bin/env php %s -h %s %s 2>&1" % (discovery_path, device_id, output)
                # TODO: Replace with command_runner
                subprocess.check_call(command, shell=True)

                elapsed_time = int(time.time() - start_time)
                print_queue.put([threading.current_thread().name, device_id, elapsed_time])
            except (KeyboardInterrupt, SystemExit):
                raise
            except:
                pass
        poll_queue.task_done()


if __name__ == '__main__':
    logger = LNMS.logger_get_logger(LOG_FILE, debug=_DEBUG)

    install_dir = os.path.dirname(os.path.realpath(__file__))
    LNMS.check_for_file(install_dir + '/.env')
    config = json.loads(LNMS.get_config_data(install_dir))

    discovery_path = config['install_dir'] + '/discovery.php'
    log_dir = config['log_dir']

    # (c) 2015, GPLv3, Daniel Preussker <f0o@devilcode.org> <<<EOC1
    if 'distributed_poller_group' in config:
        discovery_group = str(config['distributed_poller_group'])
    else:
        discovery_group = False

    if ('distributed_poller' in config and
        'distributed_poller_memcached_host' in config and
        'distributed_poller_memcached_port' in config and
        config['distributed_poller']):
        try:
            import memcache
            import uuid

            memc = memcache.Client([config['distributed_poller_memcached_host'] + ':' +
                                    str(config['distributed_poller_memcached_port'])])
            if str(memc.get("discovery.master")) == config['distributed_poller_name']:
                print("This system is already joined as the discovery master.")
                sys.exit(2)
            if memc_alive():
                if memc.get("discovery.master") is None:
                    print("Registered as Master")
                    memc.set("discovery.master", config['distributed_poller_name'], 30)
                    memc.set("discovery.nodes", 0, 3600)
                    IsNode = False
                else:
                    print("Registered as Node joining Master %s" % memc.get("discovery.master"))
                    IsNode = True
                    memc.incr("discovery.nodes")
                distdisco = True
            else:
                print("Could not connect to memcached, disabling distributed discovery.")
                distdisco = False
                IsNode = False
        except SystemExit:
            raise
        except ImportError:
            print("ERROR: missing memcache python module:")
            print("On deb systems: apt-get install python3-memcache")
            print("On other systems: pip3 install python-memcached")
            print("Disabling distributed discovery.")
            distdisco = False
    else:
        distdisco = False
    # EOC1

    s_time = time.time()
    real_duration = 0
    per_device_duration = {}
    discovered_devices = 0

    """
        Take the amount of threads we want to run in parallel from the commandline
        if None are given or the argument was garbage, fall back to default of 1
    """
    usage = "usage: %prog [options] <workers> (Default: 1 Do not set too high)"
    description = "Spawn multiple discovery.php processes in parallel."
    parser = OptionParser(usage=usage, description=description)
    parser.add_option('-d', '--debug', action='store_true', default=False,
                      help="Enable debug output. WARNING: Leaving this enabled will consume a lot of disk space.")
    (options, args) = parser.parse_args()

    debug = options.debug
    try:
        amount_of_workers = int(args[0])
    except (IndexError, ValueError):
        amount_of_workers = 1

    devices_list = []

    """
        This query specificly orders the results depending on the last_discovered_timetaken variable
        Because this way, we put the devices likely to be slow, in the top of the queue
        thus greatening our chances of completing _all_ the work in exactly the time it takes to
        discover the slowest device! cool stuff he
    """
    # (c) 2015, GPLv3, Daniel Preussker <f0o@devilcode.org> <<<EOC2
    if discovery_group is not False:
        query = "select device_id from devices where poller_group IN(" + discovery_group + ") and disabled = 0 order by last_polled_timetaken desc"
    else:
        query = "select device_id from devices where disabled = 0 order by last_polled_timetaken desc"
    # EOC2

    db = LNMS.db_open(config['db_socket'], config['db_host'], int(config['db_port']), config['db_user'], config['db_pass'], config['db_name'])
    cursor = db.cursor()
    cursor.execute(query)
    devices = cursor.fetchall()
    for row in devices:
        devices_list.append(int(row[0]))
    # (c) 2015, GPLv3, Daniel Preussker <f0o@devilcode.org> <<<EOC3
    if distdisco and not IsNode:
        query = "select max(device_id),min(device_id) from devices"
        cursor.execute(query)
        devices = cursor.fetchall()
        maxlocks = devices[0][0] or 0
        minlocks = devices[0][1] or 0
    # EOC3
    db.close()

    poll_queue = queue.Queue()
    print_queue = queue.Queue()

    print("INFO: starting the discovery at %s with %s threads, slowest devices first" % (
        time.strftime("%Y-%m-%d %H:%M:%S"),
        amount_of_workers))

    for device_id in devices_list:
        poll_queue.put(device_id)

    for i in range(amount_of_workers):
        t = threading.Thread(target=poll_worker)
        t.setDaemon(True)
        t.start()

    p = threading.Thread(target=printworker)
    p.setDaemon(True)
    p.start()

    try:
        poll_queue.join()
        print_queue.join()
    except (KeyboardInterrupt, SystemExit):
        raise

    total_time = int(time.time() - s_time)

    print("INFO: discovery-wrapper polled %s devices in %s seconds with %s workers" % (
        discovered_devices, total_time, amount_of_workers))

    # (c) 2015, GPLv3, Daniel Preussker <f0o@devilcode.org> <<<EOC6
    if distdisco or memc_alive():
        master = memc.get("discovery.master")
        if master == config['distributed_poller_name'] and not IsNode:
            print("Wait for all discovery-nodes to finish")
            nodes = memc.get("discovery.nodes")
            while nodes is not None and nodes > 0:
                try:
                    time.sleep(1)
                    nodes = memc.get("discovery.nodes")
                except:
                    pass
            print("Clearing Locks")
            x = minlocks
            while x <= maxlocks:
                memc.delete('discovery.device.' + str(x))
                x = x + 1
            print("%s Locks Cleared" % x)
            print("Clearing Nodes")
            memc.delete("discovery.master")
            memc.delete("discovery.nodes")
        else:
            memc.decr("discovery.nodes")
        print("Finished %s." % time.time())
    # EOC6

    show_stopper = False

    if total_time > 21600:
        print("WARNING: the process took more than 6 hours to finish, you need faster hardware or more threads")
        print("INFO: in sequential style discovery the elapsed time would have been: %s seconds" % real_duration)
        for device in per_device_duration:
            if per_device_duration[device] > 3600:
                print("WARNING: device %s is taking too long: %s seconds" % (device, per_device_duration[device]))
                show_stopper = True
        if show_stopper:
            print("ERROR: Some devices are taking more than 3600 seconds, the script cannot recommend you what to do.")
        else:
            recommend = int(total_time / 300.0 * amount_of_workers + 1)
            print(
                "WARNING: Consider setting a minimum of %d threads. (This does not constitute professional advice!)" % recommend)

        sys.exit(2)