#!/usr/bin/python2

import argparse
import glob
import os
import os.path
import shutil
import sys
import time

path = os.path.dirname(os.path.realpath(__file__))
path = os.path.realpath(os.path.join(path, '..', 'lib', 'python'))
if path.startswith('/opt'):
    sys.path.append(path)
import confluent.collective.manager as collective
import eventlet.green.subprocess as subprocess
import confluent.selfservice as selfservice
import confluent.util as util
import confluent.client as client
import confluent.sshutil as sshutil
import confluent.certutil as certutil
try:
    input = raw_input
except NameError:
    pass

def emprint(txt):
    if sys.stdout.isatty():
        print('\x1b[1m\x1b[4m' + txt + '\x1b[0m')
    else:
        print(txt)

fnamechars = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789.^'
def main(args):
    ap = argparse.ArgumentParser(description='Manage OS deployment resources')
    sp = ap.add_subparsers(dest='command')
    wiz = sp.add_parser('initialize', help='Do OS deployment preparation')
    wiz.add_argument('-u', help='Pull in root user key for node deployment', action='store_true')
    wiz.add_argument('-s', help='Set up SSH CA for managing node to node ssh and known hosts', action='store_true')
    wiz.add_argument('-k', help='Update local global known hosts file with confluent CA', action='store_true')
    wiz.add_argument('-t', help='Generate new TLS key for HTTPS operation and register with confluent repository', action='store_true')
    wiz.add_argument('-p', help='Copy in TFTP contents required for PXE support', action='store_true')
    wiz.add_argument('-i', help='Interactively prompt for behaviors', action='store_true')
    wiz.add_argument('-l', help='Set local management node to have SSH certificates and hosts.equiv/.shosts', action='store_true')
    osip = sp.add_parser('import', help='Import an OS image from an ISO image')
    osip.add_argument('imagefile', help='File to use for source of importing')
    cmdset = ap.parse_args()
    if cmdset.command == 'import':
        return osimport(cmdset.imagefile)
    if cmdset.command == 'initialize':
        return initialize(cmdset)
    ap.print_help()


def local_node_trust_setup():
    allnodes, domain = selfservice.get_cluster_list()
    myname = collective.get_myname()
    myprincipals = set([myname])
    neededlines = set([
        'HostbasedAuthentication yes', 'HostbasedUsesNameFromPacketOnly yes',
        'IgnoreRhosts no'])
    if not myname.endswith(domain):
        myprincipals.add('{0}.{1}'.format(myname, domain))
    for pubkey in glob.glob('/etc/ssh/ssh_host_*key.pub'):
        currpubkey = open(pubkey, 'rb').read()
        cert = sshutil.sign_host_key(currpubkey, myname, myprincipals)
        certfile = pubkey.replace('key.pub', 'key-cert.pub')
        neededlines.add('HostCertificate {0}'.format(certfile))
        if os.path.exists(certfile):
            os.unlink(certfile)
        with open(certfile, 'w') as certout:
            certout.write(cert)
    with open('/etc/ssh/sshd_config', 'r') as sshconf:
        currconfig = sshconf.read().split()
    for conline in currconfig:
        conline = conline.strip()
        neededlines.discard(conline)
    if neededlines:
        with open('/etc/ssh/sshd_config', 'a') as cfgout:
            for currline in neededlines:
                cfgout.write(currline)
                cfgout.write('\n')
    with open('/etc/ssh/shosts.equiv', 'w') as equivout:
        for node in util.natural_sort(allnodes):
            equivout.write(node + '\n')
    with open('/root/.shosts', 'w') as equivout:
        for node in util.natural_sort(allnodes):
            equivout.write(node + '\n')


def install_tftp_content():
    tftplocation = None
    candidates = ('/tftpboot', '/var/lib/tftpboot', '/srv/tftpboot', '/srv/tftp')
    for cand in candidates:
        if os.path.isdir(cand):
            tftplocation = cand
            break
    if not tftplocation:
        raise Exception('Unable to detect an installed tftp location (check that tftp server is installed)')
    if os.path.exists('/usr/lib/systemd/system/tftp.socket'):
        subprocess.check_call(['systemctl', 'enable', 'tftp.socket', '--now'])
    else:
        print(
            'Unable to automatically enable TFTP server, please enable manually')
    tftplocation = '{0}/confluent/x86_64'.format(tftplocation)
    try:
        os.makedirs(tftplocation)
    except OSError as e:
        if e.errno == 17:
            raise
    shutil.copy('/opt/confluent/lib/ipxe/ipxe.efi', tftplocation)
    shutil.copy('/opt/confluent/lib/ipxe/ipxe.kkpxe', tftplocation)


def initialize(cmdset):
    if os.getuid() != 0:
        sys.stderr.write('This command must run as root user\n')
        sys.exit(1)
    if cmdset.i:
        didsomething = True
        sys.stdout.write('Add root user key to be authorized to log into nodes (-u)? (y/n): ')
        sys.stdout.flush()
        cmdset.u = input().strip().lower().startswith('y')
        sys.stdout.write('Set up an SSH authority to help manage known_hosts and node to node ssh for all users (-s)? (y/n): ')
        cmdset.s = input().strip().lower().startswith('y')
        sys.stdout.write('Update global known hosts on this server to trust local CA certificates (-k)? (y/n): ')
        cmdset.k = input().strip().lower().startswith('y')
        sys.stdout.write('Update tftp directory with binaries to support PXE (-p) (y/n): ')
        cmdset.p = input().strip().lower().startswith('y')
        sys.stdout.write('Generate new TLS certificates for HTTP, replacing any existing certificate (-t)? (y/n): ')
        cmdset.t = input().strip().lower().startswith('y')
        if not cmdset.t:
            print(
                'In order to use your own certificate authority, make sure '
                'to put the certificate authority into '
                '/var/lib/confluent/public/site/tls/ directory as a .pem file '
                'as well as named (hash).0 where (hash) is the hash of the '
                'subject.')
    else:
        didsomething = False
    if cmdset.u:
        didsomething = True
        if not glob.glob('/root/.ssh/*.pub'):
            sys.stderr.write('No user keys for root detected, it is recommended '
                            'to run ssh-keygen -t ed25519 to generate a user '
                            'key. For optimal security, a passphrase should be '
                            'used. ssh-agent may be used to make use of a '
                            'passphrase protected ssh key easier.\n')
            sys.exit(1)
        sshutil.initialize_root_key(False)
    if cmdset.t:
        didsomething = True
        certutil.create_certificate()
        if os.path.exists('/usr/lib/systemd/system/httpd.service'):
            subprocess.check_call(['systemctl', 'try-restart', 'httpd'])
            print('HTTP server has been restarted if it was running')
        elif os.path.exists('/usr/lib/systemd/system/apache2.service'):
            subprocess.check_call(['systemctl', 'try-restart', 'apache2'])
            print('HTTP server has been restarted if it was running')
        else:
            emprint('New HTTPS certificates generated, restart the web server manually')
    if cmdset.s:
        didsomething = True
        sshutil.initialize_ca()
    if cmdset.p:
        didsomething = True
        install_tftp_content()
    if cmdset.l:
        didsomething = True
        local_node_trust_setup()
    if not didsomething:
        sys.stderr.write('Nothing was done, use initialize -i for '
                        'interactive mode, or see initialize -h for more options\n')
        sys.exit(1)
    tmpname = '/var/lib/confluent/public/site/initramfs.cpio.'
    for x in bytearray(os.urandom(22)):
        tmpname += fnamechars[x >> 2]
    topack = []
    opath = os.getcwd()
    os.chdir('/var/lib/confluent/public/site')
    topack.append('ssh/')
    for currd, _, files in os.walk('ssh'):
        for fname in files:
            topack.append(os.path.join(currd, fname))
    topack.append('tls/')
    for currd, _, files in os.walk('tls'):
        for fname in files:
            topack.append(os.path.join(currd, fname))
    with open(tmpname, 'wb') as initramfs:
        packit = subprocess.Popen(['cpio', '-H', 'newc', '-o'],
                                  stdout=initramfs, stdin=subprocess.PIPE)
        for packfile in topack:
            if not isinstance(packfile, bytes):
                packfile = packfile.encode('utf8')
            packit.stdin.write(packfile)
            packit.stdin.write(b'\n')
        packit.stdin.close()
        res = packit.wait()
        if res:
            sys.stderr.write('Error occurred while packing site initramfs')
            sys.exit(1)
    os.rename(tmpname, '/var/lib/confluent/public/site/initramfs.cpio')
    tmptarname = tmpname.replace('cpio', 'tgz')
    tarcmd = ['tar', '-czf', tmptarname, 'ssh', 'tls']
    subprocess.check_call(tarcmd)
    os.chdir(opath)
    os.rename(tmptarname, '/var/lib/confluent/public/site/initramfs.tgz')
    print('Site initramfs content packed successfully')
    if cmdset.k:
        with open('/etc/ssh/ssh_known_hosts', 'a+b') as skh:
            for cafile in glob.glob('/var/lib/confluent/public/site/ssh/*.ca'):
                cacert = open(cafile, 'rb').read()
                cacert = b'@cert-authority * ' + cacert
                skh.write(cacert)
    # ok, also need to think on how to handle getinstalldisk
    if not os.path.exists('/etc/confluent/srvcert.pem'):
        subprocess.check_call(['collective', 'gencert'])
    # TODO: check selinux and segetbool for httpd_can_network_connect
    # tftp-server available and enabled?
    # httpd available and enabled?


def osimport(imagefile):
    c = client.Command()
    imagefile = os.path.abspath(imagefile)
    importing = False
    shortname = None
    for rsp in c.create('/deployment/importing/', {'filename': imagefile}):
        if 'target' in rsp:
            importing = True
            shortname = rsp['name']
            print('Importing from {0} to {1}'.format(imagefile, rsp['target']))
        else:
            print(repr(rsp))
    while importing:
        for rsp in c.read('/deployment/importing/{0}'.format(shortname)):
            if 'progress' in rsp:
                sys.stdout.write('{0}: {1:.2f}%    \r'.format(rsp['phase'],
                                                          rsp['progress']))
                if rsp['phase'] == 'complete':
                    importing = False
                    sys.stdout.write('\n')
                    for profile in rsp['profiles']:
                        print('Deployment profile created: {0}'.format(profile))
                sys.stdout.flush()
            else:
                print(repr(rsp))
            time.sleep(0.5)
    if shortname:
        list(c.delete('/deployment/importing/{0}'.format(shortname)))

if __name__ == '__main__':
    main(sys.argv)