make_tutum_config

#!/usr/bin/python

import os
import sys
try:
    from configparser import RawConfigParser # Python 3.0+
except:
    from ConfigParser import RawConfigParser
import argparse
import tutum
from jinja2 import Environment
from hashlib import sha256

# plagiarising myself
# this is based off of https://vcs.uhoreg.ca/git/cgit/hunchback/
class hashfile(object):
    __slots__ = ['h', 'digest', 'hexdigest']

    def __init__(self):
        self.h = sha256()
        self.digest = self.h.digest
        self.hexdigest = self.h.hexdigest

    def write(self, data):
        self.h.update(data)

    def close(self):
        pass

class teefile(object):
    __slots__ = ['children']

    def __init__(self, *args):
        self.children = args

    def write(self, data):
        for child in self.children:
            child.write(data)

    def close(self):
        for child in self.children:
            child.close()

class cachedfile(object):
    __slots__ = ['contents']

    def __init__(self):
        self.contents = []

    def write(self, data):
        self.contents.append(data)

    def close(self):
        pass

    def get_contents(self):
        return ''.join(self.contents)

class ExtraOptsAction(argparse.Action):
    def __call__(self, parser, namespace, values, option_string):
        s = values.split('@', 1)
        section = s[0] or 'check_tutum'
        if len(s) == 1:
            pass #FIXME
        else:
            config = RawConfigParser()
            config.read(s[1])
            for name, value in config.items(section):
                setattr(namespace, name, value)

def indexByUri(collection):
    return dict(map(lambda x: (x.resource_uri, x), collection))

def writeConfig(f):
    env = Environment()

    clusters = indexByUri(tutum.NodeCluster.list())
    clustertempl = env.from_string('''\
define host {
    use         tutum-cluster
    host_name   {{cluster.name}}
    address     0.0.0.0
    _tutum_id   {{cluster.uuid}}
    hostgroups  docker-servers,tutum-cluster
}

define hostgroup {
    hostgroup_name {{cluster.name}}-cluster
    alias          {{cluster.name}}
}

''')
    for cluster in clusters.itervalues():
        f.write(clustertempl.render(cluster=cluster))

    # FIXME: index clusters, services, etc by resource_uri, to avoid linear search

    nodes = indexByUri(tutum.Node.list())
    nodetempl = env.from_string('''\
define host {
    use         generic-host
    host_name   {{node.external_fqdn}}
    address     {{node.public_ip}}
    _tutum_id   {{node.uuid}}
    parents     {{cluster.name}}
    hostgroups  {{cluster.name}}-cluster,docker-servers,tutum-node
}

define service {
    use generic-service
    host_name {{node.external_fqdn}}
    service_description Tutum node
    check_command check_tutum_node!$_HOSTTUTUM_ID$
}

''')
    for node in nodes.itervalues():
        cluster = clusters[node.node_cluster]
        f.write(nodetempl.render(node=node, cluster=cluster))

    stacks = indexByUri(tutum.Stack.list())
    stacktempl = env.from_string('''\
define host {
    use         tutum-stack
    host_name   {{stack.name}}.stack
    address     0.0.0.0
    _tutum_id   {{stack.uuid}}
    hostgroups  docker-servers,tutum-stack
}

define hostgroup {
    hostgroup_name {{stack.name}}-stack
    alias          {{stack.name}}
}

''')
    for stack in stacks.itervalues():
        if stack.state == 'Terminated' or stack.state == 'Terminating':
            continue
        f.write(stacktempl.render(stack=stack))

    services = indexByUri(tutum.Service.list())
    servicetempl = env.from_string('''\
define host {
    use         tutum-service
    host_name   {{service.public_dns}}
    address     0.0.0.0
    _tutum_id   {{service.uuid}}{% if stack %}
    parents     {{stack.name}}.stack{% endif %}
    hostgroups  docker-servers{% if stack %},{{stack.name}}-stack{% endif %},tutum-service
}

define hostgroup {
    hostgroup_name {{service.name}}{% if stack %}--{{stack.name}}{% endif %}-service
    alias          {{service.name}}
}

''')
    for service in services.itervalues():
        if service.state == 'Terminated' or service.state == 'Terminating':
            continue
        stack = service.stack and stacks[service.stack]
        f.write(servicetempl.render(service=service, stack=stack))

    containers = tutum.Container.list()
    containertempl = env.from_string('''\
define host {
    use         tutum-container
    host_name   {{container.public_dns}}
    address     {% if container.container_ports and container.container_ports|selectattr("published")|list|length == 0 %}{{container.private_ip}}{% else %}{{node.public_ip}}{% endif %}
    _tutum_id   {{container.uuid}}
    parents     {{service.public_dns}},{{node.external_fqdn}}
    hostgroups  {{service.name}}{% if stack %}--{{stack.name}}{% endif %}-service,docker-servers,tutum-container
}

''')
    for container in containers:
        if container.state == 'Terminated' or container.state == 'Terminating':
            continue
        service = services[container.service]
        stack = service.stack and stacks[service.stack]
        node = nodes[container.node]
        f.write(containertempl.render(container=container, node=node, service=service, stack=stack))

if __name__ == '__main__':
    argparser = argparse.ArgumentParser(usage='Usage: %(prog)s [options]',
                                        description='build nagios configuration for Tutum nodes and services')
    argparser.add_argument('--version', action='version', version='''\
%(prog)s 0.1
Copyright 2015 MuchLearning
''')
    argparser.add_argument('--extra-opts',
                           action=ExtraOptsAction,
                           metavar='[section][@file]',
                           help='Read options from an ini file. See http://nagiosplugins.org/extra-opts for usage and examples.  Note: if section is not specified, then it defaults to "check_tutum"')
    argparser.add_argument('-u', '--user',
                           action='store', dest='user',
                           metavar="USERNAME",
                           help="Tutum username")
    argparser.add_argument('-k', '--apikey',
                           action='store', dest='apikey',
                           metavar="KEY",
                           help="Tutum API key")
    argparser.add_argument('-o', '--output',
                           action='store', dest='output',
                           metavar="FILE", type=argparse.FileType('w'),
                           default=sys.stdout,
                           help="Output file (output to stdout if not specified)")
    argparser.add_argument('-f', '--follow',
                           action='store_true', dest='follow',
                           help="Listen to Tutum events and regenerate on changes")
    argparser.add_argument('-r', '--restart-cmd',
                           action='store', dest='restart_cmd',
                           help="Command to run when configuration is changed (only used when following Tutum events)")
    opts = argparser.parse_args()

    tutum.user = opts.user
    tutum.apikey = opts.apikey

    if opts.follow and opts.output == sys.stdout:
            argparser.error('Must specify an output file in follow mode')

    if opts.follow:
        hf = hashfile()
        tf = teefile(hf, opts.output)
        opts.output.truncate()
        writeConfig(tf)
        opts.output.flush()
        lasthash = hf.hexdigest()
        if opts.restart_cmd:
            os.system(opts.restart_cmd)
        print "listening for events"
        def process_event(event):
            # FIXME: filter out events that wouldn't cause configuration changes
            # FIXME: throttle writes/restarts?
            global lasthash
            print "received event"
            hf = hashfile()
            cf = cachedfile()
            tf = teefile(hf, cf)
            writeConfig(tf)
            if lasthash != hf.hexdigest():
                print "configuration changed"
                lasthash = hf.hexdigest()
                opts.output.seek(0)
                opts.output.truncate()
                opts.output.write(cf.get_contents())
                opts.output.flush()
                if opts.restart_cmd:
                    os.system(opts.restart_cmd)

        events = tutum.TutumEvents()
        events.on_message(process_event)
        events.run_forever()
    else:
        if opts.output != sys.stdout:
            opts.output.truncate()
        writeConfig(opts.output)