#!/usr/bin/env python

import os, sys, json, re
import docker
import argparse

parser = argparse.ArgumentParser( description="Manage host entries for docker containers")
parser.add_argument('-n', '--native', help='run service on docker host', action='store_true' , dest='NATIVE_MODE')
parser.add_argument('-o','-1', '--once', help='run once and exit (no continuous polling)', action='store_true', dest='ONCE_MODE')
args = parser.parse_args()

if not args.NATIVE_MODE and os.path.exists("/.dockerenv"):
    print "Found /.dockerenv... reading from /tmp/hosts"
    HOSTS_FILE="/tmp/hosts"
else:
    print "Running in native mode... reading from /etc/hosts"
    HOSTS_FILE="/etc/hosts"

if os.environ.get("DOCKER_HOST"):
    client = docker.from_env()
else:
    try:
        client = docker.DockerClient(base_url='unix://var/run/docker.sock')
    except:
        print "Sorry, I can't find a dockerd to connect to."
        sys.exit(1)

def get_ip(c):
    try:
        networks = c.attrs.get("NetworkSettings").get("Networks")

        for n in networks:
            ip = networks[n]["IPAddress"]
            if ip is not None:
                return ip
    except:
        return None

def get_aliases(c):
    aliases = []
    try:
        envs = c.attrs.get("Config").get("Env")

        for e in envs:
            m = re.match('ALIAS=(?P<aliases>.*)', e)
            if m:
                for a in m.group('aliases').split(','):
                    aliases.append(a)
        return aliases
    except KeyError as e:
        return None

def event_remove(event):
    container = client.containers.get(event.get("id"))
    container_remove(container.name)

def event_add(event):
    container = client.containers.get(event.get("id"))
    container_add(container.name)

def container_add(name):
    try:
        container = client.containers.get(name)
    except:
        return

    container_remove(container.name)
    ip = get_ip(container)
    if ip:
        entry = "%s %s" % (ip, container.name)
        aliases = get_aliases(container)
        if aliases:
            entry += " %s" % ' '.join(aliases)

        open(HOSTS_FILE, 'a').write(entry+'\n')
        print "Added: %s" % entry

def container_remove(name):
    if name is None and len(name) > 0:
        return
    entry = "(.*)\s+%s(.*)" % (name)

    lines = open(HOSTS_FILE, 'r').readlines()
    outfile = open(HOSTS_FILE, 'w')

    for line in lines:
        if re.match(entry, line) is None:
            outfile.write(line)
    outfile.close()

    print "removed entry for pattern %s" % name


### Let's start doing something useful.
# first, we catch up with any existing containers that are already deployed.
for container in client.containers.list():
    container_add(container.name)

if args.ONCE_MODE:
    sys.exit(0)

# calling events will basically poll forever, so after this we're
#  essentially in not-quite-a-daemon mode
for event_json in client.events():
    print "SAW: %s" % str(event_json)
    event = json.loads(str(event_json))

    if event.get("status") == "start":
        event_add(event)
    elif event.get("status") == "kill":
        event_remove(event)