diff --git a/.bin/Scripts/functions/disk.py b/.bin/Scripts/functions/disk.py index 257f21b0..88b26155 100644 --- a/.bin/Scripts/functions/disk.py +++ b/.bin/Scripts/functions/disk.py @@ -6,390 +6,392 @@ from settings.partition_uids import * # Regex REGEX_BAD_PARTITION = re.compile(r'(RAW|Unknown)', re.IGNORECASE) REGEX_DISK_GPT = re.compile( - r'Disk ID: {[A-Z0-9]+-[A-Z0-9]+-[A-Z0-9]+-[A-Z0-9]+-[A-Z0-9]+}', - re.IGNORECASE) + r'Disk ID: {[A-Z0-9]+-[A-Z0-9]+-[A-Z0-9]+-[A-Z0-9]+-[A-Z0-9]+}', + re.IGNORECASE) REGEX_DISK_MBR = re.compile(r'Disk ID: [A-Z0-9]+', re.IGNORECASE) REGEX_DISK_RAW = re.compile(r'Disk ID: 00000000', re.IGNORECASE) def assign_volume_letters(): - """Assign a volume letter to all available volumes.""" - remove_volume_letters() + """Assign a volume letter to all available volumes.""" + remove_volume_letters() - # Write script - script = [] - for vol in get_volumes(): - script.append('select volume {}'.format(vol['Number'])) - script.append('assign') + # Write script + script = [] + for vol in get_volumes(): + script.append('select volume {}'.format(vol['Number'])) + script.append('assign') - # Run - run_diskpart(script) + # Run + run_diskpart(script) def get_boot_mode(): - """Check if the boot mode was UEFI or legacy.""" - boot_mode = 'Legacy' - try: - reg_key = winreg.OpenKey( - winreg.HKEY_LOCAL_MACHINE, r'System\CurrentControlSet\Control') - reg_value = winreg.QueryValueEx(reg_key, 'PEFirmwareType')[0] - if reg_value == 2: - boot_mode = 'UEFI' - except: - boot_mode = 'Unknown' + """Check if the boot mode was UEFI or legacy.""" + boot_mode = 'Legacy' + try: + reg_key = winreg.OpenKey( + winreg.HKEY_LOCAL_MACHINE, r'System\CurrentControlSet\Control') + reg_value = winreg.QueryValueEx(reg_key, 'PEFirmwareType')[0] + if reg_value == 2: + boot_mode = 'UEFI' + except: + boot_mode = 'Unknown' - return boot_mode + return boot_mode def get_disk_details(disk): - """Get disk details using DiskPart.""" - details = {} - script = [ - 'select disk {}'.format(disk['Number']), - 'detail disk'] + """Get disk details using DiskPart.""" + details = {} + script = [ + 'select disk {}'.format(disk['Number']), + 'detail disk'] - # Run - try: - result = run_diskpart(script) - except subprocess.CalledProcessError: - pass - else: - output = result.stdout.decode().strip() - # Remove empty lines - tmp = [s.strip() for s in output.splitlines() if s.strip() != ''] - # Set disk name - details['Name'] = tmp[4] - # Split each line on ':' skipping those without ':' - tmp = [s.split(':') for s in tmp if ':' in s] - # Add key/value pairs to the details variable and return dict - details.update({key.strip(): value.strip() for (key, value) in tmp}) + # Run + try: + result = run_diskpart(script) + except subprocess.CalledProcessError: + pass + else: + output = result.stdout.decode().strip() + # Remove empty lines + tmp = [s.strip() for s in output.splitlines() if s.strip() != ''] + # Set disk name + details['Name'] = tmp[4] + # Split each line on ':' skipping those without ':' + tmp = [s.split(':') for s in tmp if ':' in s] + # Add key/value pairs to the details variable and return dict + details.update({key.strip(): value.strip() for (key, value) in tmp}) - return details + return details def get_disks(): - """Get list of attached disks using DiskPart.""" - disks = [] + """Get list of attached disks using DiskPart.""" + disks = [] - try: - # Run script - result = run_diskpart(['list disk']) - except subprocess.CalledProcessError: - pass - else: - # Append disk numbers - output = result.stdout.decode().strip() - for tmp in re.findall(r'Disk (\d+)\s+\w+\s+(\d+\s+\w+)', output): - num = tmp[0] - size = human_readable_size(tmp[1]) - disks.append({'Number': num, 'Size': size}) + try: + # Run script + result = run_diskpart(['list disk']) + except subprocess.CalledProcessError: + pass + else: + # Append disk numbers + output = result.stdout.decode().strip() + for tmp in re.findall(r'Disk (\d+)\s+\w+\s+(\d+\s+\w+)', output): + num = tmp[0] + size = human_readable_size(tmp[1]) + disks.append({'Number': num, 'Size': size}) - return disks + return disks def get_partition_details(disk, partition): - """Get partition details using DiskPart and fsutil.""" - details = {} - script = [ - 'select disk {}'.format(disk['Number']), - 'select partition {}'.format(partition['Number']), - 'detail partition'] + """Get partition details using DiskPart and fsutil.""" + details = {} + script = [ + 'select disk {}'.format(disk['Number']), + 'select partition {}'.format(partition['Number']), + 'detail partition'] - # Diskpart details + # Diskpart details + try: + # Run script + result = run_diskpart(script) + except subprocess.CalledProcessError: + pass + else: + # Get volume letter or RAW status + output = result.stdout.decode().strip() + tmp = re.search(r'Volume\s+\d+\s+(\w|RAW)\s+', output) + if tmp: + if tmp.group(1).upper() == 'RAW': + details['FileSystem'] = RAW + else: + details['Letter'] = tmp.group(1) + # Remove empty lines from output + tmp = [s.strip() for s in output.splitlines() if s.strip() != ''] + # Split each line on ':' skipping those without ':' + tmp = [s.split(':') for s in tmp if ':' in s] + # Add key/value pairs to the details variable and return dict + details.update({key.strip(): value.strip() for (key, value) in tmp}) + + # Get MBR type / GPT GUID for extra details on "Unknown" partitions + guid = PARTITION_UIDS.get(details.get('Type').upper(), {}) + if guid: + details.update({ + 'Description': guid.get('Description', '')[:29], + 'OS': guid.get('OS', 'Unknown')[:27]}) + + if 'Letter' in details: + # Disk usage try: - # Run script - result = run_diskpart(script) - except subprocess.CalledProcessError: - pass + tmp = psutil.disk_usage('{}:\\'.format(details['Letter'])) + except OSError as err: + details['FileSystem'] = 'Unknown' + details['Error'] = err.strerror else: - # Get volume letter or RAW status - output = result.stdout.decode().strip() - tmp = re.search(r'Volume\s+\d+\s+(\w|RAW)\s+', output) - if tmp: - if tmp.group(1).upper() == 'RAW': - details['FileSystem'] = RAW - else: - details['Letter'] = tmp.group(1) - # Remove empty lines from output - tmp = [s.strip() for s in output.splitlines() if s.strip() != ''] - # Split each line on ':' skipping those without ':' - tmp = [s.split(':') for s in tmp if ':' in s] - # Add key/value pairs to the details variable and return dict - details.update({key.strip(): value.strip() for (key, value) in tmp}) + details['Used Space'] = human_readable_size(tmp.used) - # Get MBR type / GPT GUID for extra details on "Unknown" partitions - guid = PARTITION_UIDS.get(details.get('Type').upper(), {}) - if guid: - details.update({ - 'Description': guid.get('Description', '')[:29], - 'OS': guid.get('OS', 'Unknown')[:27]}) + # fsutil details + cmd = [ + 'fsutil', + 'fsinfo', + 'volumeinfo', + '{}:'.format(details['Letter']) + ] + try: + result = run_program(cmd) + except subprocess.CalledProcessError: + pass + else: + output = result.stdout.decode().strip() + # Remove empty lines from output + tmp = [s.strip() for s in output.splitlines() if s.strip() != ''] + # Add "Feature" lines + details['File System Features'] = [s.strip() for s in tmp + if ':' not in s] + # Split each line on ':' skipping those without ':' + tmp = [s.split(':') for s in tmp if ':' in s] + # Add key/value pairs to the details variable and return dict + details.update({key.strip(): value.strip() for (key, value) in tmp}) - if 'Letter' in details: - # Disk usage - try: - tmp = psutil.disk_usage('{}:\\'.format(details['Letter'])) - except OSError as err: - details['FileSystem'] = 'Unknown' - details['Error'] = err.strerror - else: - details['Used Space'] = human_readable_size(tmp.used) + # Set Volume Name + details['Name'] = details.get('Volume Name', '') - # fsutil details - cmd = [ - 'fsutil', - 'fsinfo', - 'volumeinfo', - '{}:'.format(details['Letter']) - ] - try: - result = run_program(cmd) - except subprocess.CalledProcessError: - pass - else: - output = result.stdout.decode().strip() - # Remove empty lines from output - tmp = [s.strip() for s in output.splitlines() if s.strip() != ''] - # Add "Feature" lines - details['File System Features'] = [s.strip() for s in tmp - if ':' not in s] - # Split each line on ':' skipping those without ':' - tmp = [s.split(':') for s in tmp if ':' in s] - # Add key/value pairs to the details variable and return dict - details.update({key.strip(): value.strip() for (key, value) in tmp}) + # Set FileSystem Type + if details.get('FileSystem', '') not in ['RAW', 'Unknown']: + details['FileSystem'] = details.get('File System Name', 'Unknown') - # Set Volume Name - details['Name'] = details.get('Volume Name', '') - - # Set FileSystem Type - if details.get('FileSystem', '') not in ['RAW', 'Unknown']: - details['FileSystem'] = details.get('File System Name', 'Unknown') - - return details + return details def get_partitions(disk): - """Get list of partition using DiskPart.""" - partitions = [] - script = [ - 'select disk {}'.format(disk['Number']), - 'list partition'] + """Get list of partition using DiskPart.""" + partitions = [] + script = [ + 'select disk {}'.format(disk['Number']), + 'list partition'] - try: - # Run script - result = run_diskpart(script) - except subprocess.CalledProcessError: - pass - else: - # Append partition numbers - output = result.stdout.decode().strip() - regex = r'Partition\s+(\d+)\s+\w+\s+(\d+\s+\w+)\s+' - for tmp in re.findall(regex, output, re.IGNORECASE): - num = tmp[0] - size = human_readable_size(tmp[1]) - partitions.append({'Number': num, 'Size': size}) + try: + # Run script + result = run_diskpart(script) + except subprocess.CalledProcessError: + pass + else: + # Append partition numbers + output = result.stdout.decode().strip() + regex = r'Partition\s+(\d+)\s+\w+\s+(\d+\s+\w+)\s+' + for tmp in re.findall(regex, output, re.IGNORECASE): + num = tmp[0] + size = human_readable_size(tmp[1]) + partitions.append({'Number': num, 'Size': size}) - return partitions + return partitions def get_table_type(disk): - """Get disk partition table type using DiskPart.""" - part_type = 'Unknown' - script = [ - 'select disk {}'.format(disk['Number']), - 'uniqueid disk'] + """Get disk partition table type using DiskPart.""" + part_type = 'Unknown' + script = [ + 'select disk {}'.format(disk['Number']), + 'uniqueid disk'] - try: - result = run_diskpart(script) - except subprocess.CalledProcessError: - pass - else: - output = result.stdout.decode().strip() - if REGEX_DISK_GPT.search(output): - part_type = 'GPT' - elif REGEX_DISK_MBR.search(output): - part_type = 'MBR' - elif REGEX_DISK_RAW.search(output): - part_type = 'RAW' + try: + result = run_diskpart(script) + except subprocess.CalledProcessError: + pass + else: + output = result.stdout.decode().strip() + if REGEX_DISK_GPT.search(output): + part_type = 'GPT' + elif REGEX_DISK_MBR.search(output): + part_type = 'MBR' + elif REGEX_DISK_RAW.search(output): + part_type = 'RAW' - return part_type + return part_type def get_volumes(): - """Get list of volumes using DiskPart.""" - vols = [] - try: - result = run_diskpart(['list volume']) - except subprocess.CalledProcessError: - pass - else: - # Append volume numbers - output = result.stdout.decode().strip() - for tmp in re.findall(r'Volume (\d+)\s+([A-Za-z]?)\s+', output): - vols.append({'Number': tmp[0], 'Letter': tmp[1]}) + """Get list of volumes using DiskPart.""" + vols = [] + try: + result = run_diskpart(['list volume']) + except subprocess.CalledProcessError: + pass + else: + # Append volume numbers + output = result.stdout.decode().strip() + for tmp in re.findall(r'Volume (\d+)\s+([A-Za-z]?)\s+', output): + vols.append({'Number': tmp[0], 'Letter': tmp[1]}) - return vols + return vols def is_bad_partition(par): - """Check if the partition is accessible.""" - return 'Letter' not in par or REGEX_BAD_PARTITION.search(par['FileSystem']) + """Check if the partition is accessible.""" + return 'Letter' not in par or REGEX_BAD_PARTITION.search(par['FileSystem']) def prep_disk_for_formatting(disk=None): - """Gather details about the disk and its partitions.""" - disk['Format Warnings'] = '\n' - width = len(str(len(disk['Partitions']))) + """Gather details about the disk and its partitions.""" + disk['Format Warnings'] = '\n' + width = len(str(len(disk['Partitions']))) - # Bail early - if disk is None: - raise Exception('Disk not provided.') + # Bail early + if disk is None: + raise Exception('Disk not provided.') - # Set boot method and partition table type - disk['Use GPT'] = True - if (get_boot_mode() == 'UEFI'): - if (not ask("Setup Windows to use UEFI booting?")): - disk['Use GPT'] = False + # Set boot method and partition table type + disk['Use GPT'] = True + if (get_boot_mode() == 'UEFI'): + if (not ask("Setup Windows to use UEFI booting?")): + disk['Use GPT'] = False + else: + if (ask("Setup Windows to use BIOS/Legacy booting?")): + disk['Use GPT'] = False + + # Set Display and Warning Strings + if len(disk['Partitions']) == 0: + disk['Format Warnings'] += 'No partitions found\n' + for partition in disk['Partitions']: + display = '{size} {fs}'.format( + num = partition['Number'], + width = width, + size = partition['Size'], + fs = partition['FileSystem']) + + if is_bad_partition(partition): + # Set display string using partition description & OS type + display += '\t\t{q}{name}{q}\t{desc} ({os})'.format( + display = display, + q = '"' if partition['Name'] != '' else '', + name = partition['Name'], + desc = partition['Description'], + os = partition['OS']) else: - if (ask("Setup Windows to use BIOS/Legacy booting?")): - disk['Use GPT'] = False - - # Set Display and Warning Strings - if len(disk['Partitions']) == 0: - disk['Format Warnings'] += 'No partitions found\n' - for partition in disk['Partitions']: - display = '{size} {fs}'.format( - num = partition['Number'], - width = width, - size = partition['Size'], - fs = partition['FileSystem']) - - if is_bad_partition(partition): - # Set display string using partition description & OS type - display += '\t\t{q}{name}{q}\t{desc} ({os})'.format( - display = display, - q = '"' if partition['Name'] != '' else '', - name = partition['Name'], - desc = partition['Description'], - os = partition['OS']) - else: - # List space used instead of partition description & OS type - display += ' (Used: {used})\t{q}{name}{q}'.format( - used = partition['Used Space'], - q = '"' if partition['Name'] != '' else '', - name = partition['Name']) - # For all partitions - partition['Display String'] = display + # List space used instead of partition description & OS type + display += ' (Used: {used})\t{q}{name}{q}'.format( + used = partition['Used Space'], + q = '"' if partition['Name'] != '' else '', + name = partition['Name']) + # For all partitions + partition['Display String'] = display def reassign_volume_letter(letter, new_letter='I'): - """Assign a new letter to a volume using DiskPart.""" - if not letter: - # Ignore - return None - script = [ - 'select volume {}'.format(letter), - 'remove noerr', - 'assign letter={}'.format(new_letter)] - try: - run_diskpart(script) - except subprocess.CalledProcessError: - pass - else: - return new_letter + """Assign a new letter to a volume using DiskPart.""" + if not letter: + # Ignore + return None + script = [ + 'select volume {}'.format(letter), + 'remove noerr', + 'assign letter={}'.format(new_letter)] + try: + run_diskpart(script) + except subprocess.CalledProcessError: + pass + else: + return new_letter def remove_volume_letters(keep=None): - """Remove all assigned volume letters using DiskPart.""" - if not keep: - keep = '' + """Remove all assigned volume letters using DiskPart.""" + if not keep: + keep = '' - script = [] - for vol in get_volumes(): - if vol['Letter'].upper() != keep.upper(): - script.append('select volume {}'.format(vol['Number'])) - script.append('remove noerr') + script = [] + for vol in get_volumes(): + if vol['Letter'].upper() != keep.upper(): + script.append('select volume {}'.format(vol['Number'])) + script.append('remove noerr') - # Run script - try: - run_diskpart(script) - except subprocess.CalledProcessError: - pass + # Run script + try: + run_diskpart(script) + except subprocess.CalledProcessError: + pass def run_diskpart(script): - """Run DiskPart script.""" - tempfile = r'{}\diskpart.script'.format(global_vars['Env']['TMP']) + """Run DiskPart script.""" + tempfile = r'{}\diskpart.script'.format(global_vars['Env']['TMP']) - # Write script - with open(tempfile, 'w') as f: - for line in script: - f.write('{}\n'.format(line)) + # Write script + with open(tempfile, 'w') as f: + for line in script: + f.write('{}\n'.format(line)) - # Run script - cmd = [ - r'{}\Windows\System32\diskpart.exe'.format( - global_vars['Env']['SYSTEMDRIVE']), - '/s', tempfile] - result = run_program(cmd) - sleep(2) - return result + # Run script + cmd = [ + r'{}\Windows\System32\diskpart.exe'.format( + global_vars['Env']['SYSTEMDRIVE']), + '/s', tempfile] + result = run_program(cmd) + sleep(2) + return result def scan_disks(): - """Get details about the attached disks""" - disks = get_disks() + """Get details about the attached disks""" + disks = get_disks() - # Get disk details - for disk in disks: - # Get partition style - disk['Table'] = get_table_type(disk) + # Get disk details + for disk in disks: + # Get partition style + disk['Table'] = get_table_type(disk) - # Get disk name/model and physical details - disk.update(get_disk_details(disk)) + # Get disk name/model and physical details + disk.update(get_disk_details(disk)) - # Get partition info for disk - disk['Partitions'] = get_partitions(disk) + # Get partition info for disk + disk['Partitions'] = get_partitions(disk) - for partition in disk['Partitions']: - # Get partition details - partition.update(get_partition_details(disk, partition)) + for partition in disk['Partitions']: + # Get partition details + partition.update(get_partition_details(disk, partition)) - # Done - return disks + # Done + return disks def select_disk(title='Which disk?', disks=[]): - """Select a disk from the attached disks""" - # Build menu - disk_options = [] - for disk in disks: - display_name = '{}\t[{}] ({}) {}'.format( - disk.get('Size', ''), - disk.get('Table', ''), - disk.get('Type', ''), - disk.get('Name', 'Unknown'), - ) - pwidth=len(str(len(disk['Partitions']))) - for partition in disk['Partitions']: - # Main text - p_name = 'Partition {num:>{width}}: {size} ({fs})'.format( - num = partition['Number'], - width = pwidth, - size = partition['Size'], - fs = partition['FileSystem']) - if partition['Name']: - p_name += '\t"{}"'.format(partition['Name']) + """Select a disk from the attached disks""" + # Build menu + disk_options = [] + for disk in disks: + display_name = '{}\t[{}] ({}) {}'.format( + disk.get('Size', ''), + disk.get('Table', ''), + disk.get('Type', ''), + disk.get('Name', 'Unknown'), + ) + pwidth=len(str(len(disk['Partitions']))) + for partition in disk['Partitions']: + # Main text + p_name = 'Partition {num:>{width}}: {size} ({fs})'.format( + num = partition['Number'], + width = pwidth, + size = partition['Size'], + fs = partition['FileSystem']) + if partition['Name']: + p_name += '\t"{}"'.format(partition['Name']) - # Show unsupported partition(s) - if is_bad_partition(partition): - p_name = '{YELLOW}{p_name}{CLEAR}'.format( - p_name=p_name, **COLORS) + # Show unsupported partition(s) + if is_bad_partition(partition): + p_name = '{YELLOW}{p_name}{CLEAR}'.format( + p_name=p_name, **COLORS) - display_name += '\n\t\t\t{}'.format(p_name) - if not disk['Partitions']: - display_name += '\n\t\t\t{}No partitions found.{}'.format( - COLORS['YELLOW'], COLORS['CLEAR']) + display_name += '\n\t\t\t{}'.format(p_name) + if not disk['Partitions']: + display_name += '\n\t\t\t{}No partitions found.{}'.format( + COLORS['YELLOW'], COLORS['CLEAR']) - disk_options.append({'Name': display_name, 'Disk': disk}) - actions = [ - {'Name': 'Main Menu', 'Letter': 'M'}, - ] + disk_options.append({'Name': display_name, 'Disk': disk}) + actions = [ + {'Name': 'Main Menu', 'Letter': 'M'}, + ] - # Menu loop - selection = menu_select( - title = title, - main_entries = disk_options, - action_entries = actions) + # Menu loop + selection = menu_select( + title = title, + main_entries = disk_options, + action_entries = actions) - if (selection.isnumeric()): - return disk_options[int(selection)-1]['Disk'] - elif (selection == 'M'): - raise GenericAbort + if (selection.isnumeric()): + return disk_options[int(selection)-1]['Disk'] + elif (selection == 'M'): + raise GenericAbort if __name__ == '__main__': - print("This file is not meant to be called directly.") + print("This file is not meant to be called directly.") + +# vim: sts=2 sw=2 ts=2