2Shirt/WizardKit
2Shirt/WizardKit
1
0
Fork 0
Code Issues 2 Pull requests Projects Releases 26 Packages Wiki Activity Actions

functions\disk.py done

Browse source
This commit is contained in:
Alan Mason 2017-11-30 23:06:23 -08:00
parent a0460d6a82
commit c951380482
1 changed files with 118 additions and 121 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

239
Scripts/functions/disk.py
View file

@ -3,10 +3,18 @@
from functions.common import * from functions.common import *
import partition_uids import partition_uids
# 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)
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(): def assign_volume_letters():
with open(DISKPART_SCRIPT, 'w') as script: with open(DISKPART_SCRIPT, 'w') as script:
for vol in get_volumes(): for vol in get_volumes():
script.write('select volume {Number}\n'.format(**vol)) script.write('select volume {}\n'.format(vol['Number']))
script.write('assign\n') script.write('assign\n')
# Remove current letters # Remove current letters
@ -14,14 +22,15 @@ def assign_volume_letters():
# Run script # Run script
try: try:
run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT)) run_program(['diskpart', '/s', DISKPART_SCRIPT])
except subprocess.CalledProcessError: except subprocess.CalledProcessError:
pass pass
def get_boot_mode(): def get_boot_mode():
boot_mode = 'Legacy' boot_mode = 'Legacy'
try: try:
reg_key = winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, 'System\\CurrentControlSet\\Control') reg_key = winreg.OpenKey(
winreg.HKEY_LOCAL_MACHINE, r'System\CurrentControlSet\Control')
reg_value = winreg.QueryValueEx(reg_key, 'PEFirmwareType')[0] reg_value = winreg.QueryValueEx(reg_key, 'PEFirmwareType')[0]
if reg_value == 2: if reg_value == 2:
boot_mode = 'UEFI' boot_mode = 'UEFI'
@ -30,32 +39,25 @@ def get_boot_mode():
return boot_mode return boot_mode
def get_disk_details(disk=None): def get_disk_details(disk):
details = {} details = {}
with open(DISKPART_SCRIPT, 'w') as script:
# Bail early script.write('select disk {}\n'.format(disk['Number']))
if disk is None: script.write('detail disk\n')
raise Exception('Disk not specified.')
try: try:
# Run script # Run script
with open(DISKPART_SCRIPT, 'w') as script: output = run_program(['diskpart', '/s', DISKPART_SCRIPT])
script.write('select disk {Number}\n'.format(**disk)) output = output.stdout.decode().strip()
script.write('detail disk\n')
process_return = run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT))
process_return = process_return.stdout.decode().strip()
except subprocess.CalledProcessError: except subprocess.CalledProcessError:
pass pass
else: else:
# Remove empty lines # Remove empty lines
tmp = [s.strip() for s in process_return.splitlines() if s.strip() != ''] tmp = [s.strip() for s in output.splitlines() if s.strip() != '']
# Set disk name # Set disk name
details['Name'] = tmp[4] details['Name'] = tmp[4]
# Split each line on ':' skipping those without ':'
# Remove lines without a ':' and split each remaining line at the ':' to form a key/value pair
tmp = [s.split(':') for s in tmp if ':' in s] tmp = [s.split(':') for s in tmp if ':' in s]
# Add key/value pairs to the details variable and return dict # Add key/value pairs to the details variable and return dict
details.update({key.strip(): value.strip() for (key, value) in tmp}) details.update({key.strip(): value.strip() for (key, value) in tmp})
@ -63,90 +65,85 @@ def get_disk_details(disk=None):
def get_disks(): def get_disks():
disks = [] disks = []
with open(DISKPART_SCRIPT, 'w') as script:
script.write('list disk\n')
try: try:
# Run script # Run script
with open(DISKPART_SCRIPT, 'w') as script: output = run_program(['diskpart', '/s', DISKPART_SCRIPT])
script.write('list disk\n') output = output.stdout.decode().strip()
process_return = run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT))
process_return = process_return.stdout.decode().strip()
except subprocess.CalledProcessError: except subprocess.CalledProcessError:
pass pass
else: else:
# Append disk numbers # Append disk numbers
for tmp in re.findall(r'Disk (\d+)\s+\w+\s+(\d+\s+\w+)', process_return): for tmp in re.findall(r'Disk (\d+)\s+\w+\s+(\d+\s+\w+)', output):
_num = tmp[0] num = tmp[0]
_size = human_readable_size(tmp[1]) size = human_readable_size(tmp[1])
disks.append({'Number': _num, 'Size': _size}) disks.append({'Number': num, 'Size': size})
return disks return disks
def get_partition_details(disk=None, par=None): def get_partition_details(disk, partition):
details = {} details = {}
with open(DISKPART_SCRIPT, 'w') as script:
# Bail early script.write('select disk {}\n'.format(disk['Number']))
if disk is None: script.write('select partition {}\n'.format(partition['Number']))
raise Exception('Disk not specified.') script.write('detail partition\n')
if par is None:
raise Exception('Partition not specified.')
# Diskpart details # Diskpart details
try: try:
# Run script # Run script
with open(DISKPART_SCRIPT, 'w') as script: output = run_program(['diskpart', '/s', DISKPART_SCRIPT])
script.write('select disk {Number}\n'.format(**disk)) output = output.stdout.decode().strip()
script.write('select partition {Number}\n'.format(**par))
script.write('detail partition\n')
process_return = run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT))
process_return = process_return.stdout.decode().strip()
except subprocess.CalledProcessError: except subprocess.CalledProcessError:
pass pass
else: else:
# Get volume letter or RAW status # Get volume letter or RAW status
tmp = re.search(r'Volume\s+\d+\s+(\w|RAW)\s+', process_return) tmp = re.search(r'Volume\s+\d+\s+(\w|RAW)\s+', output)
if tmp: if tmp:
if tmp.group(1).upper() == 'RAW': if tmp.group(1).upper() == 'RAW':
details['FileSystem'] = RAW details['FileSystem'] = RAW
else: else:
details['Letter'] = tmp.group(1) details['Letter'] = tmp.group(1)
# Remove empty lines from output
# Remove empty lines from process_return tmp = [s.strip() for s in output.splitlines() if s.strip() != '']
tmp = [s.strip() for s in process_return.splitlines() if s.strip() != ''] # Split each line on ':' skipping those without ':'
# Remove lines without a ':' and split each remaining line at the ':' to form a key/value pair
tmp = [s.split(':') for s in tmp if ':' in s] tmp = [s.split(':') for s in tmp if ':' in s]
# Add key/value pairs to the details variable and return dict # Add key/value pairs to the details variable and return dict
details.update({key.strip(): value.strip() for (key, value) in tmp}) details.update({key.strip(): value.strip() for (key, value) in tmp})
# Get MBR type / GPT GUID for extra details on "Unknown" partitions # Get MBR type / GPT GUID for extra details on "Unknown" partitions
guid = partition_uids.lookup_guid(details['Type']) guid = partition_uids.lookup_guid(details['Type'])
if guid is not None: if guid:
details.update({ details.update({
'Description': guid.get('Description', ''), 'Description': guid.get('Description', ''),
'OS': guid.get('OS', '')}) 'OS': guid.get('OS', '')})
if 'Letter' in details: if 'Letter' in details:
# Disk usage # Disk usage
tmp = shutil.disk_usage('{Letter}:\\'.format(**details)) tmp = shutil.disk_usage('{}:\\'.format(details['Letter']))
details['Used Space'] = human_readable_size(tmp.used) details['Used Space'] = human_readable_size(tmp.used)
# fsutil details # fsutil details
cmd = [
'fsutil',
'fsinfo',
'volumeinfo',
'{}:'.format(details['Letter'])
]
try: try:
process_return = run_program('fsutil fsinfo volumeinfo {Letter}:'.format(**details)) output = run_program(cmd)
process_return = process_return.stdout.decode().strip() output = output.stdout.decode().strip()
except subprocess.CalledProcessError: except subprocess.CalledProcessError:
pass pass
else: else:
# Remove empty lines from process_return # Remove empty lines from output
tmp = [s.strip() for s in process_return.splitlines() if s.strip() != ''] tmp = [s.strip() for s in output.splitlines() if s.strip() != '']
# Add "Feature" lines # Add "Feature" lines
details['File System Features'] = [s.strip() for s in tmp if ':' not in s] details['File System Features'] = [s.strip() for s in tmp
if ':' not in s]
# Remove lines without a ':' and split each remaining line at the ':' to form a key/value pair # Split each line on ':' skipping those without ':'
tmp = [s.split(':') for s in tmp if ':' in s] tmp = [s.split(':') for s in tmp if ':' in s]
# Add key/value pairs to the details variable and return dict # Add key/value pairs to the details variable and return dict
details.update({key.strip(): value.strip() for (key, value) in tmp}) details.update({key.strip(): value.strip() for (key, value) in tmp})
@ -159,74 +156,72 @@ def get_partition_details(disk=None, par=None):
return details return details
def get_partitions(disk=None): def get_partitions(disk):
partitions = [] partitions = []
with open(DISKPART_SCRIPT, 'w') as script:
# Bail early script.write('select disk {}\n'.format(disk['Number']))
if disk is None: script.write('list partition\n')
raise Exception('Disk not specified.')
try: try:
# Run script # Run script
with open(DISKPART_SCRIPT, 'w') as script: output = run_program(['diskpart', '/s', DISKPART_SCRIPT])
script.write('select disk {Number}\n'.format(**disk)) output = output.stdout.decode().strip()
script.write('list partition\n')
process_return = run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT))
process_return = process_return.stdout.decode().strip()
except subprocess.CalledProcessError: except subprocess.CalledProcessError:
pass pass
else: else:
# Append partition numbers # Append partition numbers
for tmp in re.findall(r'Partition\s+(\d+)\s+\w+\s+(\d+\s+\w+)\s+', process_return, re.IGNORECASE): regex = r'Partition\s+(\d+)\s+\w+\s+(\d+\s+\w+)\s+'
_num = tmp[0] for tmp in re.findall(regex, output, re.IGNORECASE):
_size = human_readable_size(tmp[1]) num = tmp[0]
partitions.append({'Number': _num, 'Size': _size}) size = human_readable_size(tmp[1])
partitions.append({'Number': num, 'Size': size})
return partitions return partitions
def get_table_type(disk=None): def get_table_type(disk):
_type = 'Unknown' part_type = 'Unknown'
with open(DISKPART_SCRIPT, 'w') as script:
# Bail early script.write('select disk {}\n'.format(disk['Number']))
if disk is None: script.write('uniqueid disk\n')
raise Exception('Disk not specified.')
try: try:
with open(DISKPART_SCRIPT, 'w') as script: output = run_program(['diskpart', '/s', DISKPART_SCRIPT])
script.write('select disk {Number}\n'.format(**disk)) output = output.stdout.decode().strip()
script.write('uniqueid disk\n')
process_return = run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT))
process_return = process_return.stdout.decode().strip()
except subprocess.CalledProcessError: except subprocess.CalledProcessError:
pass pass
else: else:
if re.findall(r'Disk ID: {[A-Z0-9]+-[A-Z0-9]+-[A-Z0-9]+-[A-Z0-9]+-[A-Z0-9]+}', process_return, re.IGNORECASE): if REGEX_DISK_GPT.search(output):
_type = 'GPT' part_type = 'GPT'
elif re.findall(r'Disk ID: 00000000', process_return, re.IGNORECASE): elif REGEX_DISK_MBR.search(output):
_type = 'RAW' part_type = 'MBR'
elif re.findall(r'Disk ID: [A-Z0-9]+', process_return, re.IGNORECASE): elif REGEX_DISK_RAW.search(output):
_type = 'MBR' part_type = 'RAW'
else:
part_type = 'Unknown
return _type return part_type
def get_volumes(): def get_volumes():
vols = [] vols = []
with open(DISKPART_SCRIPT, 'w') as script:
script.write('list volume\n')
try: try:
# Run script # Run script
with open(DISKPART_SCRIPT, 'w') as script: output = run_program(['diskpart', '/s', DISKPART_SCRIPT])
script.write('list volume\n') output = output.stdout.decode().strip()
process_return = run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT))
process_return = process_return.stdout.decode().strip()
except subprocess.CalledProcessError: except subprocess.CalledProcessError:
pass pass
else: else:
# Append volume numbers # Append volume numbers
for tmp in re.findall(r'Volume (\d+)\s+([A-Za-z]?)\s+', process_return): for tmp in re.findall(r'Volume (\d+)\s+([A-Za-z]?)\s+', output):
vols.append({'Number': tmp[0], 'Letter': tmp[1]}) vols.append({'Number': tmp[0], 'Letter': tmp[1]})
return vols return vols
def is_bad_partition(par):
return 'Letter' not in par or REGEX_BAD_PARTITION.search(par['FileSystem'])
def prep_disk_for_formatting(disk=None): def prep_disk_for_formatting(disk=None):
disk['Format Warnings'] = '\n' disk['Format Warnings'] = '\n'
width = len(str(len(disk['Partitions']))) width = len(str(len(disk['Partitions'])))
@ -247,29 +242,29 @@ def prep_disk_for_formatting(disk=None):
# Set Display and Warning Strings # Set Display and Warning Strings
if len(disk['Partitions']) == 0: if len(disk['Partitions']) == 0:
disk['Format Warnings'] += 'No partitions found\n' disk['Format Warnings'] += 'No partitions found\n'
for par in disk['Partitions']: for partition in disk['Partitions']:
display = ' Partition {num:>{width}}:\t{size} {fs}'.format( display = ' Partition {num:>{width}}:\t{size} {fs}'.format(
num = par['Number'], num = partition['Number'],
width = width, width = width,
size = par['Size'], size = partition['Size'],
fs = par['FileSystem']) fs = partition['FileSystem'])
if 'Letter' not in par or REGEX_BAD_PARTITION.search(par['FileSystem']): if is_bad_partition(partition):
# Set display string using partition description & OS type # Set display string using partition description & OS type
display += '\t\t{q}{name}{q}\t{desc} ({os})'.format( display += '\t\t{q}{name}{q}\t{desc} ({os})'.format(
display = display, display = display,
q = '"' if par['Name'] != '' else '', q = '"' if partition['Name'] != '' else '',
name = par['Name'], name = partition['Name'],
desc = par['Description'], desc = partition['Description'],
os = par['OS']) os = partition['OS'])
else: else:
# List space used instead of partition description & OS type # List space used instead of partition description & OS type
display += ' (Used: {used})\t{q}{name}{q}'.format( display += ' (Used: {used})\t{q}{name}{q}'.format(
used = par['Used Space'], used = partition['Used Space'],
q = '"' if par['Name'] != '' else '', q = '"' if partition['Name'] != '' else '',
name = par['Name']) name = partition['Name'])
# For all partitions # For all partitions
par['Display String'] = display partition['Display String'] = display
def reassign_volume_letter(letter, new_letter='I'): def reassign_volume_letter(letter, new_letter='I'):
if not letter: if not letter:
@ -281,7 +276,7 @@ def reassign_volume_letter(letter, new_letter='I'):
script.write('select volume {}\n'.format(letter)) script.write('select volume {}\n'.format(letter))
script.write('remove noerr\n') script.write('remove noerr\n')
script.write('assign letter={}\n'.format(new_letter)) script.write('assign letter={}\n'.format(new_letter))
run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT)) run_program(['diskpart', '/s', DISKPART_SCRIPT])
except subprocess.CalledProcessError: except subprocess.CalledProcessError:
pass pass
else: else:
@ -293,12 +288,12 @@ def remove_volume_letters(keep=None):
with open(DISKPART_SCRIPT, 'w') as script: with open(DISKPART_SCRIPT, 'w') as script:
for vol in get_volumes(): for vol in get_volumes():
if vol['Letter'].upper() != keep.upper(): if vol['Letter'].upper() != keep.upper():
script.write('select volume {Number}\n'.format(**vol)) script.write('select volume {}\n'.format(vol['Number']))
script.write('remove noerr\n') script.write('remove noerr\n')
# Run script # Run script
try: try:
run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT)) run_program(['diskpart', '/s', DISKPART_SCRIPT])
except subprocess.CalledProcessError: except subprocess.CalledProcessError:
pass pass
@ -317,9 +312,9 @@ def scan_disks():
# Get partition info for disk # Get partition info for disk
disk['Partitions'] = get_partitions(disk) disk['Partitions'] = get_partitions(disk)
for par in disk['Partitions']: for partition in disk['Partitions']:
# Get partition details # Get partition details
par.update(get_partition_details(disk, par)) partition.update(get_partition_details(disk, partition))
# Done # Done
return disks return disks
@ -331,23 +326,25 @@ def select_disk(title='Which disk?', disks):
for disk in disks: for disk in disks:
display_name = '{Size}\t[{Table}] ({Type}) {Name}'.format(**disk) display_name = '{Size}\t[{Table}] ({Type}) {Name}'.format(**disk)
pwidth=len(str(len(disk['Partitions']))) pwidth=len(str(len(disk['Partitions'])))
for par in disk['Partitions']: for partition in disk['Partitions']:
# Main text # Main text
p_name = 'Partition {num:>{width}}: {size} ({fs})'.format( p_name = 'Partition {num:>{width}}: {size} ({fs})'.format(
num = par['Number'], num = partition['Number'],
width = pwidth, width = pwidth,
size = par['Size'], size = partition['Size'],
fs = par['FileSystem']) fs = partition['FileSystem'])
if par['Name']: if partition['Name']:
p_name += '\t"{}"'.format(par['Name']) p_name += '\t"{}"'.format(partition['Name'])
# Show unsupported partition(s) # Show unsupported partition(s)
if 'Letter' not in par or REGEX_BAD_PARTITION.search(par['FileSystem']): if is_bad_partition(partition):
p_display_name = '{YELLOW}{display}{CLEAR}'.format(display=p_name, **COLORS) p_display_name = '{YELLOW}{display}{CLEAR}'.format(
display=p_name, **COLORS)
display_name += '\n\t\t\t{}'.format(display_name) display_name += '\n\t\t\t{}'.format(display_name)
if not disk['Partitions']: if not disk['Partitions']:
display_name += '\n\t\t\t{YELLOW}No partitions found.{CLEAR}'.format(**COLORS) display_name += '\n\t\t\t{}No partitions found.{}'.format(
COLORS['YELLOW'], COLORS['CLEAR'])
disk_options.append({'Name': display_name, 'Disk': disk}) disk_options.append({'Name': display_name, 'Disk': disk})
actions = [ actions = [