#!/usr/bin/env python3


import argparse
import copy
import os
import subprocess
import sys
import tempfile

import yaml


CONFIG_FILENAME = os.path.expanduser("~/.config/hetznertool/hetznertool.yaml")


expect_script = """\
if {[llength $argv] != 2} {
    puts stderr "Usage: $argv0 name token
    exit 2
}

lassign $argv name token

log_user 0
spawn hcloud context create $name
expect {
	"Token:" {
		send "$token\n"
		exp_continue
	}
	timeout {close}
}
lassign [wait] wait_pid spawn_id exec_rc wait_code childkilled
# couldn't exec pesign
if {$exec_rc != 0} {
	exit 1
}
# killed by signal (e.g. timeout)
if {$childkilled == "CHILDKILLED"} {
	exit 1
}
# all good?
exit $wait_code
"""


class HetznerTool:
    def run(self):
        try:
            self.main()
        except Exception as e:
            sys.stderr.write("ERROR: {}\n".format(str(e)))

    def main(self):
        os.environ["LC_ALL"] = "C"
        config = self.create_config()
        parser = Parser(config)
        parser.execute()

    def create_config(self):
        config = Config()
        config.add("ssh-key")
        config.add("pass-dir")
        config.add("pass-subdir")
        config.add("ns1")
        config.add("domain")
        config.add("dnszone-dir")
        config.add("dnszone-file")
        config.add("ansible-inventory-dir")
        config.add("style")

        config.load(CONFIG_FILENAME)
        return config


class Config:
    def __init__(self):
        self.config_vars = []
        self.profiles = {}

    def dump(self):
        for profile in self.profiles:
            for name in self.profiles[profile]:
                print(profile, name, repr(self.get(profile, name)))

    def load(self, filename):
        with open(filename) as f:
            obj = yaml.safe_load(f)
        for profile in obj:
            for name in self.config_vars:
                if name not in obj[profile]:
                    raise Exception(
                        "Profile {} is missing variable {}".format(profile, name)
                    )
            for name in obj[profile]:
                if name not in self.config_vars:
                    raise Exception(
                        "Profile {} has extra variable {}".format(profile, name)
                    )
                self.set(profile, name, obj[profile][name])

    def add(self, name):
        assert name not in self.config_vars
        self.config_vars.append(name)

    def get(self, profile, name):
        return self.profiles.get(profile).get(name, "")

    def get_profile(self, profile):
        return {name: self.get(profile, name) for name in self.config_vars}

    def set(self, profile, name, value):
        if profile not in self.profiles:
            self.profiles[profile] = {}
        self.profiles[profile][name] = value


class Parser:
    def __init__(self, config):
        self.config = config

    def execute(self):
        parser = self.create_parser()
        args = vars(parser.parse_args())
        profile = self.config.get_profile(args["profile"])
        profile.update(args)
        func = args["func"]
        func(profile)

    def create_parser(self):
        parser = argparse.ArgumentParser(
            description="Manage VMs in Hetzner Cloud, with DNS updates."
        )

        parser.add_argument(
            "-n", "--no-act", dest="act", default=True, action="store_false"
        )
        parser.add_argument("-p", "--profile", dest="profile", required=True)

        factory = parser.add_subparsers()

        create = factory.add_parser("create")
        create.add_argument("context", help="hcloud context to create VMs in")
        create.add_argument("specfile", help="file to read VM specifications from")

        servers = factory.add_parser("list")

        contexts = factory.add_parser("contexts")

        zonefile = factory.add_parser("zonefile")

        create_context = factory.add_parser("create-context")
        create_context.add_argument("context", help="create new hcloud context")
        create_context.add_argument("token", help="token for hcloud context")

        delete = factory.add_parser("delete")
        delete.add_argument("context", help="hcloud context to delete all VMs from")

        parser.set_defaults(func=self.usage)
        servers.set_defaults(func=List().run)
        create.set_defaults(func=Create().run)
        delete.set_defaults(func=Delete().run)
        contexts.set_defaults(func=ListContexts().run)
        create_context.set_defaults(func=CreateContext().run)
        zonefile.set_defaults(func=Zonefile().run)

        return parser

    def usage(self, args):
        raise Exception("Missing operation")


class Zone:
    def __init__(self, domain):
        self.domain = domain
        self.servers = []

    def servername(self, context, server):
        return "{}-{}".format(context, server)

    def dnsname(self, context, server):
        return "{}.{}".format(self.servername(context, server), self.domain)

    def add(self, context, server, addr):
        self.servers.append((context, server, addr))

    def __iter__(self):
        return iter(self.servers)


class ZoneWriter:
    def __init__(self, zonedir, zonefile, ns1):
        self.zonedir = zonedir
        self.zonefile = zonefile
        self.ns1 = ns1

    def write(self, zone):
        git = Git(self.zonedir)
        if not git.isclean():
            raise Exception("git dir unclean: {}".format(self.zonedir))

        git.pull()

        pathname = os.path.join(self.zonedir, self.zonefile)
        with open(pathname, "w") as f:
            self.write_zone(zone, f)

        if not git.isclean():
            git.add(".")
            git.commit("update zone file")
            git.push()

    def write_zone(self, zone, f):
        raise NotImplementedError()


class ZoneWriterDnsApi(ZoneWriter):
    def write_zone(self, zone, f):
        for context, server, addr in zone:
            name = zone.dnsname(context, server)
            f.write("+{}:{}:60\n".format(name, addr))


class ZoneWriterBind9(ZoneWriter):

    preamble = """\
$TTL 30
$ORIGIN h.qvarnlabs.eu.

@ IN SOA ns1.qvarnlabs.net. ops.qvarnlabs.com ({serial} 30 30 8640000 15 )
@ IN NS ns1.qvarnlabs.net.
@ IN NS ns2.qvarnlabs.net.

"""

    def write_zone(self, zone, f):
        serial = self.increment_serial()
        f.write(self.preamble.format(serial=serial))
        for context, server, addr in zone:
            name = zone.servername(context, server)
            f.write("{} IN A {}\n".format(name, addr))
        f.flush()
        self.kick_bind9()

    def increment_serial(self):
        filename = os.path.join(self.zonedir, "serial")
        with open(filename) as f:
            serial = int(f.readline().strip())
        serial += 1
        with open(filename, "w") as f:
            f.write(str(serial))
        return serial

    def kick_bind9(self):
        target = "{}:/etc/bind/{}".format(self.ns1, self.zonefile)
        filename = os.path.join(self.zonedir, self.zonefile)
        subprocess.check_call(["scp", "-q", filename, target])
        subprocess.check_call(["ssh", self.ns1, "systemctl", "reload", "bind9"])


class Inventory:
    def __init__(self, dirname):
        self.dirname = dirname

    def write(self, zone, contexts):
        for context in contexts:
            filename = os.path.join(self.dirname, "hosts.{}".format(context))
            with open(filename, "w") as f:
                servers = [server for ctx, server, _ in zone if ctx == context]
                for server in servers:
                    f.write(
                        "{} ansible_ssh_host={}\n".format(
                            server, zone.dnsname(context, server)
                        )
                    )


class Operation:
    def __init__(self):
        self.hcloud = Hcloud()

    def run(self, profile):
        raise NotImplementedError()

    def new_pass(self, profile):
        return Pass(profile["pass-dir"])

    def create_hcloud_contexts(self, profile):
        p = self.new_pass(profile)
        pass_subdir = profile["pass-subdir"]
        pass_contexts = p.listdir(pass_subdir)
        hcloud_contexts = self.hcloud.list_contexts()
        for context in pass_contexts:
            if context not in hcloud_contexts:
                token = p.get_token(pass_subdir, context)
                self.hcloud.create_context(context, token)

    def update_zone_file(self, profile):
        zonedir = profile["dnszone-dir"]
        filename = profile["dnszone-file"]
        ns1 = profile["ns1"]

        zone, contexts = self.create_zone(profile)
        klass = self.get_zone_writer(profile)
        writer = klass(zonedir, filename, ns1)
        writer.write(zone)

        self.write_inventory(profile, zone, contexts)

    def create_zone(self, profile):
        zone = Zone(profile["domain"])

        p = self.new_pass(profile)
        pass_contexts = p.listdir(profile["pass-subdir"])
        for context in pass_contexts:
            for server, addr in self.hcloud.list_servers(context):
                zone.add(context, server, addr)

        return zone, pass_contexts

    def get_zone_writer(self, profile):
        styles = {
            "dns-api": ZoneWriterDnsApi,
            "qvarnlabs": ZoneWriterBind9,
        }

        style = profile["style"]
        if style not in styles:
            raise Exception("unknown zone file style: {}".format(style))
        return styles[style]

    def write_inventory(self, profile, zone, contexts):
        inventory = Inventory(profile["ansible-inventory-dir"])
        inventory.write(zone, contexts)


class Zonefile(Operation):
    def run(self, profile):
        self.update_zone_file(profile)


class List(Operation):
    def run(self, profile):
        self.create_hcloud_contexts(profile)

        p = self.new_pass(profile)
        pass_contexts = p.listdir(profile["pass-subdir"])
        domain = profile["domain"]
        zone = Zone(domain)
        for context in pass_contexts:
            for server, addr in self.hcloud.list_servers(context):
                print(zone.dnsname(context, server), addr)


class Create(Operation):
    def run(self, profile):
        if profile["act"]:
            self.create_hcloud_contexts(profile)

        ssh_key = profile["ssh-key"]
        context = profile["context"]
        specfile = SpecFile(profile["specfile"])
        zone = Zone(profile["domain"])
        for server in specfile.servers():
            if profile["act"]:
                self.hcloud.create_server(context, ssh_key=ssh_key, **server)
                print("created", zone.dnsname(context, server["name"]))
            else:
                print("pretending to create", zone.dnsname(context, server["name"]))

        if profile["act"]:
            self.update_zone_file(profile)


class Delete(Operation):
    def run(self, profile):
        if profile["act"]:
            self.create_hcloud_contexts(profile)

        ssh_key = profile["ssh-key"]
        context = profile["context"]
        zone = Zone(profile["domain"])
        for server, _ in self.hcloud.list_servers(context):
            if profile["act"]:
                print("deleting", zone.dnsname(context, server))
                self.hcloud.delete_server(context, server)
            else:
                print("pretend-deleting", zone.dnsname(context, server))

        if profile["act"]:
            self.update_zone_file(profile)


class ListContexts(Operation):
    def run(self, profile):
        p = self.new_pass(profile)
        pass_subdir = profile["pass-subdir"]
        for context in p.listdir(pass_subdir):
            print(context)


class CreateContext(Operation):
    def run(self, profile):
        context = profile["context"]
        token = profile["token"]
        pass_dir = profile["pass-dir"]
        pass_subdir = profile["pass-subdir"]

        git = Git(pass_dir)
        if not git.isclean():
            raise Exception("pass word store is not clean")
        git.pull()

        p = self.new_pass(profile)
        p.create(pass_subdir, context, token)

        git.push()

        self.hcloud.create_context(context, token)


class Hcloud:
    def run(self, *args):
        argv = ["hcloud"] + list(args)
        output = subprocess.check_output(argv)
        return output.decode("UTF-8").splitlines()

    def list_contexts(self):
        lines = self.run("context", "list")
        return sorted(lines[1:])

    def use_context(self, context):
        self.run("context", "use", context)

    def list_servers(self, context):
        self.use_context(context)
        servers = []
        lines = self.run("server", "list")
        for line in lines[1:]:
            words = line.split()
            servers.append((words[1], words[3]))
        return servers

    def create_server(self, context, name, image=None, ssh_key=None, server_type=None):
        assert image is not None
        assert ssh_key is not None
        assert server_type is not None
        self.use_context(context)
        self.run(
            "server",
            "create",
            "--name",
            name,
            "--ssh-key",
            ssh_key,
            "--image",
            image,
            "--type",
            server_type,
        )

    def delete_server(self, context, name):
        self.use_context(context)
        self.run("server", "delete", name)

    def create_context(self, name, token):
        fd, filename = tempfile.mkstemp()
        os.write(fd, expect_script.encode("UTF-8"))
        os.close(fd)

        argv = ["expect", "-f", filename, name, token]
        subprocess.check_output(argv)


class Pass:
    def __init__(self, pwdir):
        self.env = dict(os.environ)
        self.env["PASSWORD_STORE_DIR"] = pwdir

    def run(self, *args, **kwargs):
        argv = ["pass"] + list(args)
        return subprocess.check_output(argv, env=self.env, **kwargs)

    def run_lines(self, *args):
        text = self.run(*args)
        return text.decode("UTF-8").splitlines()

    def listdir(self, dirname):
        names = []
        for line in self.run_lines("show", dirname):
            words = line.split()
            if len(words) == 2:
                names.append(words[1])
        return sorted(names)

    def get_token(self, dirname, basename):
        bin = self.run("show", "{}/{}".format(dirname, basename))
        return bin.decode("UTF-8").strip()

    def create(self, dirname, basename, value):
        fd, filename = tempfile.mkstemp()
        os.remove(filename)
        os.write(fd, value.encode("UTF-8"))
        os.lseek(fd, 0, os.SEEK_SET)

        pathname = os.path.join(dirname, basename)
        self.run("insert", "-m", pathname, stdin=fd)

        os.close(fd)


class Git:
    def __init__(self, dirname):
        self.dirname = dirname

    def run(self, *args, **kwargs):
        argv = ["git"] + list(args)
        return subprocess.check_output(argv, cwd=self.dirname, **kwargs)

    def isclean(self):
        bin = self.run("status", "--porcelain")
        return bin == b""

    def add(self, filename):
        self.run("add", filename)

    def commit(self, msg):
        self.run("commit", "-m", msg, "-q")

    def pull(self):
        self.run("pull", "-q")

    def push(self):
        self.run("push", "-q")


class SpecFile:
    def __init__(self, filename):
        with open(filename) as f:
            self.obj = yaml.safe_load(f)

    def servers(self):
        servers = []
        for server in self.obj["hosts"]:
            server = copy.deepcopy(server)
            server.update(self.obj["defaults"])
            server["server_type"] = server.pop("type")
            servers.append(server)
        return servers


HetznerTool().run()