dt-varnish-config

#!/usr/bin/python

import argparse
import os
import socket
import re
import docker

def parse_arguments():
	class DefaultList(list):
		def __copy__(self):
			return []

	parser = argparse.ArgumentParser(description="Rewrite the backends and vhosts in the varnish config file for docker containers.", formatter_class=argparse.ArgumentDefaultsHelpFormatter)
	parser.add_argument("--container", type=str, default="varnish", help="id or name of varnish container")
	parser.add_argument("--graph-storage-dir", type=str, default="/var/lib/docker", help="docker graph storage directory")
	parser.add_argument("--backends-file", type=str, default="/etc/varnish/docker_backends.vcl", help="output file for backends in container")
	parser.add_argument("--vhosts-file", type=str, default="/etc/varnish/docker_vhosts.vcl", help="output file for vhosts in container")
	parser.add_argument("--domain", type=str, default=DefaultList([socket.getfqdn()]), action="append", dest="domains", help="domain name to append to each host")
	parser.add_argument("--dry-run", action="store_true", help="only print to stdout instead of writing to file")

	return parser.parse_args()


def get_varnish_rootfs(args, client):
	inspect = client.inspect_container(args.container)
	
	return os.path.join(args.graph_storage_dir, "containers", inspect["ID"], "rootfs")


def get_domains(args, name, inspect=None):
	domains = set()
	domains.add(name)
	for domain in args.domains:
		domains.add(name + "." + domain)
	if inspect and inspect["Config"]["Hostname"]:
		domains.add(inspect["Config"]["Hostname"])

	return domains


def get_docker_backends_and_directors(args, client):
	containers = client.containers()

	backends = []
	directors = {}
	for container in containers:
		inspect = client.inspect_container(container["Id"])
		name = inspect["Name"].lstrip("/")

		backend = {}
		backend["name"] = re.sub(r"[^\w]", "_", name) 
		backend["host"] = name
		backend["port"] = "80"

		# create director if necessary
		match = re.match(r"(.*)-\d+$", name)
		if match:
			director_name = match.group(1)

			director = directors.setdefault(director_name, {})
			director["name"] = director_name
			director.setdefault("domains", set())
			director["domains"] = director["domains"] | get_domains(args, director_name, inspect)
			director.setdefault("backend_names", [])
			director["backend_names"].append(backend["name"])
	
			# do not use hostname as domain because director uses it	
			backend["domains"] = get_domains(args, name)
		else:
			backend["domains"] = get_domains(args, name, inspect)

		backends.append(backend)

	return (backends, directors.values())


def write_backends_file(args, rootfs, backends, directors):
	output = "\n".join(["backend %(name)s { .host = \"%(host)s\"; .port = \"%(port)s\"; }" % backend for backend in backends])
	output += "\n\n"
	output += "\n".join(["director %(name)s round-robin {\n%(backends)s\n}" % {"name": director["name"], "backends": "\n".join(["  { .backend = %s; }" % backend_name for backend_name in director["backend_names"]])} for director in directors]) 

	if args.dry_run:
		print output
	else:
		with open(os.path.normpath(os.sep.join([rootfs, args.backends_file])), "w") as f:
			f.write(output)
		print "Replaced varnish backends in %s" % args.backends_file


def write_vhosts_file(args, rootfs, backends, directors):
	output = "sub vcl_recv {\n"
	output += "\n".join(["  else if (req.http.host ~ \"(?i)(%(re)s)$\") {\n    set req.backend = %(name)s;\n  }" % {"re": "|".join([re.escape(domain) for domain in backend["domains"]]), "name": backend["name"]} for backend in directors + backends]) 
	output += "\n  else {\n    error 404 \"Unknown virtual host\";\n  }\n}";

	# replace first else if with just if
	output = re.sub(r"else if", "if", output, 1)

	if args.dry_run:
		print output
	else:
		with open(os.path.normpath(os.sep.join([rootfs, args.vhosts_file])), "w") as f:
			f.write(output)
		print "Replaced varnish vhosts in %s" % args.vhosts_file


args = parse_arguments()
client = docker.Client()
rootfs = get_varnish_rootfs(args, client)
backends, directors = get_docker_backends_and_directors(args, client)
write_backends_file(args, rootfs, backends, directors)
write_vhosts_file(args, rootfs, backends, directors)