#!/usr/bin/python
# .- coding: utf-8 -.
#
# Artificial munin node that behaves in all the ways you would like
# ordinary nodes _not_ to behave.
#
# Intended use is for designing and debugging munin-server poller to handle 
# such problems.
#
# See the file MIT-LICENSE for licensing information.
#
# Copyright (C) 2011 Karstensen IT
# Written by Lasse Karstensen <lasse.karstensen@gmail.com>, Dec 2011.

import os, sys, time, random
import socket
import threading
import SocketServer
import ConfigParser

VERSION = "muninnode-from-hell v0.1"
modules = {}

class MuninPlugin:
    def __init__(self):
        self.current_load = None
        self.current_locks = None

    def sleep_fetch(self, conf):
        period = None
        if conf.get("mode") == "sleepy" and conf.get("sleepyness"):
            period = float(conf.get("sleepyness"))
        if conf.get("mode") == "exp" and conf.get("lambd"):
            period = random.expovariate(1 / float(conf.get("lambd")))

        if period:
            #print "will sleep %.3f seconds" % period
            time.sleep(period)
            
    def sleep_config(self, conf):
        return self.sleep_fetch(conf)

    def find_load(self):
        # At about a thousand node instances you get this:
        #IOError: [Errno 24] Too many open files: '/proc/loadavg'
        # cache it for a bit..
        if (not self.current_load) or random.randint(0,100) == 1:
            load = open("/proc/loadavg", "r").read()
            load, rest = load.split(" ", 1)
            self.current_load = float(load)
        return self.current_load
            
    def find_locks(self):
        if (not self.current_locks) or random.randint(0,100) == 1:
            fp = open("/proc/locks", "r")
            self.current_locks = len(fp.readlines())
        return self.current_locks 

class load(MuninPlugin):
    def fetch(self, conf):
        self.sleep_fetch(conf)
        return "load.value %.2f" % self.find_load()

    def config(self, conf):
        self.sleep_config(conf)
        return """graph_title Load average
graph_args --base 1000 -l 0
graph_vlabel load
graph_scale no
graph_category system
load.label load
graph_info The load average of the machine describes how many processes are in the run-queue (scheduled to run "immediately").
load.info 5 minute load average """
modules["load"] = load()

class locks(MuninPlugin):
    def fetch(self, conf):
        self.sleep_fetch(conf)
        return "locks.value %i" % self.find_locks()

    def config(self, conf):
        self.sleep_config(conf)
        return """graph_title Filesystem locks
graph_vlabel number of locks
graph_scale no
graph_info This graph shows file system lock info
graph_category system
locks.label number of locks
locks.info Number of active locks"""
modules["locks"] = locks()

class tarpit(MuninPlugin):
    "Nasty plugin that never responds"
    def fetch(self, conf):
        time.sleep(1000)

    def config(self, conf):
        time.sleep(1000)
modules["tarpit"] = tarpit()

class always_warning(MuninPlugin):
    conftext = """graph_title Always in LEVEL
graph_vlabel Level
graph_scale no
graph_info A simple graph that is always in LEVEL
graph_category always_LEVEL
generic.label Level
generic.info Level usually above warning level
generic.warning 5
generic.critical 10"""

    def fetch(self, conf):
        return "generic.value 10"

    def config(self, conf):
        return self.conftext.replace("LEVEL","warning")
modules["always_warning"] = always_warning()

class always_critical(always_warning):
    def fetch(self, conf):
        return "generic.value 20"

    def config(self, conf):
        return self.conftext.replace("LEVEL","critical")
modules["always_critical"] = always_critical()

class failing_plugin(MuninPlugin):
    "A really broken plugin"
    def fetch(self, conf):
        return "# Bad exit"

    def config(self, conf):
        return "# Bad exit"
modules["failing_plugin"] = failing_plugin()

class failing_plugin2(MuninPlugin):
    def fetch(self, conf):
        return "# Bad exit"

    def config(self, conf):
        return """graph_title Config works, fetch fails
graph_vlabel Level
graph_category failing
generic.label generic_label_here
generic.info never_really_used"""
modules["failing_plugin2"] = failing_plugin2()

class failing_plugin3(MuninPlugin):
    def config(self, conf):
        return """graph_title A plugin with two dses but only fetch value for one
graph_args --base 1000 -l 0
fivemin.label 1 minute load
onemin.label 5 minute load"""
    def fetch(self, conf):
        return "onemin.value 1"
modules["failing_plugin3"] = failing_plugin3()


class graph_area(MuninPlugin):
    "A plugin that uses STACK and AREA. From proc_pri. Use: testing the grapher"
    def fetch(self, conf):
        return """high.value 3
low.value 2
locked.value 1"""

    def config(self, conf):
        return """graph_title AREA and STACK
graph_order low high locked
graph_category graphtest
graph_info This graph shows nuber of processes at each priority
graph_args --base 1000 -l 0
graph_vlabel Number of processes
high.label high priority
high.draw STACK
high.info The number of high-priority processes (tasks)
low.label low priority
low.draw AREA
low.info The number of low-priority processes (tasks)
locked.label locked in memory
locked.draw STACK
locked.info The number of processes that have pages locked into memory (for real-time and custom IO)
"""
modules["graph_area"] = graph_area()

class utf8_graphcat(MuninPlugin):
    "A plugin with a graph category which has UTF-8 in it"
    def fetch(self, conf):
        return "apples.value %.2f" % self.find_load()

    def config(self, conf):
        return """graph_title Example UTF-8 graph
graph_vlabel apples
graph_category foo™
apples.label apples
graph_info Apples eaten
apples.info Apples eaten"""
modules["utf8_graphcat"] = utf8_graphcat()

class utf8_graphname(MuninPlugin):
    "A plugin with a UTF-8 name"
    def fetch(self, conf):
        return "apples.value %.2f" % self.find_load()

    def config(self, conf):
        return """graph_title Example UTF-8 graph
graph_vlabel apples
graph_category system
apples.label apples
graph_info Apples eaten
apples.info Apples eaten"""
modules["utf8_™graphname"] = utf8_graphname()


class ArgumentTCPserver(SocketServer.ThreadingTCPServer):
    def __init__(self, server_address, RequestHandlerClass, args):
        SocketServer.ThreadingTCPServer.__init__(self,server_address, RequestHandlerClass)
        self.args = args


class MuninHandler(SocketServer.StreamRequestHandler):
    """
    Munin server implementation.

    This is based on munin_node.py by Chris Holcombe / http://sourceforge.net/projects/pythonmuninnode/

    Possible commands:
    list, nodes, config, fetch, version or quit
    """

    def handle(self):
        if self.server.args.get("verbose"): print "%s: Connection from %s:%s. server args is %s" \
            % (self.server.args["name"], self.client_address[0], self.client_address[1], self.server.args)
        # slow path
        hostname = self.server.args["name"]
        full_hostname = hostname

        moduleprofile = self.server.args["pluginprofile"]
        modulenames = set(moduleprofile)

        self.wfile.write("# munin node at %s\n" % hostname)

        while True:
            line = self.rfile.readline().strip()
            try:
                cmd, args = line.split(" ", 1)
            except ValueError:
                cmd = line
                args = ""

            if not cmd or cmd == "quit":
                break

            if cmd == "list":
                # List all plugins that are available
                self.wfile.write(" ".join(self.server.args["plugins"].keys()) + "\n")
            elif cmd == "nodes":
                # We just support this host
                self.wfile.write("%s\n.\n" % full_hostname)
            elif cmd == "config":
                # display the config information of the plugin
                if not self.server.args["plugins"].has_key(args):
                    self.wfile.write("# Unknown service\n.\n" )
                else:
                    config = self.server.args["plugins"][args].config(self.server.args)
                    if config is None:
                        self.wfile.write("# Unknown service\n.\n")
                    else:
                        self.wfile.write(config + "\n.\n")
            elif cmd == "fetch":
                # display the data information as returned by the plugin
                if not self.server.args["plugins"].has_key(args):
                    self.wfile.write("# Unknown service\n.\n")
                else:
                    data = self.server.args["plugins"][args].fetch(self.server.args)
                    if data is None:
                        self.wfile.write("# Unknown service\n.\n")
                    else:
                        self.wfile.write(data + "\n.\n")
            elif cmd == "version":
                # display the server version
                self.wfile.write("munin node on %s version: %s\n" %
                                 (full_hostname, VERSION))
            else:
                self.wfile.write("# Unknown command. Try list, nodes, " \
                                 "config, fetch, version or quit\n")


def start_servers(instances):
    # TODO: Listen to IPv6
    HOST = "0.0.0.0"
    servers = {}
    for iconf in instances:
            print "Setting up instance %s at port %s" \
                % (iconf["name"], iconf["expanded_port"])

            server = ArgumentTCPserver((HOST, iconf["expanded_port"]), MuninHandler, iconf)
            server_thread = threading.Thread(target=server.serve_forever)
            server_thread.daemon = True
            server_thread.start()

            servers[iconf["name"]] = server
    return servers



def usage():
    print "Usage: %s [--run] [--verbose] [--muninconf] <configfile> <configfileN>" % sys.argv[0]

def main():
    if len(sys.argv) <= 2:
        usage()
        sys.exit(1)

    verbose = False
    if "--verbose" in sys.argv:
        verbose = True

    config = ConfigParser.RawConfigParser()
    for configfile in sys.argv[1:]:
        if not configfile.endswith(".conf"):
            continue
        if verbose:
            print "Reading config file %s" % configfile
        config.read(configfile)

    instancekeys = [ key for key in config.sections() if key.startswith("instance:") ]
    servers = {}

    instances = []

    for key in instancekeys:
        instancename = key.split(":", 2)[1]
        portrange = []
        if config.has_option(key, "port"):
            portrange = [ config.getint(key, "port") ]
        if config.has_option(key, "portrange"):
            rangestr = config.get(key, "portrange")
            ranges = rangestr.split("-")
            range_expanded = range(int(ranges[0]), int(ranges[1])+1, 1)
            portrange += range_expanded

        if len(portrange) == 0:
            print "WARN: No port or portrange defined for instance %s" \
                % instancename

        pluginprofile = "pluginprofile:%s" % config.get(key, "pluginprofile")
        if not config.has_section(pluginprofile):
            print "WARN: Definition for pluginprofile %s not found, skipping" \
                % config.get(key, "pluginprofile")
            continue

        plugins = {}
        tentative_pluginlist = config.get(pluginprofile, "plugins").split(",")
        assert(len(tentative_pluginlist) > 0)
        for tentative_plugin in tentative_pluginlist:
            tentative_plugin = tentative_plugin.strip()
            if not modules.has_key(tentative_plugin):
                print "WARN: Pluginprofile %s specifies unknown plugin %s" \
                    % (pluginprofile, tentative_plugin)
                continue

            # support more than one instanciation of the same plugin.
            plugininstancename = tentative_plugin
            i=2
            while (plugins.has_key(plugininstancename)):
                plugininstancename = tentative_plugin + str(i)
                i += 1

            plugins[plugininstancename] = modules[tentative_plugin]

        for portinstance in portrange:
            instanceconfig = dict()

            for k,v in config.items(key):
                instanceconfig[k] = v

            instanceconfig["plugins"] = plugins
            instanceconfig["verbose"] = verbose

            instanceconfig["name"] = "%s-%s" % (instancename, portinstance)
            instanceconfig["expanded_port"] = portinstance

            instances.append(instanceconfig)
            # XXX: need to store what handlers we should have.
    print instances

    # output sample munin config for the poller
    if "--muninconf" in sys.argv:
        for i in instances:
            print "[%s;%s]\n\taddress %s\n\tuse_node_name yes\n\tport %s\n" \
                % ( "fromhell", i["name"], config.get("base","hostname"), i["port"])


    if "--run" in sys.argv:
        if verbose: print "Starting up.."
        servers = start_servers(instances)

        try:
            while True:
                time.sleep(0.5)
        except KeyboardInterrupt:
            print "Caught Ctrl-c, shutting down.."
            for port, server in servers.items():
                server.shutdown()
                sys.exit(0)

if __name__ == "__main__":
    main()