#!/usr/bin/python3 import json import subprocess import os import shutil class SilentException(Exception): pass # List of models that are generally used for OS install disk (generally 2-disk mirroring) PRIORITY_MODELS = ('m.2 nvme 2-bay raid kit', 'thinksystem_m.2_vd', 'thinksystem m.2', 'thinksystem_m.2', 'b540p-2hs', # 7mm or M.2 in rear 'b540d-2hs', # M.2 in E3.S or 2.5inch 'b550d-2hs', # Front M.2 'b550p-2hs', # Rear M.2 'raid b540p-2hs', 'raid b540d-2hs', 'raid b550d-2hs', 'raid b550p-2hs') class DiskInfo(object): def __init__(self, devname): if devname.startswith('nvme') and 'c' in devname: raise Exception("Skipping multipath devname") self.name = devname self.wwn = None self.path = None self.model = '' self.size = 0 self.driver = '' self.mdcontainer = '' self.subsystype = '' devnode = '/dev/{0}'.format(devname) qprop = subprocess.check_output( ['udevadm', 'info', '--query=property', devnode]) if not isinstance(qprop, str): qprop = qprop.decode('utf8') for prop in qprop.split('\n'): if '=' not in prop: continue k, v = prop.split('=', 1) if k == 'DEVTYPE' and v != 'disk': if v == 'partition': raise SilentException('Partition') raise Exception('Not a disk') elif k == 'DM_NAME': raise SilentException('Device Mapper') elif k == 'ID_MODEL': self.model = v elif k == 'DEVPATH': self.path = v elif k == 'ID_WWN': self.wwn = v elif k == 'MD_CONTAINER': self.mdcontainer = v attrs = subprocess.check_output(['udevadm', 'info', '-a', devnode]) if not isinstance(attrs, str): attrs = attrs.decode('utf8') for attr in attrs.split('\n'): if '==' not in attr: continue k, v = attr.split('==', 1) k = k.strip() if k == 'ATTRS{size}': self.size = v.replace('"', '') elif (k == 'DRIVERS' and not self.driver and v not in ('"sd"', '""')): self.driver = v.replace('"', '') elif k == 'ATTRS{subsystype}': self.subsystype = v.replace('"', '') elif k == 'ATTR{ro}' and v == '"1"': raise Exception("Device is read-only") elif k == 'ATTRS{removable}' and v == '"1"': raise Exception("Device is removable") elif k == 'ATTRS{address}': self.busaddr = v.replace('"', '') if not self.driver and 'imsm' not in self.mdcontainer and self.subsystype != 'nvm': raise Exception("No driver detected") if self.driver == 'sr': raise Exception('cd/dvd') if os.path.exists('/sys/block/{0}/size'.format(self.name)): with open('/sys/block/{0}/size'.format(self.name), 'r') as sizesrc: self.size = int(sizesrc.read()) * 512 if int(self.size) < 2147483648: raise Exception("Device too small for install ({}MiB)".format(int(self.size)/1024/1024)) @property def priority(self): if self.model.lower() in PRIORITY_MODELS: return 0 if 'imsm' in self.mdcontainer: return 1 if self.driver == 'ahci': return 2 if self.driver.startswith('megaraid'): return 3 if self.driver.startswith('mpt'): return 4 return 99 def __repr__(self): return repr({ 'name': self.name, 'path': self.path, 'wwn': self.wwn, 'driver': self.driver, 'size': self.size, 'model': self.model, }) def get_dmi_type9_info(): dmidecode = shutil.which('dmidecode') if not dmidecode: return {} try: dmio = subprocess.check_output( [dmidecode, '--type', '9'], stderr=subprocess.DEVNULL) except Exception: return {} if not isinstance(dmio, str): dmio = dmio.decode('utf8', errors='ignore') slots = {} currentslot = None for line in dmio.split('\n'): stripped = line.strip() if not stripped: if currentslot: busaddr = currentslot.get('bus_address', None) if busaddr: slots[busaddr] = currentslot currentslot = None continue if stripped.startswith('Handle ') and ', DMI type 9,' in stripped: if currentslot: busaddr = currentslot.get('bus_address', None) if busaddr: slots[busaddr] = currentslot currentslot = {'handle': stripped.split(',', 1)[0].split(' ', 1)[1]} continue if currentslot is None: continue if ':' not in stripped: continue key, val = stripped.split(':', 1) key = key.strip().lower().replace(' ', '_') currentslot[key] = val.strip() if currentslot: busaddr = currentslot.get('bus_address', None) if busaddr: slots[busaddr] = currentslot return slots def get_m2_slot_info(): slots = get_dmi_type9_info() m2slots = {} for slot in slots.values(): if 'm.2 socket' not in slot.get('type', '').lower(): continue busaddr = slot.get('bus_address', None) if not busaddr: continue m2slots[busaddr] = slot return m2slots def get_deployment_storage(): apiclient = None for clientpath in ('/opt/confluent/bin/apiclient', '/etc/confluent/apiclient'): if os.path.exists(clientpath): apiclient = clientpath break if not apiclient: return None for attriburl in ('/confluent-api/self/myattrib', '/confluent-api/self/myattribs'): try: attribs = subprocess.check_output( ['python3', apiclient, attriburl, '-j'], stderr=subprocess.DEVNULL) if not isinstance(attribs, str): attribs = attribs.decode('utf8') attribs = json.loads(attribs) if not isinstance(attribs, dict): continue storage = attribs.get('deployment.storage', None) if isinstance(storage, str) and storage.strip(): return storage.strip() except Exception: continue return None def filter_deployment_storage(disks, storagespec, slotinfo=None): if not storagespec: return disks spec = storagespec.strip().lower() for disk in list(disks): if spec == 'm2': if disk.model.lower() in PRIORITY_MODELS: continue if not getattr(disk, 'busaddr', None): disks.remove(disk) continue if slotinfo is None: slotinfo = get_m2_slot_info() if disk.busaddr not in slotinfo: disks.remove(disk) elif spec.startswith('/dev/'): bspec = os.path.basename(spec) if disk.name.lower() == bspec: return [disk] else: raise Exception("Unknown deployment.storage specification: {0}".format(storagespec)) if not disks: raise Exception("No disks match deployment.storage specification: {0}".format(storagespec)) return disks def main(): disks = [] for disk in sorted(os.listdir('/sys/class/block')): try: disk = DiskInfo(disk) disks.append(disk) except SilentException: pass except Exception as e: print("Skipping {0}: {1}".format(disk, str(e))) slotinfo = get_m2_slot_info() disks = filter_deployment_storage(disks, get_deployment_storage(), slotinfo=slotinfo) nd = [x.name for x in sorted(disks, key=lambda x: [x.priority, x.size])] if nd: open('/tmp/installdisk', 'w').write(nd[0]) else: raise Exception("No suitable install disk found") if __name__ == '__main__': main()