From 705f7a9b290ff98f50b4431b9e1cc7076abfeaa1 Mon Sep 17 00:00:00 2001 From: Joshua Boniface Date: Thu, 14 Jun 2018 11:57:36 -0400 Subject: [PATCH] Move all the stuff from pvcf.py into pvc.py to keep everything in a single "binary" --- pvc.py | 509 +++++++++++++++++++++++++++++++++++++++++++++++--------- pvcf.py | 373 ----------------------------------------- 2 files changed, 432 insertions(+), 450 deletions(-) delete mode 100644 pvcf.py diff --git a/pvc.py b/pvc.py index b67d2431..91f54a78 100755 --- a/pvc.py +++ b/pvc.py @@ -20,10 +20,361 @@ # ############################################################################### -import kazoo.client, os, socket, time, click, lxml.objectify, pvcf, ansiiprint, configparser +import kazoo.client, os, socket, time, click, lxml.objectify, ansiiprint, configparser -myhostname = socket.gethostname() -zk_host = '' +############################################################################### +# Supplemental functions +############################################################################### + +# +# Validate a UUID +# +def validateUUID(dom_uuid): + try: + uuid.UUID(dom_uuid) + return True + except: + return False + + +# +# Connect and disconnect from Zookeeper +# +def startZKConnection(zk_host): + zk = kazoo.client.KazooClient(hosts=zk_host) + zk.start() + return zk + +def stopZKConnection(zk): + zk.stop() + zk.close() + return 0 + + +# +# XML information parsing functions +# + +# Get the main details for a VM object from XML +def getDomainMainDetails(parsed_xml): + # Get the information we want from it + duuid = str(parsed_xml.uuid) + dname = str(parsed_xml.name) + dmemory = str(parsed_xml.memory) + dmemory_unit = str(parsed_xml.memory.attrib['unit']) + if dmemory_unit == 'KiB': + dmemory = str(int(dmemory) * 1024) + elif dmemory_unit == 'GiB': + dmemory = str(int(dmemory) / 1024) + dvcpu = str(parsed_xml.vcpu) + try: + dvcputopo = '{}/{}/{}'.format(parsed_xml.cpu.topology.attrib['sockets'], parsed_xml.cpu.topology.attrib['cores'], parsed_xml.cpu.topology.attrib['threads']) + except: + dvcputopo = 'N/A' + + return duuid, dname, dmemory, dvcpu, dvcputopo + +# Get long-format details +def getDomainExtraDetails(parsed_xml): + dtype = parsed_xml.os.type + darch = parsed_xml.os.type.attrib['arch'] + dmachine = parsed_xml.os.type.attrib['machine'] + dconsole = parsed_xml.devices.console.attrib['type'] + demulator = parsed_xml.devices.emulator + + return dtype, darch, dmachine, dconsole, demulator + +# Get CPU features +def getDomainCPUFeatures(parsed_xml): + dfeatures = [] + for feature in parsed_xml.features.getchildren(): + dfeatures.append(feature.tag) + + return dfeatures + +# Get disk devices +def getDomainDisks(parsed_xml): + ddisks = [] + for device in parsed_xml.devices.getchildren(): + if device.tag == 'disk': + disk_attrib = device.source.attrib + disk_target = device.target.attrib + disk_type = device.attrib['type'] + if disk_type == 'network': + disk_obj = { 'type': disk_attrib.get('protocol'), 'name': disk_attrib.get('name'), 'dev': disk_target.get('dev'), 'bus': disk_target.get('bus') } + elif disk_type == 'file': + disk_obj = { 'type': 'file', 'name': disk_attrib.get('file'), 'dev': disk_target.get('dev'), 'bus': disk_target.get('bus') } + else: + disk_obj = {} + ddisks.append(disk_obj) + + return ddisks + +# Get network devices +def getDomainNetworks(parsed_xml): + dnets = [] + for device in parsed_xml.devices.getchildren(): + if device.tag == 'interface': + net_type = device.attrib['type'] + net_mac = device.mac.attrib['address'] + net_bridge = device.source.attrib[net_type] + net_model = device.model.attrib['type'] + net_obj = { 'type': net_type, 'mac': net_mac, 'source': net_bridge, 'model': net_model } + dnets.append(net_obj) + + return dnets + +# Get controller devices +def getDomainControllers(parsed_xml): + dcontrollers = [] + for device in parsed_xml.devices.getchildren(): + if device.tag == 'controller': + controller_type = device.attrib['type'] + try: + controller_model = device.attrib['model'] + except KeyError: + controller_model = 'none' + controller_obj = { 'type': controller_type, 'model': controller_model } + dcontrollers.append(controller_obj) + + return dcontrollers + +# Parse an XML object +def getDomainXML(zk, dom_uuid): + try: + xml = zk.get('/domains/%s/xml' % dom_uuid)[0].decode('ascii') + except: + return None + + # Parse XML using lxml.objectify + parsed_xml = lxml.objectify.fromstring(xml) + return parsed_xml + +# Root functions +def getInformationFromNode(zk, node_name, long_output): + node_daemon_state = zk.get('/nodes/{}/daemonstate'.format(node_name))[0].decode('ascii') + node_domain_state = zk.get('/nodes/{}/domainstate'.format(node_name))[0].decode('ascii') + node_cpu_count = zk.get('/nodes/{}/staticdata'.format(node_name))[0].decode('ascii').split()[0] + node_arch = zk.get('/nodes/{}/staticdata'.format(node_name))[0].decode('ascii').split()[1] + node_os = zk.get('/nodes/{}/staticdata'.format(node_name))[0].decode('ascii').split()[2] + node_kernel = zk.get('/nodes/{}/staticdata'.format(node_name))[0].decode('ascii').split()[3] + node_mem_used = zk.get('/nodes/{}/memused'.format(node_name))[0].decode('ascii') + node_mem_free = zk.get('/nodes/{}/memfree'.format(node_name))[0].decode('ascii') + node_mem_total = int(node_mem_used) + int(node_mem_free) + node_domains_count = zk.get('/nodes/{}/domainscount'.format(node_name))[0].decode('ascii') + node_running_domains = zk.get('/nodes/{}/runningdomains'.format(node_name))[0].decode('ascii').split() + node_mem_allocated = 0 + for domain in node_running_domains: + parsed_xml = getDomainXML(zk, domain) + duuid, dname, dmemory, dvcpu, dvcputopo = getDomainMainDetails(parsed_xml) + node_mem_allocated += int(dmemory) + + if node_daemon_state == 'run': + daemon_state_colour = ansiiprint.green() + elif node_daemon_state == 'stop': + daemon_state_colour = ansiiprint.red() + elif node_daemon_state == 'init': + daemon_state_colour = ansiiprint.yellow() + elif node_daemon_state == 'dead': + daemon_state_colour = ansiiprint.red() + ansiiprint.bold() + else: + daemon_state_colour = ansiiprint.blue() + + if node_domain_state == 'ready': + domain_state_colour = ansiiprint.green() + else: + domain_state_colour = ansiiprint.blue() + + # Format a nice output; do this line-by-line then concat the elements at the end + ainformation = [] + ainformation.append('{}Hypervisor Node information:{}'.format(ansiiprint.bold(), ansiiprint.end())) + ainformation.append('') + # Basic information + ainformation.append('{}Name:{} {}'.format(ansiiprint.purple(), ansiiprint.end(), node_name)) + ainformation.append('{}Daemon State:{} {}{}{}'.format(ansiiprint.purple(), ansiiprint.end(), daemon_state_colour, node_daemon_state, ansiiprint.end())) + ainformation.append('{}Domain State:{} {}{}{}'.format(ansiiprint.purple(), ansiiprint.end(), domain_state_colour, node_domain_state, ansiiprint.end())) + ainformation.append('{}Active Domain Count:{} {}'.format(ansiiprint.purple(), ansiiprint.end(), node_domains_count)) + if long_output == True: + ainformation.append('') + ainformation.append('{}Architecture:{} {}'.format(ansiiprint.purple(), ansiiprint.end(), node_arch)) + ainformation.append('{}Operating System:{} {}'.format(ansiiprint.purple(), ansiiprint.end(), node_os)) + ainformation.append('{}Kernel Version:{} {}'.format(ansiiprint.purple(), ansiiprint.end(), node_kernel)) + ainformation.append('') + ainformation.append('{}CPUs:{} {}'.format(ansiiprint.purple(), ansiiprint.end(), node_cpu_count)) + ainformation.append('{}Total RAM (MiB):{} {}'.format(ansiiprint.purple(), ansiiprint.end(), node_mem_total)) + ainformation.append('{}Used RAM (MiB):{} {}'.format(ansiiprint.purple(), ansiiprint.end(), node_mem_used)) + ainformation.append('{}Free RAM (MiB):{} {}'.format(ansiiprint.purple(), ansiiprint.end(), node_mem_free)) + ainformation.append('{}Allocated RAM (MiB):{} {}'.format(ansiiprint.purple(), ansiiprint.end(), node_mem_allocated)) + + # Join it all together + information = '\n'.join(ainformation) + return information + + +def getInformationFromXML(zk, uuid, long_output): + # Obtain the contents of the XML from Zookeeper + try: + dstate = zk.get('/domains/{}/state'.format(uuid))[0].decode('ascii') + dhypervisor = zk.get('/domains/{}/hypervisor'.format(uuid))[0].decode('ascii') + dlasthypervisor = zk.get('/domains/{}/lasthypervisor'.format(uuid))[0].decode('ascii') + except: + return None + + if dlasthypervisor == '': + dlasthypervisor = 'N/A' + + parsed_xml = getDomainXML(zk, uuid) + + duuid, dname, dmemory, dvcpu, dvcputopo = getDomainMainDetails(parsed_xml) + if long_output == True: + dtype, darch, dmachine, dconsole, demulator = getDomainExtraDetails(parsed_xml) + dfeatures = getDomainCPUFeatures(parsed_xml) + ddisks = getDomainDisks(parsed_xml) + dnets = getDomainNetworks(parsed_xml) + dcontrollers = getDomainControllers(parsed_xml) + + # Format a nice output; do this line-by-line then concat the elements at the end + ainformation = [] + ainformation.append('{}Virtual machine information:{}'.format(ansiiprint.bold(), ansiiprint.end())) + ainformation.append('') + # Basic information + ainformation.append('{}UUID:{} {}'.format(ansiiprint.purple(), ansiiprint.end(), duuid)) + ainformation.append('{}Name:{} {}'.format(ansiiprint.purple(), ansiiprint.end(), dname)) + ainformation.append('{}Memory (MiB):{} {}'.format(ansiiprint.purple(), ansiiprint.end(), dmemory)) + ainformation.append('{}vCPUs:{} {}'.format(ansiiprint.purple(), ansiiprint.end(), dvcpu)) + ainformation.append('{}Topology (S/C/T):{} {}'.format(ansiiprint.purple(), ansiiprint.end(), dvcputopo)) + + if long_output == True: + # Virtualization information + ainformation.append('') + ainformation.append('{}Emulator:{} {}'.format(ansiiprint.purple(), ansiiprint.end(), demulator)) + ainformation.append('{}Type:{} {}'.format(ansiiprint.purple(), ansiiprint.end(), dtype)) + ainformation.append('{}Arch:{} {}'.format(ansiiprint.purple(), ansiiprint.end(), darch)) + ainformation.append('{}Machine:{} {}'.format(ansiiprint.purple(), ansiiprint.end(), dmachine)) + ainformation.append('{}Features:{} {}'.format(ansiiprint.purple(), ansiiprint.end(), ' '.join(dfeatures))) + + # PVC cluster information + ainformation.append('') + dstate_colour = { + 'start': ansiiprint.green(), + 'restart': ansiiprint.yellow(), + 'shutdown': ansiiprint.yellow(), + 'stop': ansiiprint.red(), + 'migrate': ansiiprint.blue(), + 'unmigrate': ansiiprint.blue() + } + ainformation.append('{}State:{} {}{}{}'.format(ansiiprint.purple(), ansiiprint.end(), dstate_colour[dstate], dstate, ansiiprint.end())) + ainformation.append('{}Active Hypervisor:{} {}'.format(ansiiprint.purple(), ansiiprint.end(), dhypervisor)) + ainformation.append('{}Last Hypervisor:{} {}'.format(ansiiprint.purple(), ansiiprint.end(), dlasthypervisor)) + + if long_output == True: + # Disk list + ainformation.append('') + name_length = 0 + for disk in ddisks: + _name_length = len(disk['name']) + 1 + if _name_length > name_length: + name_length = _name_length + ainformation.append('{0}Disks:{1} {2}ID Type {3: <{width}} Dev Bus{4}'.format(ansiiprint.purple(), ansiiprint.end(), ansiiprint.bold(), 'Name', ansiiprint.end(), width=name_length)) + for disk in ddisks: + ainformation.append(' {0: <3} {1: <5} {2: <{width}} {3: <4} {4: <5}'.format(ddisks.index(disk), disk['type'], disk['name'], disk['dev'], disk['bus'], width=name_length)) + # Network list + ainformation.append('') + ainformation.append('{}Interfaces:{} {}ID Type Source Model MAC{}'.format(ansiiprint.purple(), ansiiprint.end(), ansiiprint.bold(), ansiiprint.end())) + for net in dnets: + ainformation.append(' {0: <3} {1: <8} {2: <10} {3: <8} {4}'.format(dnets.index(net), net['type'], net['source'], net['model'], net['mac'])) + # Controller list + ainformation.append('') + ainformation.append('{}Controllers:{} {}ID Type Model{}'.format(ansiiprint.purple(), ansiiprint.end(), ansiiprint.bold(), ansiiprint.end())) + for controller in dcontrollers: + ainformation.append(' {0: <3} {1: <14} {2: <8}'.format(dcontrollers.index(controller), controller['type'], controller['model'])) + + # Join it all together + information = '\n'.join(ainformation) + return information + + +# +# Cluster search functions +# +def getClusterDomainList(zk): + # Get a list of UUIDs by listing the children of /domains + uuid_list = zk.get_children('/domains') + name_list = [] + # For each UUID, get the corresponding name from the data + for uuid in uuid_list: + name_list.append(zk.get('/domains/%s' % uuid)[0].decode('ascii')) + return uuid_list, name_list + +def searchClusterByUUID(zk, uuid): + try: + # Get the lists + uuid_list, name_list = getClusterDomainList(zk) + # We're looking for UUID, so find that element ID + index = uuid_list.index(uuid) + # Get the name_list element at that index + name = name_list[index] + except ValueError: + # We didn't find anything + return None + + return name + +def searchClusterByName(zk, name): + try: + # Get the lists + uuid_list, name_list = getClusterDomainList(zk) + # We're looking for name, so find that element ID + index = name_list.index(name) + # Get the uuid_list element at that index + uuid = uuid_list[index] + except ValueError: + # We didn't find anything + return None + + return uuid + + +# +# Allow mutually exclusive options in Click +# +class MutuallyExclusiveOption(click.Option): + def __init__(self, *args, **kwargs): + meargs = kwargs.pop('mutually_exclusive', []) + _me_arg = [] + _me_func = [] + + for arg in meargs: + _me_arg.append(arg['argument']) + _me_func.append(arg['function']) + + self.me_arg = set(_me_arg) + self.me_func = set(_me_func) + + help = kwargs.get('help', '') + if self.me_func: + ex_str = ', '.join(self.me_arg) + kwargs['help'] = help + ( + ' Mutually exclusive with `' + ex_str + '`.' + ) + + super(MutuallyExclusiveOption, self).__init__(*args, **kwargs) + + def handle_parse_result(self, ctx, opts, args): + if self.me_func.intersection(opts) and self.name in opts: + raise click.UsageError( + "Illegal usage: `{}` is mutually exclusive with " + "arguments `{}`.".format( + self.opts[-1], + ', '.join(self.me_arg) + ) + ) + + return super(MutuallyExclusiveOption, self).handle_parse_result( + ctx, + opts, + args + ) ######################## ######################## @@ -33,6 +384,9 @@ zk_host = '' ######################## ######################## +myhostname = socket.gethostname() +zk_host = '' + CONTEXT_SETTINGS = dict(help_option_names=['-h', '--help'], max_content_width=120) ############################################################################### @@ -64,7 +418,7 @@ def flush_host(node_name): """ # Open a Zookeeper connection - zk = pvcf.startZKConnection(zk_host) + zk = startZKConnection(zk_host) # Verify node is valid try: @@ -81,7 +435,7 @@ def flush_host(node_name): results = transaction.commit() # Close the Zookeeper connection - pvcf.stopZKConnection(zk) + stopZKConnection(zk) ############################################################################### @@ -102,7 +456,7 @@ def ready_host(node_name): """ # Open a Zookeeper connection - zk = pvcf.startZKConnection(zk_host) + zk = startZKConnection(zk_host) # Verify node is valid try: @@ -119,7 +473,7 @@ def ready_host(node_name): results = transaction.commit() # Close the Zookeeper connection - pvcf.stopZKConnection(zk) + stopZKConnection(zk) ############################################################################### @@ -140,7 +494,7 @@ def node_info(node_name, long_output): """ # Open a Zookeeper connection - zk = pvcf.startZKConnection(zk_host) + zk = startZKConnection(zk_host) # Verify node is valid try: @@ -150,7 +504,7 @@ def node_info(node_name, long_output): exit(1) # Get information about node in a pretty format - information = pvcf.getInformationFromNode(zk, node_name, long_output) + information = getInformationFromNode(zk, node_name, long_output) if information == None: click.echo('ERROR: Could not find a domain matching that name or UUID.') @@ -166,7 +520,7 @@ def node_info(node_name, long_output): _vm_list(node_name) # Close the Zookeeper connection - pvcf.stopZKConnection(zk) + stopZKConnection(zk) ############################################################################### @@ -179,7 +533,7 @@ def node_list(): """ # Open a Zookeeper connection - zk = pvcf.startZKConnection(zk_host) + zk = startZKConnection(zk_host) node_list = zk.get_children('/nodes') node_list_output = [] @@ -206,8 +560,8 @@ def node_list(): node_running_domains[node_name] = zk.get('/nodes/{}/runningdomains'.format(node_name))[0].decode('ascii').split() node_mem_allocated[node_name] = 0 for domain in node_running_domains[node_name]: - parsed_xml = pvcf.getDomainXML(zk, domain) - duuid, dname, dmemory, dvcpu, dvcputopo = pvcf.getDomainMainDetails(parsed_xml) + parsed_xml = getDomainXML(zk, domain) + duuid, dname, dmemory, dvcpu, dvcputopo = getDomainMainDetails(parsed_xml) node_mem_allocated[node_name] += int(dmemory) # Determine optimal column widths @@ -287,7 +641,7 @@ RAM (MiB): {node_mem_total: <6} {node_mem_used: <6} {node_mem_free: <6} {node_me click.echo('\n'.join(sorted(node_list_output))) # Close the Zookeeper connection - pvcf.stopZKConnection(zk) + stopZKConnection(zk) ############################################################################### @@ -334,7 +688,7 @@ def define_vm(xml_config_file, target_hypervisor): click.echo('Adding new VM with Name "{}" and UUID "{}" to database.'.format(dom_name, dom_uuid)) # Open a Zookeeper connection - zk = pvcf.startZKConnection(zk_host) + zk = startZKConnection(zk_host) # Add the new domain to Zookeeper transaction = zk.transaction() @@ -346,7 +700,7 @@ def define_vm(xml_config_file, target_hypervisor): results = transaction.commit() # Close the Zookeeper connection - pvcf.stopZKConnection(zk) + stopZKConnection(zk) ############################################################################### @@ -355,13 +709,13 @@ def define_vm(xml_config_file, target_hypervisor): @click.command(name='undefine', short_help='Undefine and stop a virtual machine.') @click.option( '-n', '--name', 'dom_name', - cls=pvcf.MutuallyExclusiveOption, + cls=MutuallyExclusiveOption, mutually_exclusive=[{ 'function': 'dom_uuid', 'argument': '--uuid' }], help='Search for this human-readable name.' ) @click.option( '-u', '--uuid', 'dom_uuid', - cls=pvcf.MutuallyExclusiveOption, + cls=MutuallyExclusiveOption, mutually_exclusive=[{ 'function': 'dom_name', 'argument': '--name' }], help='Search for this UUID.' ) @@ -376,14 +730,14 @@ def undefine_vm(dom_name, dom_uuid): return # Open a Zookeeper connection - zk = pvcf.startZKConnection(zk_host) + zk = startZKConnection(zk_host) # If the --name value was passed, get the UUID if dom_name != None: - dom_uuid = pvcf.searchClusterByName(zk, dom_name) + dom_uuid = searchClusterByName(zk, dom_name) # Verify we got a result or abort - if not pvcf.validateUUID(dom_uuid): + if not validateUUID(dom_uuid): if dom_name != None: message_name = dom_name else: @@ -414,7 +768,7 @@ def undefine_vm(dom_name, dom_uuid): transaction.commit() # Close the Zookeeper connection - pvcf.stopZKConnection(zk) + stopZKConnection(zk) ############################################################################### @@ -423,13 +777,13 @@ def undefine_vm(dom_name, dom_uuid): @click.command(name='start', short_help='Start up a defined virtual machine.') @click.option( '-n', '--name', 'dom_name', - cls=pvcf.MutuallyExclusiveOption, + cls=MutuallyExclusiveOption, mutually_exclusive=[{ 'function': 'dom_uuid', 'argument': '--uuid' }], help='Search for this human-readable name.' ) @click.option( '-u', '--uuid', 'dom_uuid', - cls=pvcf.MutuallyExclusiveOption, + cls=MutuallyExclusiveOption, mutually_exclusive=[{ 'function': 'dom_name', 'argument': '--name' }], help='Search for this UUID.' ) @@ -444,14 +798,14 @@ def start_vm(dom_name, dom_uuid): return # Open a Zookeeper connection - zk = pvcf.startZKConnection(zk_host) + zk = startZKConnection(zk_host) # If the --name value was passed, get the UUID if dom_name != None: - dom_uuid = pvcf.searchClusterByName(zk, dom_name) + dom_uuid = searchClusterByName(zk, dom_name) # Verify we got a result or abort - if not pvcf.validateUUID(dom_uuid): + if not validateUUID(dom_uuid): if dom_name != None: message_name = dom_name else: @@ -464,7 +818,7 @@ def start_vm(dom_name, dom_uuid): zk.set('/domains/%s/state' % dom_uuid, 'start'.encode('ascii')) # Close the Zookeeper connection - pvcf.stopZKConnection(zk) + stopZKConnection(zk) ############################################################################### @@ -473,13 +827,13 @@ def start_vm(dom_name, dom_uuid): @click.command(name='restart', short_help='Restart virtual machine.') @click.option( '-n', '--name', 'dom_name', - cls=pvcf.MutuallyExclusiveOption, + cls=MutuallyExclusiveOption, mutually_exclusive=[{ 'function': 'dom_uuid', 'argument': '--uuid' }], help='Search for this human-readable name.' ) @click.option( '-u', '--uuid', 'dom_uuid', - cls=pvcf.MutuallyExclusiveOption, + cls=MutuallyExclusiveOption, mutually_exclusive=[{ 'function': 'dom_name', 'argument': '--name' }], help='Search for this UUID.' ) @@ -494,14 +848,14 @@ def restart_vm(dom_name, dom_uuid): return # Open a Zookeeper connection - zk = pvcf.startZKConnection(zk_host) + zk = startZKConnection(zk_host) # If the --name value was passed, get the UUID if dom_name != None: - dom_uuid = pvcf.searchClusterByName(zk, dom_name) + dom_uuid = searchClusterByName(zk, dom_name) # Verify we got a result or abort - if not pvcf.validateUUID(dom_uuid): + if not validateUUID(dom_uuid): if dom_name != None: message_name = dom_name else: @@ -514,7 +868,7 @@ def restart_vm(dom_name, dom_uuid): zk.set('/domains/%s/state' % dom_uuid, 'restart'.encode('ascii')) # Close the Zookeeper connection - pvcf.stopZKConnection(zk) + stopZKConnection(zk) ############################################################################### @@ -523,13 +877,13 @@ def restart_vm(dom_name, dom_uuid): @click.command(name='shutdown', short_help='Gracefully shut down a running virtual machine.') @click.option( '-n', '--name', 'dom_name', - cls=pvcf.MutuallyExclusiveOption, + cls=MutuallyExclusiveOption, mutually_exclusive=[{ 'function': 'dom_uuid', 'argument': '--uuid' }], help='Search for this human-readable name.' ) @click.option( '-u', '--uuid', 'dom_uuid', - cls=pvcf.MutuallyExclusiveOption, + cls=MutuallyExclusiveOption, mutually_exclusive=[{ 'function': 'dom_name', 'argument': '--name' }], help='Search for this UUID.' ) @@ -544,14 +898,14 @@ def shutdown_vm(dom_name, dom_uuid): return # Open a Zookeeper connection - zk = pvcf.startZKConnection(zk_host) + zk = startZKConnection(zk_host) # If the --name value was passed, get the UUID if dom_name != None: - dom_uuid = pvcf.searchClusterByName(zk, dom_name) + dom_uuid = searchClusterByName(zk, dom_name) # Verify we got a result or abort - if not pvcf.validateUUID(dom_uuid): + if not validateUUID(dom_uuid): if dom_name != None: message_name = dom_name else: @@ -564,7 +918,7 @@ def shutdown_vm(dom_name, dom_uuid): zk.set('/domains/%s/state' % dom_uuid, 'shutdown'.encode('ascii')) # Close the Zookeeper connection - pvcf.stopZKConnection(zk) + stopZKConnection(zk) ############################################################################### @@ -573,13 +927,13 @@ def shutdown_vm(dom_name, dom_uuid): @click.command(name='stop', short_help='Forcibly halt a running virtual machine.') @click.option( '-n', '--name', 'dom_name', - cls=pvcf.MutuallyExclusiveOption, + cls=MutuallyExclusiveOption, mutually_exclusive=[{ 'function': 'dom_uuid', 'argument': '--uuid' }], help='Search for this human-readable name.' ) @click.option( '-u', '--uuid', 'dom_uuid', - cls=pvcf.MutuallyExclusiveOption, + cls=MutuallyExclusiveOption, mutually_exclusive=[{ 'function': 'dom_name', 'argument': '--name' }], help='Search for this UUID.' ) @@ -594,14 +948,14 @@ def stop_vm(dom_name, dom_uuid): return # Open a Zookeeper connection - zk = pvcf.startZKConnection(zk_host) + zk = startZKConnection(zk_host) # If the --name value was passed, get the UUID if dom_name != None: - dom_uuid = pvcf.searchClusterByName(zk, dom_name) + dom_uuid = searchClusterByName(zk, dom_name) # Verify we got a result or abort - if not pvcf.validateUUID(dom_uuid): + if not validateUUID(dom_uuid): if dom_name != None: message_name = dom_name else: @@ -614,7 +968,7 @@ def stop_vm(dom_name, dom_uuid): zk.set('/domains/%s/state' % dom_uuid, 'stop'.encode('ascii')) # Close the Zookeeper connection - pvcf.stopZKConnection(zk) + stopZKConnection(zk) ############################################################################### @@ -623,13 +977,13 @@ def stop_vm(dom_name, dom_uuid): @click.command(name='move', short_help='Permanently move a virtual machine to another node.') @click.option( '-n', '--name', 'dom_name', - cls=pvcf.MutuallyExclusiveOption, + cls=MutuallyExclusiveOption, mutually_exclusive=[{ 'function': 'dom_uuid', 'argument': '--uuid' }], help='Search for this human-readable name.' ) @click.option( '-u', '--uuid', 'dom_uuid', - cls=pvcf.MutuallyExclusiveOption, + cls=MutuallyExclusiveOption, mutually_exclusive=[{ 'function': 'dom_name', 'argument': '--name' }], help='Search for this UUID.' ) @@ -648,14 +1002,14 @@ def move_vm(dom_name, dom_uuid, target_hypervisor): return # Open a Zookeeper connection - zk = pvcf.startZKConnection(zk_host) + zk = startZKConnection(zk_host) # If the --name value was passed, get the UUID if dom_name != None: - dom_uuid = pvcf.searchClusterByName(zk, dom_name) + dom_uuid = searchClusterByName(zk, dom_name) # Verify we got a result or abort - if not pvcf.validateUUID(dom_uuid): + if not validateUUID(dom_uuid): if dom_name != None: message_name = dom_name else: @@ -699,7 +1053,7 @@ def move_vm(dom_name, dom_uuid, target_hypervisor): transaction.commit() # Close the Zookeeper connection - pvcf.stopZKConnection(zk) + stopZKConnection(zk) ############################################################################### @@ -708,13 +1062,13 @@ def move_vm(dom_name, dom_uuid, target_hypervisor): @click.command(name='migrate', short_help='Migrate a virtual machine to another node.') @click.option( '-n', '--name', 'dom_name', - cls=pvcf.MutuallyExclusiveOption, + cls=MutuallyExclusiveOption, mutually_exclusive=[{ 'function': 'dom_uuid', 'argument': '--uuid' }], help='Search for this human-readable name.' ) @click.option( '-u', '--uuid', 'dom_uuid', - cls=pvcf.MutuallyExclusiveOption, + cls=MutuallyExclusiveOption, mutually_exclusive=[{ 'function': 'dom_name', 'argument': '--name' }], help='Search for this UUID.' ) @@ -737,14 +1091,14 @@ def migrate_vm(dom_name, dom_uuid, target_hypervisor, force_migrate): return # Open a Zookeeper connection - zk = pvcf.startZKConnection(zk_host) + zk = startZKConnection(zk_host) # If the --name value was passed, get the UUID if dom_name != None: - dom_uuid = pvcf.searchClusterByName(zk, dom_name) + dom_uuid = searchClusterByName(zk, dom_name) # Verify we got a result or abort - if not pvcf.validateUUID(dom_uuid): + if not validateUUID(dom_uuid): if dom_name != None: message_name = dom_name else: @@ -789,7 +1143,7 @@ def migrate_vm(dom_name, dom_uuid, target_hypervisor, force_migrate): transaction.commit() # Close the Zookeeper connection - pvcf.stopZKConnection(zk) + stopZKConnection(zk) ############################################################################### @@ -798,13 +1152,13 @@ def migrate_vm(dom_name, dom_uuid, target_hypervisor, force_migrate): @click.command(name='unmigrate', short_help='Restore a migrated virtual machine to its original node.') @click.option( '-n', '--name', 'dom_name', - cls=pvcf.MutuallyExclusiveOption, + cls=MutuallyExclusiveOption, mutually_exclusive=[{ 'function': 'dom_uuid', 'argument': '--uuid' }], help='Search for this human-readable name.' ) @click.option( '-u', '--uuid', 'dom_uuid', - cls=pvcf.MutuallyExclusiveOption, + cls=MutuallyExclusiveOption, mutually_exclusive=[{ 'function': 'dom_name', 'argument': '--name' }], help='Search for this UUID.' ) @@ -819,14 +1173,14 @@ def unmigrate_vm(dom_name, dom_uuid): return # Open a Zookeeper connection - zk = pvcf.startZKConnection(zk_host) + zk = startZKConnection(zk_host) # If the --name value was passed, get the UUID if dom_name != None: - dom_uuid = pvcf.searchClusterByName(zk, dom_name) + dom_uuid = searchClusterByName(zk, dom_name) # Verify we got a result or abort - if not pvcf.validateUUID(dom_uuid): + if not validateUUID(dom_uuid): if dom_name != None: message_name = dom_name else: @@ -848,7 +1202,7 @@ def unmigrate_vm(dom_name, dom_uuid): transaction.commit() # Close the Zookeeper connection - pvcf.stopZKConnection(zk) + stopZKConnection(zk) ############################################################################### @@ -857,13 +1211,13 @@ def unmigrate_vm(dom_name, dom_uuid): @click.command(name='info', short_help='Show details of a VM object') @click.option( '-n', '--name', 'dom_name', - cls=pvcf.MutuallyExclusiveOption, + cls=MutuallyExclusiveOption, mutually_exclusive=[{ 'function': 'dom_uuid', 'argument': '--uuid' }], help='Search for this human-readable name.' ) @click.option( '-u', '--uuid', 'dom_uuid', - cls=pvcf.MutuallyExclusiveOption, + cls=MutuallyExclusiveOption, mutually_exclusive=[{ 'function': 'dom_name', 'argument': '--name' }], help='Search for this UUID.' ) @@ -881,20 +1235,20 @@ def vm_info(dom_name, dom_uuid, long_output): click.echo("ERROR: You must specify either a `--name` or `--uuid` value.") return - zk = pvcf.startZKConnection(zk_host) + zk = startZKConnection(zk_host) if dom_name != None: - dom_uuid = pvcf.searchClusterByName(zk, dom_name) + dom_uuid = searchClusterByName(zk, dom_name) if dom_uuid != None: - dom_name = pvcf.searchClusterByUUID(zk, dom_uuid) + dom_name = searchClusterByUUID(zk, dom_uuid) - information = pvcf.getInformationFromXML(zk, dom_uuid, long_output) + information = getInformationFromXML(zk, dom_uuid, long_output) if information == None: click.echo('ERROR: Could not find a domain matching that name or UUID.') return click.echo(information) - pvcf.stopZKConnection(zk) + stopZKConnection(zk) ############################################################################### @@ -915,7 +1269,7 @@ def _vm_list(hypervisor): """ # Open a Zookeeper connection - zk = pvcf.startZKConnection(zk_host) + zk = startZKConnection(zk_host) vm_list_raw = zk.get_children('/domains') vm_list = [] @@ -948,8 +1302,8 @@ def _vm_list(hypervisor): else: vm_migrated[vm] = 'no' - vm_xml = pvcf.getDomainXML(zk, vm) - vm_uuid[vm], vm_name[vm], vm_memory[vm], vm_vcpu[vm], vm_vcputopo = pvcf.getDomainMainDetails(vm_xml) + vm_xml = getDomainXML(zk, vm) + vm_uuid[vm], vm_name[vm], vm_memory[vm], vm_vcpu[vm], vm_vcputopo = getDomainMainDetails(vm_xml) # Determine optimal column widths # Dynamic columns: node_name, hypervisor, migrated @@ -1034,7 +1388,7 @@ def _vm_list(hypervisor): click.echo('\n'.join(sorted(vm_list_output))) # Close the Zookeeper connection - pvcf.stopZKConnection(zk) + stopZKConnection(zk) ############################################################################### @@ -1052,7 +1406,7 @@ def init_cluster(): click.echo('Initializing a new cluster with Zookeeper address "{}".'.format(zk_host)) # Open a Zookeeper connection - zk = pvcf.startZKConnection(zk_host) + zk = startZKConnection(zk_host) # Destroy the existing data try: @@ -1068,7 +1422,7 @@ def init_cluster(): transaction.commit() # Close the Zookeeper connection - pvcf.stopZKConnection(zk) + stopZKConnection(zk) click.echo('Successfully initialized new cluster. Any running PVC daemons will need to be restarted.') @@ -1122,3 +1476,4 @@ def main(): if __name__ == '__main__': main() + diff --git a/pvcf.py b/pvcf.py deleted file mode 100644 index 5b5e1983..00000000 --- a/pvcf.py +++ /dev/null @@ -1,373 +0,0 @@ -#!/usr/bin/env python3 - -# pvcf.py - Supplemental functions for the PVC CLI client -# Part of the Parallel Virtual Cluster (PVC) system -# -# Copyright (C) 2018 Joshua M. Boniface -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU General Public License as published by -# the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU General Public License for more details. -# -# You should have received a copy of the GNU General Public License -# along with this program. If not, see . -# -############################################################################### - -import os, sys, libvirt, uuid, kazoo.client, lxml.objectify, click, ansiiprint - -# -# Validate a UUID -# -def validateUUID(dom_uuid): - try: - uuid.UUID(dom_uuid) - return True - except: - return False - - -# -# Connect and disconnect from Zookeeper -# -def startZKConnection(zk_host): - zk = kazoo.client.KazooClient(hosts=zk_host) - zk.start() - return zk - -def stopZKConnection(zk): - zk.stop() - zk.close() - return 0 - - -# -# XML information parsing functions -# - -# Get the main details for a VM object from XML -def getDomainMainDetails(parsed_xml): - # Get the information we want from it - duuid = str(parsed_xml.uuid) - dname = str(parsed_xml.name) - dmemory = str(parsed_xml.memory) - dmemory_unit = str(parsed_xml.memory.attrib['unit']) - if dmemory_unit == 'KiB': - dmemory = str(int(dmemory) * 1024) - elif dmemory_unit == 'GiB': - dmemory = str(int(dmemory) / 1024) - dvcpu = str(parsed_xml.vcpu) - try: - dvcputopo = '{}/{}/{}'.format(parsed_xml.cpu.topology.attrib['sockets'], parsed_xml.cpu.topology.attrib['cores'], parsed_xml.cpu.topology.attrib['threads']) - except: - dvcputopo = 'N/A' - - return duuid, dname, dmemory, dvcpu, dvcputopo - -# Get long-format details -def getDomainExtraDetails(parsed_xml): - dtype = parsed_xml.os.type - darch = parsed_xml.os.type.attrib['arch'] - dmachine = parsed_xml.os.type.attrib['machine'] - dconsole = parsed_xml.devices.console.attrib['type'] - demulator = parsed_xml.devices.emulator - - return dtype, darch, dmachine, dconsole, demulator - -# Get CPU features -def getDomainCPUFeatures(parsed_xml): - dfeatures = [] - for feature in parsed_xml.features.getchildren(): - dfeatures.append(feature.tag) - - return dfeatures - -# Get disk devices -def getDomainDisks(parsed_xml): - ddisks = [] - for device in parsed_xml.devices.getchildren(): - if device.tag == 'disk': - disk_attrib = device.source.attrib - disk_target = device.target.attrib - disk_type = device.attrib['type'] - if disk_type == 'network': - disk_obj = { 'type': disk_attrib.get('protocol'), 'name': disk_attrib.get('name'), 'dev': disk_target.get('dev'), 'bus': disk_target.get('bus') } - elif disk_type == 'file': - disk_obj = { 'type': 'file', 'name': disk_attrib.get('file'), 'dev': disk_target.get('dev'), 'bus': disk_target.get('bus') } - else: - disk_obj = {} - ddisks.append(disk_obj) - - return ddisks - -# Get network devices -def getDomainNetworks(parsed_xml): - dnets = [] - for device in parsed_xml.devices.getchildren(): - if device.tag == 'interface': - net_type = device.attrib['type'] - net_mac = device.mac.attrib['address'] - net_bridge = device.source.attrib[net_type] - net_model = device.model.attrib['type'] - net_obj = { 'type': net_type, 'mac': net_mac, 'source': net_bridge, 'model': net_model } - dnets.append(net_obj) - - return dnets - -# Get controller devices -def getDomainControllers(parsed_xml): - dcontrollers = [] - for device in parsed_xml.devices.getchildren(): - if device.tag == 'controller': - controller_type = device.attrib['type'] - try: - controller_model = device.attrib['model'] - except KeyError: - controller_model = 'none' - controller_obj = { 'type': controller_type, 'model': controller_model } - dcontrollers.append(controller_obj) - - return dcontrollers - -# Parse an XML object -def getDomainXML(zk, dom_uuid): - try: - xml = zk.get('/domains/%s/xml' % dom_uuid)[0].decode('ascii') - except: - return None - - # Parse XML using lxml.objectify - parsed_xml = lxml.objectify.fromstring(xml) - return parsed_xml - -# Root functions -def getInformationFromNode(zk, node_name, long_output): - node_daemon_state = zk.get('/nodes/{}/daemonstate'.format(node_name))[0].decode('ascii') - node_domain_state = zk.get('/nodes/{}/domainstate'.format(node_name))[0].decode('ascii') - node_cpu_count = zk.get('/nodes/{}/staticdata'.format(node_name))[0].decode('ascii').split()[0] - node_arch = zk.get('/nodes/{}/staticdata'.format(node_name))[0].decode('ascii').split()[1] - node_os = zk.get('/nodes/{}/staticdata'.format(node_name))[0].decode('ascii').split()[2] - node_kernel = zk.get('/nodes/{}/staticdata'.format(node_name))[0].decode('ascii').split()[3] - node_mem_used = zk.get('/nodes/{}/memused'.format(node_name))[0].decode('ascii') - node_mem_free = zk.get('/nodes/{}/memfree'.format(node_name))[0].decode('ascii') - node_mem_total = int(node_mem_used) + int(node_mem_free) - node_domains_count = zk.get('/nodes/{}/domainscount'.format(node_name))[0].decode('ascii') - node_running_domains = zk.get('/nodes/{}/runningdomains'.format(node_name))[0].decode('ascii').split() - node_mem_allocated = 0 - for domain in node_running_domains: - parsed_xml = getDomainXML(zk, domain) - duuid, dname, dmemory, dvcpu, dvcputopo = getDomainMainDetails(parsed_xml) - node_mem_allocated += int(dmemory) - - if node_daemon_state == 'run': - daemon_state_colour = ansiiprint.green() - elif node_daemon_state == 'stop': - daemon_state_colour = ansiiprint.red() - elif node_daemon_state == 'init': - daemon_state_colour = ansiiprint.yellow() - elif node_daemon_state == 'dead': - daemon_state_colour = ansiiprint.red() + ansiiprint.bold() - else: - daemon_state_colour = ansiiprint.blue() - - if node_domain_state == 'ready': - domain_state_colour = ansiiprint.green() - else: - domain_state_colour = ansiiprint.blue() - - # Format a nice output; do this line-by-line then concat the elements at the end - ainformation = [] - ainformation.append('{}Hypervisor Node information:{}'.format(ansiiprint.bold(), ansiiprint.end())) - ainformation.append('') - # Basic information - ainformation.append('{}Name:{} {}'.format(ansiiprint.purple(), ansiiprint.end(), node_name)) - ainformation.append('{}Daemon State:{} {}{}{}'.format(ansiiprint.purple(), ansiiprint.end(), daemon_state_colour, node_daemon_state, ansiiprint.end())) - ainformation.append('{}Domain State:{} {}{}{}'.format(ansiiprint.purple(), ansiiprint.end(), domain_state_colour, node_domain_state, ansiiprint.end())) - ainformation.append('{}Active Domain Count:{} {}'.format(ansiiprint.purple(), ansiiprint.end(), node_domains_count)) - if long_output == True: - ainformation.append('') - ainformation.append('{}Architecture:{} {}'.format(ansiiprint.purple(), ansiiprint.end(), node_arch)) - ainformation.append('{}Operating System:{} {}'.format(ansiiprint.purple(), ansiiprint.end(), node_os)) - ainformation.append('{}Kernel Version:{} {}'.format(ansiiprint.purple(), ansiiprint.end(), node_kernel)) - ainformation.append('') - ainformation.append('{}CPUs:{} {}'.format(ansiiprint.purple(), ansiiprint.end(), node_cpu_count)) - ainformation.append('{}Total RAM (MiB):{} {}'.format(ansiiprint.purple(), ansiiprint.end(), node_mem_total)) - ainformation.append('{}Used RAM (MiB):{} {}'.format(ansiiprint.purple(), ansiiprint.end(), node_mem_used)) - ainformation.append('{}Free RAM (MiB):{} {}'.format(ansiiprint.purple(), ansiiprint.end(), node_mem_free)) - ainformation.append('{}Allocated RAM (MiB):{} {}'.format(ansiiprint.purple(), ansiiprint.end(), node_mem_allocated)) - - # Join it all together - information = '\n'.join(ainformation) - return information - - -def getInformationFromXML(zk, uuid, long_output): - # Obtain the contents of the XML from Zookeeper - try: - dstate = zk.get('/domains/{}/state'.format(uuid))[0].decode('ascii') - dhypervisor = zk.get('/domains/{}/hypervisor'.format(uuid))[0].decode('ascii') - dlasthypervisor = zk.get('/domains/{}/lasthypervisor'.format(uuid))[0].decode('ascii') - except: - return None - - if dlasthypervisor == '': - dlasthypervisor = 'N/A' - - parsed_xml = getDomainXML(zk, uuid) - - duuid, dname, dmemory, dvcpu, dvcputopo = getDomainMainDetails(parsed_xml) - if long_output == True: - dtype, darch, dmachine, dconsole, demulator = getDomainExtraDetails(parsed_xml) - dfeatures = getDomainCPUFeatures(parsed_xml) - ddisks = getDomainDisks(parsed_xml) - dnets = getDomainNetworks(parsed_xml) - dcontrollers = getDomainControllers(parsed_xml) - - # Format a nice output; do this line-by-line then concat the elements at the end - ainformation = [] - ainformation.append('{}Virtual machine information:{}'.format(ansiiprint.bold(), ansiiprint.end())) - ainformation.append('') - # Basic information - ainformation.append('{}UUID:{} {}'.format(ansiiprint.purple(), ansiiprint.end(), duuid)) - ainformation.append('{}Name:{} {}'.format(ansiiprint.purple(), ansiiprint.end(), dname)) - ainformation.append('{}Memory (MiB):{} {}'.format(ansiiprint.purple(), ansiiprint.end(), dmemory)) - ainformation.append('{}vCPUs:{} {}'.format(ansiiprint.purple(), ansiiprint.end(), dvcpu)) - ainformation.append('{}Topology (S/C/T):{} {}'.format(ansiiprint.purple(), ansiiprint.end(), dvcputopo)) - - if long_output == True: - # Virtualization information - ainformation.append('') - ainformation.append('{}Emulator:{} {}'.format(ansiiprint.purple(), ansiiprint.end(), demulator)) - ainformation.append('{}Type:{} {}'.format(ansiiprint.purple(), ansiiprint.end(), dtype)) - ainformation.append('{}Arch:{} {}'.format(ansiiprint.purple(), ansiiprint.end(), darch)) - ainformation.append('{}Machine:{} {}'.format(ansiiprint.purple(), ansiiprint.end(), dmachine)) - ainformation.append('{}Features:{} {}'.format(ansiiprint.purple(), ansiiprint.end(), ' '.join(dfeatures))) - - # PVC cluster information - ainformation.append('') - dstate_colour = { - 'start': ansiiprint.green(), - 'restart': ansiiprint.yellow(), - 'shutdown': ansiiprint.yellow(), - 'stop': ansiiprint.red(), - 'migrate': ansiiprint.blue(), - 'unmigrate': ansiiprint.blue() - } - ainformation.append('{}State:{} {}{}{}'.format(ansiiprint.purple(), ansiiprint.end(), dstate_colour[dstate], dstate, ansiiprint.end())) - ainformation.append('{}Active Hypervisor:{} {}'.format(ansiiprint.purple(), ansiiprint.end(), dhypervisor)) - ainformation.append('{}Last Hypervisor:{} {}'.format(ansiiprint.purple(), ansiiprint.end(), dlasthypervisor)) - - if long_output == True: - # Disk list - ainformation.append('') - name_length = 0 - for disk in ddisks: - _name_length = len(disk['name']) + 1 - if _name_length > name_length: - name_length = _name_length - ainformation.append('{0}Disks:{1} {2}ID Type {3: <{width}} Dev Bus{4}'.format(ansiiprint.purple(), ansiiprint.end(), ansiiprint.bold(), 'Name', ansiiprint.end(), width=name_length)) - for disk in ddisks: - ainformation.append(' {0: <3} {1: <5} {2: <{width}} {3: <4} {4: <5}'.format(ddisks.index(disk), disk['type'], disk['name'], disk['dev'], disk['bus'], width=name_length)) - # Network list - ainformation.append('') - ainformation.append('{}Interfaces:{} {}ID Type Source Model MAC{}'.format(ansiiprint.purple(), ansiiprint.end(), ansiiprint.bold(), ansiiprint.end())) - for net in dnets: - ainformation.append(' {0: <3} {1: <8} {2: <10} {3: <8} {4}'.format(dnets.index(net), net['type'], net['source'], net['model'], net['mac'])) - # Controller list - ainformation.append('') - ainformation.append('{}Controllers:{} {}ID Type Model{}'.format(ansiiprint.purple(), ansiiprint.end(), ansiiprint.bold(), ansiiprint.end())) - for controller in dcontrollers: - ainformation.append(' {0: <3} {1: <14} {2: <8}'.format(dcontrollers.index(controller), controller['type'], controller['model'])) - - # Join it all together - information = '\n'.join(ainformation) - return information - - -# -# Cluster search functions -# -def getClusterDomainList(zk): - # Get a list of UUIDs by listing the children of /domains - uuid_list = zk.get_children('/domains') - name_list = [] - # For each UUID, get the corresponding name from the data - for uuid in uuid_list: - name_list.append(zk.get('/domains/%s' % uuid)[0].decode('ascii')) - return uuid_list, name_list - -def searchClusterByUUID(zk, uuid): - try: - # Get the lists - uuid_list, name_list = getClusterDomainList(zk) - # We're looking for UUID, so find that element ID - index = uuid_list.index(uuid) - # Get the name_list element at that index - name = name_list[index] - except ValueError: - # We didn't find anything - return None - - return name - -def searchClusterByName(zk, name): - try: - # Get the lists - uuid_list, name_list = getClusterDomainList(zk) - # We're looking for name, so find that element ID - index = name_list.index(name) - # Get the uuid_list element at that index - uuid = uuid_list[index] - except ValueError: - # We didn't find anything - return None - - return uuid - - -# -# Allow mutually exclusive options in Click -# -class MutuallyExclusiveOption(click.Option): - def __init__(self, *args, **kwargs): - meargs = kwargs.pop('mutually_exclusive', []) - _me_arg = [] - _me_func = [] - - for arg in meargs: - _me_arg.append(arg['argument']) - _me_func.append(arg['function']) - - self.me_arg = set(_me_arg) - self.me_func = set(_me_func) - - help = kwargs.get('help', '') - if self.me_func: - ex_str = ', '.join(self.me_arg) - kwargs['help'] = help + ( - ' Mutually exclusive with `' + ex_str + '`.' - ) - - super(MutuallyExclusiveOption, self).__init__(*args, **kwargs) - - def handle_parse_result(self, ctx, opts, args): - if self.me_func.intersection(opts) and self.name in opts: - raise click.UsageError( - "Illegal usage: `{}` is mutually exclusive with " - "arguments `{}`.".format( - self.opts[-1], - ', '.join(self.me_arg) - ) - ) - - return super(MutuallyExclusiveOption, self).handle_parse_result( - ctx, - opts, - args - )