Updated disk.py
This commit is contained in:
parent
07e43307c5
commit
72eac47524
GPG key ID:
152FAC923B0E132C
1 changed files with 319 additions and 317 deletions
|
|
@ -6,390 +6,392 @@ from settings.partition_uids import *
|
||||||
# Regex
|
# Regex
|
||||||
REGEX_BAD_PARTITION = re.compile(r'(RAW|Unknown)', re.IGNORECASE)
|
REGEX_BAD_PARTITION = re.compile(r'(RAW|Unknown)', re.IGNORECASE)
|
||||||
REGEX_DISK_GPT = re.compile(
|
REGEX_DISK_GPT = re.compile(
|
||||||
r'Disk ID: {[A-Z0-9]+-[A-Z0-9]+-[A-Z0-9]+-[A-Z0-9]+-[A-Z0-9]+}',
|
r'Disk ID: {[A-Z0-9]+-[A-Z0-9]+-[A-Z0-9]+-[A-Z0-9]+-[A-Z0-9]+}',
|
||||||
re.IGNORECASE)
|
re.IGNORECASE)
|
||||||
REGEX_DISK_MBR = re.compile(r'Disk ID: [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)
|
REGEX_DISK_RAW = re.compile(r'Disk ID: 00000000', re.IGNORECASE)
|
||||||
|
|
||||||
def assign_volume_letters():
|
def assign_volume_letters():
|
||||||
"""Assign a volume letter to all available volumes."""
|
"""Assign a volume letter to all available volumes."""
|
||||||
remove_volume_letters()
|
remove_volume_letters()
|
||||||
|
|
||||||
# Write script
|
# Write script
|
||||||
script = []
|
script = []
|
||||||
for vol in get_volumes():
|
for vol in get_volumes():
|
||||||
script.append('select volume {}'.format(vol['Number']))
|
script.append('select volume {}'.format(vol['Number']))
|
||||||
script.append('assign')
|
script.append('assign')
|
||||||
|
|
||||||
# Run
|
# Run
|
||||||
run_diskpart(script)
|
run_diskpart(script)
|
||||||
|
|
||||||
def get_boot_mode():
|
def get_boot_mode():
|
||||||
"""Check if the boot mode was UEFI or legacy."""
|
"""Check if the boot mode was UEFI or legacy."""
|
||||||
boot_mode = 'Legacy'
|
boot_mode = 'Legacy'
|
||||||
try:
|
try:
|
||||||
reg_key = winreg.OpenKey(
|
reg_key = winreg.OpenKey(
|
||||||
winreg.HKEY_LOCAL_MACHINE, r'System\CurrentControlSet\Control')
|
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'
|
||||||
except:
|
except:
|
||||||
boot_mode = 'Unknown'
|
boot_mode = 'Unknown'
|
||||||
|
|
||||||
return boot_mode
|
return boot_mode
|
||||||
|
|
||||||
def get_disk_details(disk):
|
def get_disk_details(disk):
|
||||||
"""Get disk details using DiskPart."""
|
"""Get disk details using DiskPart."""
|
||||||
details = {}
|
details = {}
|
||||||
script = [
|
script = [
|
||||||
'select disk {}'.format(disk['Number']),
|
'select disk {}'.format(disk['Number']),
|
||||||
'detail disk']
|
'detail disk']
|
||||||
|
|
||||||
# Run
|
# Run
|
||||||
try:
|
try:
|
||||||
result = run_diskpart(script)
|
result = run_diskpart(script)
|
||||||
except subprocess.CalledProcessError:
|
except subprocess.CalledProcessError:
|
||||||
pass
|
pass
|
||||||
else:
|
else:
|
||||||
output = result.stdout.decode().strip()
|
output = result.stdout.decode().strip()
|
||||||
# Remove empty lines
|
# Remove empty lines
|
||||||
tmp = [s.strip() for s in output.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 ':'
|
# 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})
|
||||||
|
|
||||||
return details
|
return details
|
||||||
|
|
||||||
def get_disks():
|
def get_disks():
|
||||||
"""Get list of attached disks using DiskPart."""
|
"""Get list of attached disks using DiskPart."""
|
||||||
disks = []
|
disks = []
|
||||||
|
|
||||||
try:
|
try:
|
||||||
# Run script
|
# Run script
|
||||||
result = run_diskpart(['list disk'])
|
result = run_diskpart(['list disk'])
|
||||||
except subprocess.CalledProcessError:
|
except subprocess.CalledProcessError:
|
||||||
pass
|
pass
|
||||||
else:
|
else:
|
||||||
# Append disk numbers
|
# Append disk numbers
|
||||||
output = result.stdout.decode().strip()
|
output = result.stdout.decode().strip()
|
||||||
for tmp in re.findall(r'Disk (\d+)\s+\w+\s+(\d+\s+\w+)', output):
|
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, partition):
|
def get_partition_details(disk, partition):
|
||||||
"""Get partition details using DiskPart and fsutil."""
|
"""Get partition details using DiskPart and fsutil."""
|
||||||
details = {}
|
details = {}
|
||||||
script = [
|
script = [
|
||||||
'select disk {}'.format(disk['Number']),
|
'select disk {}'.format(disk['Number']),
|
||||||
'select partition {}'.format(partition['Number']),
|
'select partition {}'.format(partition['Number']),
|
||||||
'detail partition']
|
'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:
|
try:
|
||||||
# Run script
|
tmp = psutil.disk_usage('{}:\\'.format(details['Letter']))
|
||||||
result = run_diskpart(script)
|
except OSError as err:
|
||||||
except subprocess.CalledProcessError:
|
details['FileSystem'] = 'Unknown'
|
||||||
pass
|
details['Error'] = err.strerror
|
||||||
else:
|
else:
|
||||||
# Get volume letter or RAW status
|
details['Used Space'] = human_readable_size(tmp.used)
|
||||||
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
|
# fsutil details
|
||||||
guid = PARTITION_UIDS.get(details.get('Type').upper(), {})
|
cmd = [
|
||||||
if guid:
|
'fsutil',
|
||||||
details.update({
|
'fsinfo',
|
||||||
'Description': guid.get('Description', '')[:29],
|
'volumeinfo',
|
||||||
'OS': guid.get('OS', 'Unknown')[:27]})
|
'{}:'.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:
|
# Set Volume Name
|
||||||
# Disk usage
|
details['Name'] = details.get('Volume Name', '')
|
||||||
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)
|
|
||||||
|
|
||||||
# fsutil details
|
# Set FileSystem Type
|
||||||
cmd = [
|
if details.get('FileSystem', '') not in ['RAW', 'Unknown']:
|
||||||
'fsutil',
|
details['FileSystem'] = details.get('File System Name', 'Unknown')
|
||||||
'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 Volume Name
|
return details
|
||||||
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
|
|
||||||
|
|
||||||
def get_partitions(disk):
|
def get_partitions(disk):
|
||||||
"""Get list of partition using DiskPart."""
|
"""Get list of partition using DiskPart."""
|
||||||
partitions = []
|
partitions = []
|
||||||
script = [
|
script = [
|
||||||
'select disk {}'.format(disk['Number']),
|
'select disk {}'.format(disk['Number']),
|
||||||
'list partition']
|
'list partition']
|
||||||
|
|
||||||
try:
|
try:
|
||||||
# Run script
|
# Run script
|
||||||
result = run_diskpart(script)
|
result = run_diskpart(script)
|
||||||
except subprocess.CalledProcessError:
|
except subprocess.CalledProcessError:
|
||||||
pass
|
pass
|
||||||
else:
|
else:
|
||||||
# Append partition numbers
|
# Append partition numbers
|
||||||
output = result.stdout.decode().strip()
|
output = result.stdout.decode().strip()
|
||||||
regex = r'Partition\s+(\d+)\s+\w+\s+(\d+\s+\w+)\s+'
|
regex = r'Partition\s+(\d+)\s+\w+\s+(\d+\s+\w+)\s+'
|
||||||
for tmp in re.findall(regex, output, re.IGNORECASE):
|
for tmp in re.findall(regex, output, re.IGNORECASE):
|
||||||
num = tmp[0]
|
num = tmp[0]
|
||||||
size = human_readable_size(tmp[1])
|
size = human_readable_size(tmp[1])
|
||||||
partitions.append({'Number': num, 'Size': size})
|
partitions.append({'Number': num, 'Size': size})
|
||||||
|
|
||||||
return partitions
|
return partitions
|
||||||
|
|
||||||
def get_table_type(disk):
|
def get_table_type(disk):
|
||||||
"""Get disk partition table type using DiskPart."""
|
"""Get disk partition table type using DiskPart."""
|
||||||
part_type = 'Unknown'
|
part_type = 'Unknown'
|
||||||
script = [
|
script = [
|
||||||
'select disk {}'.format(disk['Number']),
|
'select disk {}'.format(disk['Number']),
|
||||||
'uniqueid disk']
|
'uniqueid disk']
|
||||||
|
|
||||||
try:
|
try:
|
||||||
result = run_diskpart(script)
|
result = run_diskpart(script)
|
||||||
except subprocess.CalledProcessError:
|
except subprocess.CalledProcessError:
|
||||||
pass
|
pass
|
||||||
else:
|
else:
|
||||||
output = result.stdout.decode().strip()
|
output = result.stdout.decode().strip()
|
||||||
if REGEX_DISK_GPT.search(output):
|
if REGEX_DISK_GPT.search(output):
|
||||||
part_type = 'GPT'
|
part_type = 'GPT'
|
||||||
elif REGEX_DISK_MBR.search(output):
|
elif REGEX_DISK_MBR.search(output):
|
||||||
part_type = 'MBR'
|
part_type = 'MBR'
|
||||||
elif REGEX_DISK_RAW.search(output):
|
elif REGEX_DISK_RAW.search(output):
|
||||||
part_type = 'RAW'
|
part_type = 'RAW'
|
||||||
|
|
||||||
return part_type
|
return part_type
|
||||||
|
|
||||||
def get_volumes():
|
def get_volumes():
|
||||||
"""Get list of volumes using DiskPart."""
|
"""Get list of volumes using DiskPart."""
|
||||||
vols = []
|
vols = []
|
||||||
try:
|
try:
|
||||||
result = run_diskpart(['list volume'])
|
result = run_diskpart(['list volume'])
|
||||||
except subprocess.CalledProcessError:
|
except subprocess.CalledProcessError:
|
||||||
pass
|
pass
|
||||||
else:
|
else:
|
||||||
# Append volume numbers
|
# Append volume numbers
|
||||||
output = result.stdout.decode().strip()
|
output = result.stdout.decode().strip()
|
||||||
for tmp in re.findall(r'Volume (\d+)\s+([A-Za-z]?)\s+', output):
|
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):
|
def is_bad_partition(par):
|
||||||
"""Check if the partition is accessible."""
|
"""Check if the partition is accessible."""
|
||||||
return 'Letter' not in par or REGEX_BAD_PARTITION.search(par['FileSystem'])
|
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):
|
||||||
"""Gather details about the disk and its partitions."""
|
"""Gather details about the disk and its partitions."""
|
||||||
disk['Format Warnings'] = '\n'
|
disk['Format Warnings'] = '\n'
|
||||||
width = len(str(len(disk['Partitions'])))
|
width = len(str(len(disk['Partitions'])))
|
||||||
|
|
||||||
# Bail early
|
# Bail early
|
||||||
if disk is None:
|
if disk is None:
|
||||||
raise Exception('Disk not provided.')
|
raise Exception('Disk not provided.')
|
||||||
|
|
||||||
# Set boot method and partition table type
|
# Set boot method and partition table type
|
||||||
disk['Use GPT'] = True
|
disk['Use GPT'] = True
|
||||||
if (get_boot_mode() == 'UEFI'):
|
if (get_boot_mode() == 'UEFI'):
|
||||||
if (not ask("Setup Windows to use UEFI booting?")):
|
if (not ask("Setup Windows to use UEFI booting?")):
|
||||||
disk['Use GPT'] = False
|
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:
|
else:
|
||||||
if (ask("Setup Windows to use BIOS/Legacy booting?")):
|
# List space used instead of partition description & OS type
|
||||||
disk['Use GPT'] = False
|
display += ' (Used: {used})\t{q}{name}{q}'.format(
|
||||||
|
used = partition['Used Space'],
|
||||||
# Set Display and Warning Strings
|
q = '"' if partition['Name'] != '' else '',
|
||||||
if len(disk['Partitions']) == 0:
|
name = partition['Name'])
|
||||||
disk['Format Warnings'] += 'No partitions found\n'
|
# For all partitions
|
||||||
for partition in disk['Partitions']:
|
partition['Display String'] = display
|
||||||
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
|
|
||||||
|
|
||||||
def reassign_volume_letter(letter, new_letter='I'):
|
def reassign_volume_letter(letter, new_letter='I'):
|
||||||
"""Assign a new letter to a volume using DiskPart."""
|
"""Assign a new letter to a volume using DiskPart."""
|
||||||
if not letter:
|
if not letter:
|
||||||
# Ignore
|
# Ignore
|
||||||
return None
|
return None
|
||||||
script = [
|
script = [
|
||||||
'select volume {}'.format(letter),
|
'select volume {}'.format(letter),
|
||||||
'remove noerr',
|
'remove noerr',
|
||||||
'assign letter={}'.format(new_letter)]
|
'assign letter={}'.format(new_letter)]
|
||||||
try:
|
try:
|
||||||
run_diskpart(script)
|
run_diskpart(script)
|
||||||
except subprocess.CalledProcessError:
|
except subprocess.CalledProcessError:
|
||||||
pass
|
pass
|
||||||
else:
|
else:
|
||||||
return new_letter
|
return new_letter
|
||||||
|
|
||||||
def remove_volume_letters(keep=None):
|
def remove_volume_letters(keep=None):
|
||||||
"""Remove all assigned volume letters using DiskPart."""
|
"""Remove all assigned volume letters using DiskPart."""
|
||||||
if not keep:
|
if not keep:
|
||||||
keep = ''
|
keep = ''
|
||||||
|
|
||||||
script = []
|
script = []
|
||||||
for vol in get_volumes():
|
for vol in get_volumes():
|
||||||
if vol['Letter'].upper() != keep.upper():
|
if vol['Letter'].upper() != keep.upper():
|
||||||
script.append('select volume {}'.format(vol['Number']))
|
script.append('select volume {}'.format(vol['Number']))
|
||||||
script.append('remove noerr')
|
script.append('remove noerr')
|
||||||
|
|
||||||
# Run script
|
# Run script
|
||||||
try:
|
try:
|
||||||
run_diskpart(script)
|
run_diskpart(script)
|
||||||
except subprocess.CalledProcessError:
|
except subprocess.CalledProcessError:
|
||||||
pass
|
pass
|
||||||
|
|
||||||
def run_diskpart(script):
|
def run_diskpart(script):
|
||||||
"""Run DiskPart script."""
|
"""Run DiskPart script."""
|
||||||
tempfile = r'{}\diskpart.script'.format(global_vars['Env']['TMP'])
|
tempfile = r'{}\diskpart.script'.format(global_vars['Env']['TMP'])
|
||||||
|
|
||||||
# Write script
|
# Write script
|
||||||
with open(tempfile, 'w') as f:
|
with open(tempfile, 'w') as f:
|
||||||
for line in script:
|
for line in script:
|
||||||
f.write('{}\n'.format(line))
|
f.write('{}\n'.format(line))
|
||||||
|
|
||||||
# Run script
|
# Run script
|
||||||
cmd = [
|
cmd = [
|
||||||
r'{}\Windows\System32\diskpart.exe'.format(
|
r'{}\Windows\System32\diskpart.exe'.format(
|
||||||
global_vars['Env']['SYSTEMDRIVE']),
|
global_vars['Env']['SYSTEMDRIVE']),
|
||||||
'/s', tempfile]
|
'/s', tempfile]
|
||||||
result = run_program(cmd)
|
result = run_program(cmd)
|
||||||
sleep(2)
|
sleep(2)
|
||||||
return result
|
return result
|
||||||
|
|
||||||
def scan_disks():
|
def scan_disks():
|
||||||
"""Get details about the attached disks"""
|
"""Get details about the attached disks"""
|
||||||
disks = get_disks()
|
disks = get_disks()
|
||||||
|
|
||||||
# Get disk details
|
# Get disk details
|
||||||
for disk in disks:
|
for disk in disks:
|
||||||
# Get partition style
|
# Get partition style
|
||||||
disk['Table'] = get_table_type(disk)
|
disk['Table'] = get_table_type(disk)
|
||||||
|
|
||||||
# Get disk name/model and physical details
|
# Get disk name/model and physical details
|
||||||
disk.update(get_disk_details(disk))
|
disk.update(get_disk_details(disk))
|
||||||
|
|
||||||
# Get partition info for disk
|
# Get partition info for disk
|
||||||
disk['Partitions'] = get_partitions(disk)
|
disk['Partitions'] = get_partitions(disk)
|
||||||
|
|
||||||
for partition in disk['Partitions']:
|
for partition in disk['Partitions']:
|
||||||
# Get partition details
|
# Get partition details
|
||||||
partition.update(get_partition_details(disk, partition))
|
partition.update(get_partition_details(disk, partition))
|
||||||
|
|
||||||
# Done
|
# Done
|
||||||
return disks
|
return disks
|
||||||
|
|
||||||
def select_disk(title='Which disk?', disks=[]):
|
def select_disk(title='Which disk?', disks=[]):
|
||||||
"""Select a disk from the attached disks"""
|
"""Select a disk from the attached disks"""
|
||||||
# Build menu
|
# Build menu
|
||||||
disk_options = []
|
disk_options = []
|
||||||
for disk in disks:
|
for disk in disks:
|
||||||
display_name = '{}\t[{}] ({}) {}'.format(
|
display_name = '{}\t[{}] ({}) {}'.format(
|
||||||
disk.get('Size', ''),
|
disk.get('Size', ''),
|
||||||
disk.get('Table', ''),
|
disk.get('Table', ''),
|
||||||
disk.get('Type', ''),
|
disk.get('Type', ''),
|
||||||
disk.get('Name', 'Unknown'),
|
disk.get('Name', 'Unknown'),
|
||||||
)
|
)
|
||||||
pwidth=len(str(len(disk['Partitions'])))
|
pwidth=len(str(len(disk['Partitions'])))
|
||||||
for partition 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 = partition['Number'],
|
num = partition['Number'],
|
||||||
width = pwidth,
|
width = pwidth,
|
||||||
size = partition['Size'],
|
size = partition['Size'],
|
||||||
fs = partition['FileSystem'])
|
fs = partition['FileSystem'])
|
||||||
if partition['Name']:
|
if partition['Name']:
|
||||||
p_name += '\t"{}"'.format(partition['Name'])
|
p_name += '\t"{}"'.format(partition['Name'])
|
||||||
|
|
||||||
# Show unsupported partition(s)
|
# Show unsupported partition(s)
|
||||||
if is_bad_partition(partition):
|
if is_bad_partition(partition):
|
||||||
p_name = '{YELLOW}{p_name}{CLEAR}'.format(
|
p_name = '{YELLOW}{p_name}{CLEAR}'.format(
|
||||||
p_name=p_name, **COLORS)
|
p_name=p_name, **COLORS)
|
||||||
|
|
||||||
display_name += '\n\t\t\t{}'.format(p_name)
|
display_name += '\n\t\t\t{}'.format(p_name)
|
||||||
if not disk['Partitions']:
|
if not disk['Partitions']:
|
||||||
display_name += '\n\t\t\t{}No partitions found.{}'.format(
|
display_name += '\n\t\t\t{}No partitions found.{}'.format(
|
||||||
COLORS['YELLOW'], COLORS['CLEAR'])
|
COLORS['YELLOW'], COLORS['CLEAR'])
|
||||||
|
|
||||||
disk_options.append({'Name': display_name, 'Disk': disk})
|
disk_options.append({'Name': display_name, 'Disk': disk})
|
||||||
actions = [
|
actions = [
|
||||||
{'Name': 'Main Menu', 'Letter': 'M'},
|
{'Name': 'Main Menu', 'Letter': 'M'},
|
||||||
]
|
]
|
||||||
|
|
||||||
# Menu loop
|
# Menu loop
|
||||||
selection = menu_select(
|
selection = menu_select(
|
||||||
title = title,
|
title = title,
|
||||||
main_entries = disk_options,
|
main_entries = disk_options,
|
||||||
action_entries = actions)
|
action_entries = actions)
|
||||||
|
|
||||||
if (selection.isnumeric()):
|
if (selection.isnumeric()):
|
||||||
return disk_options[int(selection)-1]['Disk']
|
return disk_options[int(selection)-1]['Disk']
|
||||||
elif (selection == 'M'):
|
elif (selection == 'M'):
|
||||||
raise GenericAbort
|
raise GenericAbort
|
||||||
|
|
||||||
if __name__ == '__main__':
|
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
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue