#!/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',
                   'thinksystem 7mm',
                   'thinksystem_7mm',
                   'thinksystem_7mm_vd',                  
                   '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',
                   'raid b540i-2i',
                   'raid_b540i-2i',
                   'raid b545i-2i',
                   'raid b545i-2i',
                   'raid_b545-2i')

BOOT_SLOT_KEYWORDS = ('m.2 socket', '7mm')

class DiskInfo(object):
    def __init__(self, devname, slotinfo=None):
        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 = ''
        self.kernnames = []
        self.vrocmembers = []
        self.is_m2 = False
        self.in_vroc = False
        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('"', '')
            elif k == 'KERNELS':
                self.kernnames.append(v.replace('"', ''))
        if not getattr(self, 'busaddr', None):
            for kernname in getattr(self, 'kernnames', []):
                if kernname in slotinfo:
                    self.is_m2 = True
                    break
        if getattr(self, 'busaddr', None) and self.busaddr in slotinfo:
            self.is_m2 = True
        if not self.driver and 'imsm' not in self.mdcontainer and self.subsystype != 'nvm':
            raise Exception("No driver detected")
        if 'imsm' in self.mdcontainer:
            try:
                mdinfo = subprocess.check_output(
                    ['mdadm', '--detail', '-Y', devnode],
                    stderr=subprocess.DEVNULL)
                if not isinstance(mdinfo, str):
                    mdinfo = mdinfo.decode('utf8', errors='ignore')
                for line in mdinfo.splitlines():
                    key, separator, value = line.partition('=')
                    if separator and key.endswith('_DEV'):
                        member = value.strip()
                        member = member.replace('/dev/', '')
                        if member and member not in self.vrocmembers:
                            self.vrocmembers.append(member)
            except Exception:
                pass
        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.in_vroc:
            return 99  # prefer the assembled array over the member disks
        if self.is_m2:
            return 0
        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():
        label = '{0} {1}'.format(
            slot.get('type', ''), slot.get('designation', '')).lower()
        if not any(x in label for x in BOOT_SLOT_KEYWORDS):
            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.is_m2:
                continue
            if disk.model.lower() in PRIORITY_MODELS:
                continue
            if slotinfo is None:
                slotinfo = get_m2_slot_info()
            if not getattr(disk, 'busaddr', None):
                for kernname in getattr(disk, 'kernnames', []):
                    if kernname in slotinfo:
                        break
                else:
                    disks.remove(disk)
                continue
            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 attach_member_disks(disks):
    for disk in disks:
        num_vroc_members = len(getattr(disk, 'vrocmembers', []))
        for vrocmember in getattr(disk, 'vrocmembers', []):
            for memberdisk in disks:
                if memberdisk.name == vrocmember:
                    memberdisk.in_vroc = True
                    if memberdisk.is_m2:
                        disk.is_m2 = True
                    else:
                        if num_vroc_members == 2 and 'nvme' in vrocmember:
                            # This is not assured, but is likely
                            # Unfortunately, some VMD configurations blow past any way to detect form factor
                            # as DMI table may not include
                            disk.is_m2 = True
                    

def main():
    disks = []
    slotinfo = get_m2_slot_info()
    for disk in sorted(os.listdir('/sys/class/block')):
        try:
            disk = DiskInfo(disk, slotinfo=slotinfo)
            disks.append(disk)
        except SilentException:
            pass
        except Exception as e:
            print("Skipping {0}: {1}".format(disk, str(e)))
    attach_member_disks(disks)
    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()
