#!/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/hetznertool2.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):
        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.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))
        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()