confluent/confluent_server/bin/osimage

#!/usr/bin/python2

import argparse
import eventlet.green.subprocess as subprocess
import glob
import os
import os.path
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.client as client
import confluent.sshutil as sshutil
import confluent.certutil as certutil
try:
    input = raw_input
except NameError:
    pass

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('-t', help='Generate new TLS key for HTTPS operation and register with confluent repository', action='store_true')
    wiz.add_argument('-i', help='Interactively prompt for behaviors', 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 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('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:
        if not glob.glob('/root/.ssh/*.pub'):
            didsomething = True
            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 cmdset.s:
        didsomething = True
        sshutil.initialize_ca()

    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')
    for currd, _, files in os.walk('tls'):
        for fname in files:
            topack.append(os.path.join(currd, fname))
    for currd, _, files in os.walk('ssh'):
        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.chdir(opath)
    os.rename(tmpname, '/var/lib/confluent/public/site/initramfs.cpio')

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)
    list(c.delete('/deployment/importing/{0}'.format(shortname)))

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