#!/usr/bin/env python

from builtins import object

import datetime
from io import BytesIO
import json
import logging
import os
import tempfile
import subprocess
import sys

from cryptography import x509
from cryptography.hazmat.backends import default_backend
import fabric

from tools import secretstore

import ca


local_root = os.getenv('hscloud_root')
if local_root is None:
    raise Exception("Please source env.sh")


cluster = 'k0.hswaw.net'
remote_root = '/opt/hscloud'
ss = secretstore.SecretStore(
    plain_root=os.path.join(local_root, 'cluster/secrets/plain'),
    cipher_root=os.path.join(local_root, 'cluster/secrets/cipher'))


logger = logging.getLogger()
logger.setLevel(logging.INFO)
formatter = logging.Formatter('%(levelname)s - %(message)s')
sh = logging.StreamHandler()
sh.setFormatter(formatter)
logger.addHandler(sh)


def pki_config(key):
    raw = subprocess.check_output([
        'nix', 'eval', '--raw',
        '( (import ' + local_root + '/nix/toplevel.nix ).pki.' + key + '.json )',
    ])
    return json.loads(raw)


def _file_exists(c, filename):
    res = c.run('stat "{}"'.format(filename), warn=True, hide=True)
    return res.exited == 0


def configure_k8s(username, ca, cert, key):
    subprocess.check_call([
        'kubectl', 'config',
        'set-cluster', cluster,
        '--certificate-authority=' + ca,
        '--embed-certs=true',
        '--server=https://' + cluster + ':4001',
    ])
    subprocess.check_call([
        'kubectl', 'config',
        'set-credentials', username,
        '--client-certificate=' + cert,
        '--client-key=' + key,
        '--embed-certs=true',
    ])
    subprocess.check_call([
        'kubectl', 'config',
        'set-context', cluster,
        '--cluster=' + cluster,
        '--user=' + username,
    ])
    subprocess.check_call([
        'kubectl', 'config',
        'use-context', cluster,
    ])


def admincreds(args):
    if len(args) != 1:
        sys.stderr.write("Usage: admincreds q3k\n")
        return 1
    username = args[0]

    ## Make kube certificates.
    certs_root = os.path.join(local_root, 'cluster/certs')
    ca_kube = ca.CA(ss, certs_root, 'kube', 'kubernetes main CA')

    local_key = os.path.join(local_root, '.kubectl/admin.key')
    local_crt = os.path.join(local_root, '.kubectl/admin.crt')

    kubectl = os.path.join(local_root, '.kubectl')
    if not os.path.exists(kubectl):
        os.mkdir(kubectl)

    generate_cert = False
    if not os.path.exists(local_key):
        generate_cert = True

    if os.path.exists(local_crt):
        with open(local_crt, 'rb') as f:
            b = f.read()
            cert = x509.load_pem_x509_certificate(b, default_backend())
            delta = cert.not_valid_after - datetime.datetime.now()
            logger.info("admin: existing cert expiry: {}".format(delta))
            if delta.total_seconds() < 3600 * 24:
                logger.info("admin: expires soon, regenerating")
                generate_cert = True
    else:
        generate_cert = True

    if not generate_cert:
        return configure_k8s(username, ca_kube._cert, local_crt, local_key)

    key, csr = ca_kube.gen_key(hosts=['admin', username], o='system:masters', ou='Kube Admin Account')
    crt = ca_kube.sign(csr)

    with open(local_key, 'w') as f:
        f.write(key)

    with open(local_crt, 'w') as f:
        f.write(crt)

    configure_k8s(username, ca_kube._cert, local_crt, local_key)


def nodestrap(args, nocerts=False):
    if len(args) != 1:
        sys.stderr.write("Usage: nodestrap bc01n01.hswaw.net\n")
        return 1
    fqdn = args[0]

    logger.info("Nodestrapping {}...".format(fqdn))
    r = fabric.Connection('root@{}'.format(fqdn))

    if not nocerts:
        cfg = dict((k, pki_config(k)) for k in [
            'etcdPeer', 'etcd.server', 'etcd.kube'
        ])
        certs_root = os.path.join(local_root, 'cluster/certs')

        # Make etcd peer certificate for node.
        ca_etcd_peer = ca.CA(ss, certs_root, 'etcdpeer', 'etcd peer ca')
        ca_etcd_peer.upload(r, cfg['etcdPeer']['ca'])
        c = ca_etcd_peer.make_cert('etcdpeer-{}'.format(fqdn), hosts=[fqdn], ou='node etcd peer certificate')
        c.upload_pki(r, cfg['etcdPeer'])

        # Make etcd server certificate for node and client certificate for kube.
        ca_etcd = ca.CA(ss, certs_root, 'etcd', 'etcd ca')
        ca_etcd.upload(r, cfg['etcd.server']['ca'])

        c = ca_etcd.make_cert('etcd-{}'.format(fqdn), hosts=[fqdn], ou='node etcd server certificate')
        c.upload_pki(r, cfg['etcd.server'])

        c = ca_etcd.make_cert('etcd-kube', hosts=['kube'], ou='kube etcd client certificate')
        c.upload_pki(r, cfg['etcd.kube'])

        # Make root etcd client (do not upload).
        ca_etcd.make_cert('etcd-root', hosts=['root'], ou='root etcd client certificate')

        # Make calico etcd client (do not upload, used by jsonnet).
        ca_etcd.make_cert('etcd-calico', hosts=['calico'], ou='root etcd client certificate')

        ## Make kube certificates.
        ca_kube = ca.CA(ss, certs_root, 'kube', 'kubernetes main CA')

        # Make kubelet certificate (per node).
        c = ca_kube.make_cert('kube-kubelet-'+fqdn, o='system:nodes', ou='Kubelet', hosts=['system:node:'+fqdn, fqdn])
        c.upload_pki(r, pki_config('kube.kubelet'))

        # Make apiserver certificate.
        c = ca_kube.make_cert('kube-apiserver', ou='Kubernetes API', hosts=[cluster, '10.10.12.1'])
        c.upload_pki(r, pki_config('kube.apiserver'), concat_ca=True)

        # Make service accounts decryption key (as cert for consistency).
        c = ca_kube.make_cert('kube-serviceaccounts', ou='Kubernetes Service Accounts Signer', hosts=['serviceaccounts'])
        c.upload_pki(r, pki_config('kube.serviceaccounts'))

        # Make kube component certificates.
        kube_components = ['controllermanager', 'scheduler', 'proxy']
        cfg = dict((k, pki_config('kube.' + k)) for k in kube_components)
        for k in kube_components:
            ca_kube.upload(r, cfg[k]['ca'])
            # meh 
            if k == 'controllermanager':
                o = 'system:kube-controller-manager'
            else:
                o = 'system:kube-'+k
            ou = 'Kubernetes Component '+k
            c = ca_kube.make_cert('kube-'+k, ou=ou, o=o, hosts=[o,])
            c.upload_pki(r, cfg[k])

        ## Make kubefront certificates.
        ca_kubefront = ca.CA(ss, certs_root, 'kubefront', 'kubernetes frontend CA')
        ca_kubefront.upload(r, pki_config('kubeFront.apiserver')['ca'])
        c = ca_kubefront.make_cert('kubefront-apiserver', ou='Kubernetes Frontend', hosts=['apiserver'])
        c.upload_pki(r, pki_config('kubeFront.apiserver'))

    # Upload NixOS config
    for f in ['toplevel', 'cluster-configuration']:
        r.put(local=os.path.join(local_root, 'nix/{}.nix'.format(f)),
              remote='/etc/nixos/{}.nix'.format(f))

    r.run('nixos-rebuild switch')


def usage():
    sys.stderr.write("Usage: {} <nodestrap|admincreds|config>\n".format(sys.argv[0]))


def main():
    if len(sys.argv) < 2:
        usage()
        return 1

    mode = sys.argv[1]
    if mode == "nodestrap":
        return nodestrap(sys.argv[2:])
    elif mode == "nodestrap-nocerts":
        return nodestrap(sys.argv[2:], nocerts=True)
    elif mode == "admincreds":
        return admincreds(sys.argv[2:])
    elif mode == "config":
        print('etcd peer:')
        print(json.dumps(pki_config('etcdPeer'), indent=2))
        print('etcd client:')
        print(json.dumps(pki_config('etcdClient'), indent=2))
    else:
        usage()
        return 1

if __name__ == '__main__':
    sys.exit(main() or 0)