#!/usr/bin/env python
# vim: syn=python:ts=4:sw=4:et:ai

# Usage: consul_tag_helper
#
# This program is in the Public Domain.
# 
# Place your consul configuration files in /etc/consul.boilerplate.d, instead
# of /etc/consul.d.  Boilerplate JSON configuration files can contain an extra
# key under 'service' called 'tag_program'.  'tag_program' is a path to an 
# arbitrary program that emits tags, one per line, to standard output.
#
# consul_tag_helper will iterate through each boilerplate configuration file,
# run the tag helper script, insert the emitted tags into the boilerplate
# configuration, save the generated configuration to /etc/consul.d, and 
# reload the Consul agent.

import os
import shlex
import subprocess
import tempfile
import json
import hashlib
import time

BOILERPLATE_DIR = '/etc/consul.boilerplate.d'
CONSUL_CONFDIR = '/etc/consul.d'

def service_change_pending(config, filename):
    pending_json = json.dumps(config, indent=3, sort_keys=True)
    pending_digest = hashlib.md5(pending_json).digest()
    current_digest = hashlib.md5(open(filename, 'r').read()).digest()
    if pending_digest != current_digest:
        return True
    else:
        return False


def update_service(config, filename, uid=0, gid=0, mode=0644):
    print 'Service %s has changed tags.  Rewriting configuration file and reloading consul.' % config['service']['name']
    try:
        tmp_consul_config_file = tempfile.NamedTemporaryFile(prefix=CONSUL_CONFDIR + '/', delete=False)
        json.dump(config, tmp_consul_config_file, indent=3, sort_keys=True)
        tmp_consul_config_file.close()
        os.chown(tmp_consul_config_file.name, uid, gid)
        os.chmod(tmp_consul_config_file.name, mode)
        os.rename(tmp_consul_config_file.name, filename)
        subprocess.check_call(['consul', 'reload'])
    finally:
        if os.path.exists(tmp_consul_config_file.name):
            os.remove(tmp_consul_config_file.name)


while True:
    if os.path.exists(BOILERPLATE_DIR):
        for basename in os.listdir(BOILERPLATE_DIR):
            boilerplate_filename = os.path.join(BOILERPLATE_DIR, basename)
            consul_config_filename = os.path.join(CONSUL_CONFDIR, basename)
            stat = os.stat(boilerplate_filename)
            with open(boilerplate_filename, 'r') as boilerplate_file:
                config = json.load(boilerplate_file)
                if 'service' in config:
                    service = config['service']
                    if 'tag_program' in service:
                        tag_program = service['tag_program']
                        del service['tag_program']
                        tag_program = shlex.split(tag_program.encode('ascii', 'replace'))   # Need to encode as shelx.split() in Python 2.6.x only accepts non unicode strings
                        tags = subprocess.Popen(tag_program, stdout=subprocess.PIPE).communicate()[0].splitlines()
                        if 'tags' not in service:
                            service['tags'] = []
                        service['tags'] += tags
                        if not os.path.exists(consul_config_filename):
                            update_service(config, consul_config_filename, uid=stat.st_uid, gid=stat.st_gid, mode=stat.st_mode)
                        elif service_change_pending(config, consul_config_filename):
                            update_service(config, consul_config_filename, uid=stat.st_uid, gid=stat.st_gid, mode=stat.st_mode)
    time.sleep(5)
done