#!/usr/bin/python import eventlet import eventlet.green.select as select import eventlet.green.subprocess as subprocess import glob import logging logging.getLogger('libarchive').addHandler(logging.NullHandler()) import libarchive import hashlib import os import shutil import sys import time import yaml COPY = 1 EXTRACT = 2 READFILES = set([ 'README.diskdefines', 'media.1/products', 'media.2/products', '.discinfo', ]) HEADERSUMS = set([b'\x85\xeddW\x86\xc5\xbdhx\xbe\x81\x18X\x1e\xb4O\x14\x9d\x11\xb7C8\x9b\x97R\x0c-\xb8Ht\xcb\xb3']) HASHPRINTS = { '69d5f1c5e4474d70b0fb5374bfcb29bf57ba828ff00a55237cd757e61ed71048': {'name': 'cumulus-broadcom-amd64-4.0.0', 'method': COPY}, } from ctypes import byref, c_longlong, c_size_t, c_void_p from libarchive.ffi import ( write_disk_new, write_disk_set_options, write_free, write_header, read_data_block, write_data_block, write_finish_entry, ARCHIVE_EOF ) def update_boot(profiledir): profile = {} if profiledir.endswith('/'): profiledir = profiledir[:-1] profname = os.path.basename(profiledir) with open('{0}/profile.yaml'.format(profiledir)) as profileinfo: profile = yaml.safe_load(profileinfo) label = profile.get('label', profname) kernelargs = profile.get('kernelargs', '') grubcfg = "set timeout=5\nmenuentry '" grubcfg += label grubcfg += "' {\n linuxefi /kernel " + kernelargs + "\n" initrds = [] for initramfs in glob.glob(profiledir + '/boot/initramfs/*.cpio'): initramfs = os.path.basename(initramfs) initrds.append(initramfs) for initramfs in os.listdir(profiledir + '/boot/initramfs'): if initramfs not in initrds: initrds.append(initramfs) grubcfg += " initrdefi " for initramfs in initrds: grubcfg += " /initramfs/{0}".format(initramfs) grubcfg += "\n}\n" with open(profiledir + '/boot/efi/boot/grub.cfg', 'w') as grubout: grubout.write(grubcfg) ipxeargs = kernelargs for initramfs in initrds: ipxeargs += " initrd=" + initramfs with open(profiledir + '/boot/boot.ipxe', 'w') as ipxeout: ipxeout.write('#!ipxe\n') ipxeout.write('imgfetch kernel ' + ipxeargs + '\n') for initramfs in initrds: ipxeout.write('imgfetch initramfs/{0}\n'.format(initramfs)) ipxeout.write('imgload kernel\nimgexec kernel\n') subprocess.check_call( ['dir2img', '{0}/boot'.format(profiledir), '{0}/boot.img'.format(profiledir)]) def extract_entries(entries, flags=0, callback=None, totalsize=None, extractlist=None): """Extracts the given archive entries into the current directory. """ buff, size, offset = c_void_p(), c_size_t(), c_longlong() buff_p, size_p, offset_p = byref(buff), byref(size), byref(offset) sizedone = 0 printat = 0 with libarchive.extract.new_archive_write_disk(flags) as write_p: for entry in entries: if str(entry).endswith('TRANS.TBL'): continue if extractlist and str(entry) not in extractlist: continue write_header(write_p, entry._entry_p) read_p = entry._archive_p while 1: r = read_data_block(read_p, buff_p, size_p, offset_p) sizedone += size.value if callback and time.time() > printat: callback({'progress': float(sizedone) / float(totalsize)}) printat = time.time() + 0.5 if r == ARCHIVE_EOF: break write_data_block(write_p, buff, size, offset) write_finish_entry(write_p) if callback: callback({'progress': float(sizedone) / float(totalsize)}) def extract_file(filepath, flags=0, callback=lambda x: None, imginfo=(), extractlist=None): """Extracts an archive from a file into the current directory.""" totalsize = 0 for img in imginfo: if not imginfo[img]: continue totalsize += imginfo[img] with libarchive.file_reader(filepath) as archive: extract_entries(archive, flags, callback, totalsize, extractlist) def check_centos(isoinfo): ver = None arch = None for entry in isoinfo[0]: if 'centos-release-7' in entry: dotsplit = entry.split('.') arch = dotsplit[-2] ver = dotsplit[0].split('release-')[-1].replace('-', '.') break elif 'centos-release-8' in entry: ver = entry.split('-')[2] arch = entry.split('.')[-2] break else: return None return {'name': 'centos-{0}-{1}'.format(ver, arch), 'method': EXTRACT} def check_ubuntu(isoinfo): if 'README.diskdefines' not in isoinfo[1]: return None arch = None variant = None ver = None diskdefs = isoinfo[1]['README.diskdefines'] for info in diskdefs.split(b'\n'): if not info: continue _, key, val = info.split(b' ', 2) val = val.strip() if key == b'ARCH': arch = val if arch == b'amd64': arch = b'x86_64' elif key == b'DISKNAME': variant, ver, _ = val.split(b' ', 2) if variant != b'Ubuntu-Server': return None if variant: if not isinstance(ver, str): ver = ver.decode('utf8') if not isinstance(arch, str): arch = arch.decode('utf8') major = '.'.join(ver.split('.', 2)[:2]) return {'name': 'ubuntu-{0}-{1}'.format(ver, arch), 'method': EXTRACT|COPY, 'extractlist': ['casper/vmlinuz', 'casper/initrd', 'EFI/BOOT/BOOTx64.EFI', 'EFI/BOOT/grubx64.efi' ], 'copyto': 'install.iso', 'category': 'ubuntu{0}'.format(major)} def check_sles(isoinfo): ver = None arch = 'x86_64' disk = None distro = '' if 'media.1/products' in isoinfo[1]: medianame = 'media.1/products' elif 'media.2/products' in isoinfo[1]: medianame = 'media.2/products' else: return None prodinfo = isoinfo[1][medianame] if not isinstance(prodinfo, str): prodinfo = prodinfo.decode('utf8') prodinfo = prodinfo.split('\n') hline = prodinfo[0].split(' ') ver = hline[-1].split('-')[0] major = ver.split('.', 2)[0] if hline[-1].startswith('15'): if hline[1] == 'openSUSE-Leap': distro = 'opensuse_leap' else: distro = 'sle' if hline[0] == '/': disk = '1' elif hline[0].startswith('/Module'): disk = '2' elif hline[-1].startswith('12'): if 'SLES' in hline[1]: distro = 'sles' if '.1' in medianame: disk = '1' elif '.2' in medianame: disk = '2' if disk and distro: return {'name': '{0}-{1}-{2}'.format(distro, ver, arch), 'method': EXTRACT, 'subname': disk, 'category': 'suse{0}'.format(major)} return None def check_rhel(isoinfo): ver = None arch = None for entry in isoinfo[0]: if 'redhat-release-7' in entry: dotsplit = entry.split('.') arch = dotsplit[-2] ver = dotsplit[0].split('release-')[-1].replace('-', '.') break elif 'redhat-release-8' in entry: ver = entry.split('-')[2] arch = entry.split('.')[-2] break else: return None major = ver.split('.', 1)[0] return {'name': 'rhel-{0}-{1}'.format(ver, arch), 'method': EXTRACT, 'category': 'el{0}'.format(major)} def scan_iso(filename): filesizes = {} filecontents = {} with libarchive.file_reader(filename) as reader: for ent in reader: if str(ent).endswith('TRANS.TBL'): continue eventlet.sleep(0) filesizes[str(ent)] = ent.size if str(ent) in READFILES: filecontents[str(ent)] = b'' for block in ent.get_blocks(): filecontents[str(ent)] += bytes(block) return filesizes, filecontents def fingerprint(filename): with open(filename, 'rb') as archive: header = archive.read(32768) archive.seek(32769) if archive.read(6) == b'CD001\x01': # ISO image isoinfo = scan_iso(filename) name = None for fun in globals(): if fun.startswith('check_'): name = globals()[fun](isoinfo) if name: return name, isoinfo[0] return None else: sum = hashlib.sha256(header) if sum.digest() in HEADERSUMS: archive.seek(32768) chunk = archive.read(32768) while chunk: sum.update(chunk) chunk = archive.read(32768) imginfo = HASHPRINTS.get(sum.hexdigest(), None) if imginfo: return imginfo, None def import_image(filename, callback, backend=False): identity = fingerprint(filename) if not identity: return -1 identity, imginfo = identity targpath = identity['name'] if identity.get('subname', None): targpath += '/' + identity['subname'] targpath = '/var/lib/confluent/distributions/' + targpath try: os.makedirs(targpath) except OSError as e: if e.errno != 17: raise filename = os.path.abspath(filename) os.chdir(targpath) if not backend: print('Importing OS to ' + targpath + ':') printit({'progress': 0.0}) if EXTRACT & identity['method']: extract_file(filename, callback=callback, imginfo=imginfo, extractlist=identity.get('extractlist', None)) if COPY & identity['method']: basename = identity.get('copyto', os.path.basename(filename)) targpath = os.path.join(targpath, basename) shutil.copyfile(filename, targpath) printit({'progress': 1.0}) sys.stdout.write('\n') def printit(info): sys.stdout.write(' \r{:.2f}%'.format(100 * info['progress'])) sys.stdout.flush() importing = {} class MediaImporter(object): def __init__(self, media): identity = fingerprint(media) identity, imginfo = identity if not identity: raise Exception('Unrecognized OS Media') importkey = ','.join((identity['name'], identity.get('subname', ''))) if importkey in importing: raise Exception('Media import already in progress for this media') importing[importkey] = self self.osname = identity['name'] self.oscategory = identity.get('category', None) targpath = identity['name'] if identity.get('subname', None): targpath += '/' + identity['subname'] self.targpath = '/var/lib/confluent/distributions/' + targpath try: os.makedirs(self.targpath) except OSError as e: if e.errno != 17: raise self.filename = os.path.abspath(media) self.importer = eventlet.spawn(self.importmedia) self.profiles = [] def importmedia(self): wkr = subprocess.Popen( [sys.executable, __file__, self.filename, '-b'], stdin=subprocess.DEVNULL, stdout=subprocess.PIPE) currline = b'' while wkr.poll() is None: currline += wkr.stdout.read(1) if b'\r' in currline: currline = b'' a = wkr.stdout.read(1) while a: currline += a if b'\r' in currline: val = currline.split(b'%')[0] if val: self.percent = float(val) currline = b'' a = wkr.stdout.read(1) bootupdates = [] if self.oscategory: defprofile = '/opt/confluent/lib/osdeploy/{0}'.format( self.oscategory) osd, osversion, arch = self.osname.split('-') for prof in os.listdir('{0}/profiles'.format(defprofile)): srcname = '{0}/profiles/{1}'.format(defprofile, prof) profname = '{0}-{1}'.format(self.osname, prof) dirname = '/var/lib/confluent/public/os/{0}'.format(profname) shutil.copytree(srcname, dirname) profdata = None try: os.makedirs('{0}/boot/initramfs'.format(dirname)) except OSError as e: if e.errno != 17: raise with open('{0}/profile.yaml'.format(dirname)) as yin: profdata = yin.read() profdata = profdata.replace('%%DISTRO%%', osd) profdata = profdata.replace('%%VERSION%%', osversion) profdata = profdata.replace('%%ARCH%%', arch) profdata = profdata.replace('%%PROFILE%%', profname) if profdata: with open('{0}/profile.yaml'.format(dirname), 'w') as yout: yout.write(profdata) for initrd in os.listdir('{0}/initramfs'.format(defprofile)): fullpath = '{0}/initramfs/{1}'.format(defprofile, initrd) os.symlink(fullpath, '{0}/boot/initramfs/{1}'.format(dirname, initrd)) os.symlink( '/var/lib/confluent/public/site/initramfs.cpio', '{0}/boot/initramfs/site.cpio'.format(dirname)) os.symlink(self.targpath, '{0}/distribution'.format(dirname)) subprocess.check_call( ['sh', '{0}/initprofile.sh'.format(dirname), self.targpath, dirname]) bootupdates.append(eventlet.spawn(update_boot, dirname)) self.profiles.append(profname) for upd in bootupdates: upd.wait() if __name__ == '__main__': if len(sys.argv) > 2: sys.exit(import_image(sys.argv[1], callback=printit, backend=True)) else: sys.exit(import_image(sys.argv[1], callback=printit))