apply

#!/usr/bin/python3

import os
import shutil
import subprocess
import yaml
from jinja2 import Template

template_dir = os.path.join(os.environ["SNAP"], "templates")
addon_dir = os.path.join(os.environ["SNAP_USER_DATA"], "addons")

def main():
    render_templates()
    apply_addons()

def render_templates():
    shutil.rmtree(addon_dir, ignore_errors=True)
    os.mkdir(addon_dir)
    context = {
        "arch": get_snap_config("arch"),
        "pillar": {
            "dns_domain": get_snap_config("dns-domain"),
            "num_nodes": get_node_count()
        }
    }
    if get_snap_config("enable-kube-dns") == "true":
        render_template("kube-dns.yaml", context)
    if get_snap_config("enable-dashboard") == "true":
        render_template("kubernetes-dashboard.yaml", context)
        render_template("influxdb-grafana-controller.yaml", context)
        render_template("influxdb-service.yaml", context)
        render_template("grafana-service.yaml", context)
        render_template("heapster-rbac.yaml", context)
        render_template("heapster-controller.yaml", context)
        render_template("heapster-service.yaml", context)

def render_template(file, context, required=True):
    source = os.path.join(template_dir, file)
    dest = os.path.join(addon_dir, file)
    if not os.path.exists(source) and not required:
        return
    with open(source) as f:
        template = Template(f.read())
    content = template.render(context)

    # apply cdk-addons=true label
    data = [part for part in yaml.load_all(content) if part]
    for part in data:
        part["metadata"].setdefault("labels", {})
        part["metadata"]["labels"]["cdk-addons"] = "true"
        # Adding kubernetes.io/cluster-service=true label for now. We should
        # probably remove this later and stop using it after the
        # cdk-addons=true label has been live for a couple releases.
        part["metadata"]["labels"]["kubernetes.io/cluster-service"] = "true"
    content = yaml.dump_all(data)

    with open(dest, "w") as f:
        f.write(content)

def apply_addons():
    # Apply dns service first and then recursively the rest.
    if get_snap_config("enable-kube-dns") == "true":
        dns_svc = os.path.join(addon_dir, "kube-dns.yaml")
        args = ["apply",
                "-f", dns_svc,
                "--namespace=kube-system",
                "-l", "kubernetes.io/cluster-service=true",
                "--force"]
        kubectl(*args)
    args = ["apply",
            "-f", addon_dir,
            "--recursive",
            "--namespace=kube-system",
            "-l", "kubernetes.io/cluster-service=true",
            "--prune=true",
            "--force"]
    kubectl(*args)


def kubectl(*args):
    cmd = [os.path.join(os.environ["SNAP"], "kubectl")]
    kubeconfig = get_snap_config("kubeconfig", required=False)
    if kubeconfig:
        cmd += ["--kubeconfig", kubeconfig]
    cmd += list(args)
    return subprocess.check_output(cmd)

def get_node_count():
    '''Return the number of Kubernetes nodes in the cluster'''
    output = kubectl("get", "nodes", "-o", "name")
    node_count = len(output.splitlines())
    return node_count

def get_snap_config(name, required=True):
    path = os.path.join(os.environ["SNAP_DATA"], "config", name)
    with open(path) as f:
        value = f.read().rstrip()
    if not value and required:
        raise MissingSnapConfig("%s is required" % name)
    return value

class MissingSnapConfig(Exception):
    pass

if __name__ == "__main__":
    main()