Function separation done.
* Split backup.py into disk & backup * disk.py is for lower level disk management * Renamed functions\data.py's select_disk() to select_volume() * Avoid name collision with functions\disk.py's select_disk() *
This commit is contained in:
Alan Mason
parent
9a3234c822
commit
80cb9b8cea
6 changed files with 372 additions and 365 deletions
|
|
@ -1,18 +1,6 @@
|
|||
# Wizard Kit PE: Functions - Backup
|
||||
|
||||
from functions.common import *
|
||||
import partition_uids
|
||||
|
||||
def assign_volume_letters():
|
||||
try:
|
||||
# Run script
|
||||
with open(DISKPART_SCRIPT, 'w') as script:
|
||||
for vol in get_volumes():
|
||||
script.write('select volume {Number}\n'.format(**vol))
|
||||
script.write('assign\n')
|
||||
run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT))
|
||||
except subprocess.CalledProcessError:
|
||||
pass
|
||||
|
||||
def backup_partition(bin=None, disk=None, par=None):
|
||||
# Bail early
|
||||
|
|
@ -40,232 +28,6 @@ def backup_partition(bin=None, disk=None, par=None):
|
|||
par['Error'] = err.stderr.decode().splitlines()
|
||||
raise BackupError
|
||||
|
||||
def get_attached_disk_info():
|
||||
"""Get details about the attached disks"""
|
||||
disks = []
|
||||
print_info('Getting drive info...')
|
||||
|
||||
# Assign all the letters
|
||||
assign_volume_letters()
|
||||
|
||||
# Get disks
|
||||
disks = get_disks()
|
||||
|
||||
# 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 partition info for disk
|
||||
disk['Partitions'] = get_partitions(disk)
|
||||
|
||||
for par in disk['Partitions']:
|
||||
# Get partition details
|
||||
par.update(get_partition_details(disk, par))
|
||||
|
||||
# Done
|
||||
return disks
|
||||
|
||||
def get_disk_details(disk=None):
|
||||
details = {}
|
||||
|
||||
# Bail early
|
||||
if disk is None:
|
||||
raise Exception('Disk not specified.')
|
||||
|
||||
try:
|
||||
# Run script
|
||||
with open(DISKPART_SCRIPT, 'w') as script:
|
||||
script.write('select disk {Number}\n'.format(**disk))
|
||||
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:
|
||||
pass
|
||||
else:
|
||||
# Remove empty lines
|
||||
tmp = [s.strip() for s in process_return.splitlines() if s.strip() != '']
|
||||
|
||||
# Set disk name
|
||||
details['Name'] = tmp[4]
|
||||
|
||||
# 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]
|
||||
|
||||
# Add key/value pairs to the details variable and return dict
|
||||
details.update({key.strip(): value.strip() for (key, value) in tmp})
|
||||
|
||||
return details
|
||||
|
||||
def get_disks():
|
||||
disks = []
|
||||
|
||||
try:
|
||||
# Run script
|
||||
with open(DISKPART_SCRIPT, 'w') as script:
|
||||
script.write('list disk\n')
|
||||
process_return = run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT))
|
||||
process_return = process_return.stdout.decode().strip()
|
||||
except subprocess.CalledProcessError:
|
||||
pass
|
||||
else:
|
||||
# Append disk numbers
|
||||
for tmp in re.findall(r'Disk (\d+)\s+\w+\s+(\d+\s+\w+)', process_return):
|
||||
_num = tmp[0]
|
||||
_size = human_readable_size(tmp[1])
|
||||
disks.append({'Number': _num, 'Size': _size})
|
||||
|
||||
return disks
|
||||
|
||||
def get_partition_details(disk=None, par=None):
|
||||
details = {}
|
||||
|
||||
# Bail early
|
||||
if disk is None:
|
||||
raise Exception('Disk not specified.')
|
||||
if par is None:
|
||||
raise Exception('Partition not specified.')
|
||||
|
||||
# Diskpart details
|
||||
try:
|
||||
# Run script
|
||||
with open(DISKPART_SCRIPT, 'w') as script:
|
||||
script.write('select disk {Number}\n'.format(**disk))
|
||||
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:
|
||||
pass
|
||||
else:
|
||||
# Get volume letter or RAW status
|
||||
tmp = re.search(r'Volume\s+\d+\s+(\w|RAW)\s+', process_return)
|
||||
if tmp:
|
||||
if tmp.group(1).upper() == 'RAW':
|
||||
details['FileSystem'] = RAW
|
||||
else:
|
||||
details['Letter'] = tmp.group(1)
|
||||
|
||||
# Remove empty lines from process_return
|
||||
tmp = [s.strip() for s in process_return.splitlines() if s.strip() != '']
|
||||
|
||||
# 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]
|
||||
|
||||
# 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.lookup_guid(details['Type'])
|
||||
if guid is not None:
|
||||
details.update({
|
||||
'Description': guid.get('Description', ''),
|
||||
'OS': guid.get('OS', '')})
|
||||
|
||||
if 'Letter' in details:
|
||||
# Disk usage
|
||||
tmp = shutil.disk_usage('{Letter}:\\'.format(**details))
|
||||
details['Used Space'] = human_readable_size(tmp.used)
|
||||
|
||||
# fsutil details
|
||||
try:
|
||||
process_return = run_program('fsutil fsinfo volumeinfo {Letter}:'.format(**details))
|
||||
process_return = process_return.stdout.decode().strip()
|
||||
except subprocess.CalledProcessError:
|
||||
pass
|
||||
else:
|
||||
# Remove empty lines from process_return
|
||||
tmp = [s.strip() for s in process_return.splitlines() if s.strip() != '']
|
||||
|
||||
# Add "Feature" lines
|
||||
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
|
||||
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
|
||||
details['Name'] = details.get('Volume Name', '')
|
||||
|
||||
# Set FileSystem Type
|
||||
if details.get('FileSystem', '') != 'RAW':
|
||||
details['FileSystem'] = details.get('File System Name', 'Unknown')
|
||||
|
||||
return details
|
||||
|
||||
def get_partitions(disk=None):
|
||||
partitions = []
|
||||
|
||||
# Bail early
|
||||
if disk is None:
|
||||
raise Exception('Disk not specified.')
|
||||
|
||||
try:
|
||||
# Run script
|
||||
with open(DISKPART_SCRIPT, 'w') as script:
|
||||
script.write('select disk {Number}\n'.format(**disk))
|
||||
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:
|
||||
pass
|
||||
else:
|
||||
# Append partition numbers
|
||||
for tmp in re.findall(r'Partition\s+(\d+)\s+\w+\s+(\d+\s+\w+)\s+', process_return, re.IGNORECASE):
|
||||
_num = tmp[0]
|
||||
_size = human_readable_size(tmp[1])
|
||||
partitions.append({'Number': _num, 'Size': _size})
|
||||
|
||||
return partitions
|
||||
|
||||
def get_table_type(disk=None):
|
||||
_type = 'Unknown'
|
||||
|
||||
# Bail early
|
||||
if disk is None:
|
||||
raise Exception('Disk not specified.')
|
||||
|
||||
try:
|
||||
with open(DISKPART_SCRIPT, 'w') as script:
|
||||
script.write('select disk {Number}\n'.format(**disk))
|
||||
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:
|
||||
pass
|
||||
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):
|
||||
_type = 'GPT'
|
||||
elif re.findall(r'Disk ID: 00000000', process_return, re.IGNORECASE):
|
||||
_type = 'RAW'
|
||||
elif re.findall(r'Disk ID: [A-Z0-9]+', process_return, re.IGNORECASE):
|
||||
_type = 'MBR'
|
||||
|
||||
return _type
|
||||
|
||||
def get_volumes():
|
||||
vols = []
|
||||
|
||||
try:
|
||||
# Run script
|
||||
with open(DISKPART_SCRIPT, 'w') as script:
|
||||
script.write('list volume\n')
|
||||
process_return = run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT))
|
||||
process_return = process_return.stdout.decode().strip()
|
||||
except subprocess.CalledProcessError:
|
||||
pass
|
||||
else:
|
||||
# Append volume numbers
|
||||
for tmp in re.findall(r'Volume (\d+)\s+([A-Za-z]?)\s+', process_return):
|
||||
vols.append({'Number': tmp[0], 'Letter': tmp[1]})
|
||||
|
||||
return vols
|
||||
|
||||
def prep_disk_for_backup(dest=None, disk=None, ticket_id=None):
|
||||
disk['Backup Warnings'] = '\n'
|
||||
disk['Clobber Risk'] = []
|
||||
|
|
@ -342,54 +104,6 @@ def prep_disk_for_backup(dest=None, disk=None, ticket_id=None):
|
|||
if len(disk['Clobber Risk']) + len(disk['Bad Partitions']) == 1:
|
||||
disk['Backup Warnings'] += '\n{YELLOW}If you continue the partition marked above will NOT be backed up.{CLEAR}\n'.format(**COLORS)
|
||||
|
||||
def prep_disk_for_formatting(disk=None):
|
||||
disk['Format Warnings'] = '\n'
|
||||
width = len(str(len(disk['Partitions'])))
|
||||
|
||||
# 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
|
||||
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 par in disk['Partitions']:
|
||||
if 'Letter' not in par or re.search(r'(RAW|Unknown)', par['FileSystem'], re.IGNORECASE):
|
||||
# FileSystem not accessible to WinPE. List partition type / OS info for technician
|
||||
par['Display String'] = ' Partition {Number:>{width}}:\t{Size} {FileSystem}\t\t{q}{Name}{q}\t{Description} ({OS})'.format(
|
||||
width=width,
|
||||
q='"' if par['Name'] != '' else '',
|
||||
**par)
|
||||
else:
|
||||
# FileSystem accessible to WinPE. List space used instead of partition type / OS info for technician
|
||||
par['Display String'] = ' Partition {Number:>{width}}:\t{Size} {FileSystem} (Used: {Used Space})\t{q}{Name}{q}'.format(
|
||||
width=width,
|
||||
q='"' if par['Name'] != '' else '',
|
||||
**par)
|
||||
|
||||
def remove_volume_letters(keep=''):
|
||||
if keep is None:
|
||||
keep = ''
|
||||
try:
|
||||
# Run script
|
||||
with open(DISKPART_SCRIPT, 'w') as script:
|
||||
for vol in get_volumes():
|
||||
if vol['Letter'].upper() != keep.upper():
|
||||
script.write('select volume {Number}\n'.format(**vol))
|
||||
script.write('remove noerr\n')
|
||||
run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT))
|
||||
except subprocess.CalledProcessError:
|
||||
pass
|
||||
|
||||
def select_backup_destination():
|
||||
# Build menu
|
||||
dests = []
|
||||
|
|
@ -420,47 +134,6 @@ def select_backup_destination():
|
|||
print_warning('No backup destinations found.')
|
||||
return None
|
||||
|
||||
def select_disk(prompt='Which disk?'):
|
||||
"""Select a disk from the attached disks"""
|
||||
disks = get_attached_disk_info()
|
||||
|
||||
# Build menu
|
||||
disk_options = []
|
||||
for disk in disks:
|
||||
display_name = '{Size}\t[{Table}] ({Type}) {Name}'.format(**disk)
|
||||
if len(disk['Partitions']) > 0:
|
||||
pwidth=len(str(len(disk['Partitions'])))
|
||||
for par in disk['Partitions']:
|
||||
# Show unsupported partition(s) in RED
|
||||
par_skip = False
|
||||
if 'Letter' not in par or re.search(r'(RAW|Unknown)', par['FileSystem'], re.IGNORECASE):
|
||||
par_skip = True
|
||||
if par_skip:
|
||||
display_name += COLORS['YELLOW']
|
||||
|
||||
# Main text
|
||||
display_name += '\n\t\t\tPartition {Number:>{pwidth}}: {Size} ({FileSystem})'.format(pwidth=pwidth, **par)
|
||||
if par['Name'] != '':
|
||||
display_name += '\t"{Name}"'.format(**par)
|
||||
|
||||
# Clear color (if set above)
|
||||
if par_skip:
|
||||
display_name += COLORS['CLEAR']
|
||||
else:
|
||||
display_name += '{YELLOW}\n\t\t\tNo partitions found.{CLEAR}'.format(**COLORS)
|
||||
disk_options.append({'Name': display_name, 'Disk': disk})
|
||||
actions = [
|
||||
{'Name': 'Main Menu', 'Letter': 'M'},
|
||||
]
|
||||
|
||||
# Menu loop
|
||||
selection = menu_select(prompt, disk_options, actions)
|
||||
|
||||
if (selection.isnumeric()):
|
||||
return disk_options[int(selection)-1]['Disk']
|
||||
elif (selection == 'M'):
|
||||
abort_to_main_menu()
|
||||
|
||||
def verify_wim_backup(bin=None, par=None):
|
||||
# Bail early
|
||||
if bin is None:
|
||||
|
|
|
|||
|
|
@ -374,7 +374,7 @@ def scan_source_wim(source_wim, dest_path, rel_path=None, interactive=True):
|
|||
|
||||
def select_destination(folder_path, prompt='Select destination'):
|
||||
"""Select destination drive, returns path as string."""
|
||||
disk = select_disk(prompt)
|
||||
disk = select_volume(prompt)
|
||||
if 'fixed' not in disk['Disk'].opts:
|
||||
folder_path = folder_path.replace('\\', '-')
|
||||
path = '{disk}{folder_path}_{Date}'.format(
|
||||
|
|
@ -388,7 +388,7 @@ def select_destination(folder_path, prompt='Select destination'):
|
|||
|
||||
return path
|
||||
|
||||
def select_disk(title='Select disk', auto_select=True):
|
||||
def select_volume(title='Select disk', auto_select=True):
|
||||
"""Select disk from attached disks. returns dict."""
|
||||
actions = [{'Name': 'Quit', 'Letter': 'Q'}]
|
||||
disks = []
|
||||
|
|
|
|||
345
Scripts/functions/disk.py
Normal file
345
Scripts/functions/disk.py
Normal file
|
|
@ -0,0 +1,345 @@
|
|||
# Wizard Kit PE: Functions - Disk
|
||||
|
||||
from functions.common import *
|
||||
import partition_uids
|
||||
|
||||
def assign_volume_letters():
|
||||
try:
|
||||
# Run script
|
||||
with open(DISKPART_SCRIPT, 'w') as script:
|
||||
for vol in get_volumes():
|
||||
script.write('select volume {Number}\n'.format(**vol))
|
||||
script.write('assign\n')
|
||||
run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT))
|
||||
except subprocess.CalledProcessError:
|
||||
pass
|
||||
|
||||
def get_attached_disk_info():
|
||||
"""Get details about the attached disks"""
|
||||
disks = []
|
||||
print_info('Getting drive info...')
|
||||
|
||||
# Assign all the letters
|
||||
assign_volume_letters()
|
||||
|
||||
# Get disks
|
||||
disks = get_disks()
|
||||
|
||||
# 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 partition info for disk
|
||||
disk['Partitions'] = get_partitions(disk)
|
||||
|
||||
for par in disk['Partitions']:
|
||||
# Get partition details
|
||||
par.update(get_partition_details(disk, par))
|
||||
|
||||
# Done
|
||||
return disks
|
||||
|
||||
def get_boot_mode():
|
||||
boot_mode = 'Legacy'
|
||||
try:
|
||||
reg_key = winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, '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
|
||||
|
||||
def get_disk_details(disk=None):
|
||||
details = {}
|
||||
|
||||
# Bail early
|
||||
if disk is None:
|
||||
raise Exception('Disk not specified.')
|
||||
|
||||
try:
|
||||
# Run script
|
||||
with open(DISKPART_SCRIPT, 'w') as script:
|
||||
script.write('select disk {Number}\n'.format(**disk))
|
||||
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:
|
||||
pass
|
||||
else:
|
||||
# Remove empty lines
|
||||
tmp = [s.strip() for s in process_return.splitlines() if s.strip() != '']
|
||||
|
||||
# Set disk name
|
||||
details['Name'] = tmp[4]
|
||||
|
||||
# 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]
|
||||
|
||||
# Add key/value pairs to the details variable and return dict
|
||||
details.update({key.strip(): value.strip() for (key, value) in tmp})
|
||||
|
||||
return details
|
||||
|
||||
def get_disks():
|
||||
disks = []
|
||||
|
||||
try:
|
||||
# Run script
|
||||
with open(DISKPART_SCRIPT, 'w') as script:
|
||||
script.write('list disk\n')
|
||||
process_return = run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT))
|
||||
process_return = process_return.stdout.decode().strip()
|
||||
except subprocess.CalledProcessError:
|
||||
pass
|
||||
else:
|
||||
# Append disk numbers
|
||||
for tmp in re.findall(r'Disk (\d+)\s+\w+\s+(\d+\s+\w+)', process_return):
|
||||
_num = tmp[0]
|
||||
_size = human_readable_size(tmp[1])
|
||||
disks.append({'Number': _num, 'Size': _size})
|
||||
|
||||
return disks
|
||||
|
||||
def get_partition_details(disk=None, par=None):
|
||||
details = {}
|
||||
|
||||
# Bail early
|
||||
if disk is None:
|
||||
raise Exception('Disk not specified.')
|
||||
if par is None:
|
||||
raise Exception('Partition not specified.')
|
||||
|
||||
# Diskpart details
|
||||
try:
|
||||
# Run script
|
||||
with open(DISKPART_SCRIPT, 'w') as script:
|
||||
script.write('select disk {Number}\n'.format(**disk))
|
||||
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:
|
||||
pass
|
||||
else:
|
||||
# Get volume letter or RAW status
|
||||
tmp = re.search(r'Volume\s+\d+\s+(\w|RAW)\s+', process_return)
|
||||
if tmp:
|
||||
if tmp.group(1).upper() == 'RAW':
|
||||
details['FileSystem'] = RAW
|
||||
else:
|
||||
details['Letter'] = tmp.group(1)
|
||||
|
||||
# Remove empty lines from process_return
|
||||
tmp = [s.strip() for s in process_return.splitlines() if s.strip() != '']
|
||||
|
||||
# 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]
|
||||
|
||||
# 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.lookup_guid(details['Type'])
|
||||
if guid is not None:
|
||||
details.update({
|
||||
'Description': guid.get('Description', ''),
|
||||
'OS': guid.get('OS', '')})
|
||||
|
||||
if 'Letter' in details:
|
||||
# Disk usage
|
||||
tmp = shutil.disk_usage('{Letter}:\\'.format(**details))
|
||||
details['Used Space'] = human_readable_size(tmp.used)
|
||||
|
||||
# fsutil details
|
||||
try:
|
||||
process_return = run_program('fsutil fsinfo volumeinfo {Letter}:'.format(**details))
|
||||
process_return = process_return.stdout.decode().strip()
|
||||
except subprocess.CalledProcessError:
|
||||
pass
|
||||
else:
|
||||
# Remove empty lines from process_return
|
||||
tmp = [s.strip() for s in process_return.splitlines() if s.strip() != '']
|
||||
|
||||
# Add "Feature" lines
|
||||
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
|
||||
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
|
||||
details['Name'] = details.get('Volume Name', '')
|
||||
|
||||
# Set FileSystem Type
|
||||
if details.get('FileSystem', '') != 'RAW':
|
||||
details['FileSystem'] = details.get('File System Name', 'Unknown')
|
||||
|
||||
return details
|
||||
|
||||
def get_partitions(disk=None):
|
||||
partitions = []
|
||||
|
||||
# Bail early
|
||||
if disk is None:
|
||||
raise Exception('Disk not specified.')
|
||||
|
||||
try:
|
||||
# Run script
|
||||
with open(DISKPART_SCRIPT, 'w') as script:
|
||||
script.write('select disk {Number}\n'.format(**disk))
|
||||
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:
|
||||
pass
|
||||
else:
|
||||
# Append partition numbers
|
||||
for tmp in re.findall(r'Partition\s+(\d+)\s+\w+\s+(\d+\s+\w+)\s+', process_return, re.IGNORECASE):
|
||||
_num = tmp[0]
|
||||
_size = human_readable_size(tmp[1])
|
||||
partitions.append({'Number': _num, 'Size': _size})
|
||||
|
||||
return partitions
|
||||
|
||||
def get_table_type(disk=None):
|
||||
_type = 'Unknown'
|
||||
|
||||
# Bail early
|
||||
if disk is None:
|
||||
raise Exception('Disk not specified.')
|
||||
|
||||
try:
|
||||
with open(DISKPART_SCRIPT, 'w') as script:
|
||||
script.write('select disk {Number}\n'.format(**disk))
|
||||
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:
|
||||
pass
|
||||
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):
|
||||
_type = 'GPT'
|
||||
elif re.findall(r'Disk ID: 00000000', process_return, re.IGNORECASE):
|
||||
_type = 'RAW'
|
||||
elif re.findall(r'Disk ID: [A-Z0-9]+', process_return, re.IGNORECASE):
|
||||
_type = 'MBR'
|
||||
|
||||
return _type
|
||||
|
||||
def get_volumes():
|
||||
vols = []
|
||||
|
||||
try:
|
||||
# Run script
|
||||
with open(DISKPART_SCRIPT, 'w') as script:
|
||||
script.write('list volume\n')
|
||||
process_return = run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT))
|
||||
process_return = process_return.stdout.decode().strip()
|
||||
except subprocess.CalledProcessError:
|
||||
pass
|
||||
else:
|
||||
# Append volume numbers
|
||||
for tmp in re.findall(r'Volume (\d+)\s+([A-Za-z]?)\s+', process_return):
|
||||
vols.append({'Number': tmp[0], 'Letter': tmp[1]})
|
||||
|
||||
return vols
|
||||
|
||||
def prep_disk_for_formatting(disk=None):
|
||||
disk['Format Warnings'] = '\n'
|
||||
width = len(str(len(disk['Partitions'])))
|
||||
|
||||
# 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
|
||||
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 par in disk['Partitions']:
|
||||
if 'Letter' not in par or re.search(r'(RAW|Unknown)', par['FileSystem'], re.IGNORECASE):
|
||||
# FileSystem not accessible to WinPE. List partition type / OS info for technician
|
||||
par['Display String'] = ' Partition {Number:>{width}}:\t{Size} {FileSystem}\t\t{q}{Name}{q}\t{Description} ({OS})'.format(
|
||||
width=width,
|
||||
q='"' if par['Name'] != '' else '',
|
||||
**par)
|
||||
else:
|
||||
# FileSystem accessible to WinPE. List space used instead of partition type / OS info for technician
|
||||
par['Display String'] = ' Partition {Number:>{width}}:\t{Size} {FileSystem} (Used: {Used Space})\t{q}{Name}{q}'.format(
|
||||
width=width,
|
||||
q='"' if par['Name'] != '' else '',
|
||||
**par)
|
||||
|
||||
def remove_volume_letters(keep=''):
|
||||
if keep is None:
|
||||
keep = ''
|
||||
try:
|
||||
# Run script
|
||||
with open(DISKPART_SCRIPT, 'w') as script:
|
||||
for vol in get_volumes():
|
||||
if vol['Letter'].upper() != keep.upper():
|
||||
script.write('select volume {Number}\n'.format(**vol))
|
||||
script.write('remove noerr\n')
|
||||
run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT))
|
||||
except subprocess.CalledProcessError:
|
||||
pass
|
||||
|
||||
def select_disk(prompt='Which disk?'):
|
||||
"""Select a disk from the attached disks"""
|
||||
disks = get_attached_disk_info()
|
||||
|
||||
# Build menu
|
||||
disk_options = []
|
||||
for disk in disks:
|
||||
display_name = '{Size}\t[{Table}] ({Type}) {Name}'.format(**disk)
|
||||
if len(disk['Partitions']) > 0:
|
||||
pwidth=len(str(len(disk['Partitions'])))
|
||||
for par in disk['Partitions']:
|
||||
# Show unsupported partition(s) in RED
|
||||
par_skip = False
|
||||
if 'Letter' not in par or re.search(r'(RAW|Unknown)', par['FileSystem'], re.IGNORECASE):
|
||||
par_skip = True
|
||||
if par_skip:
|
||||
display_name += COLORS['YELLOW']
|
||||
|
||||
# Main text
|
||||
display_name += '\n\t\t\tPartition {Number:>{pwidth}}: {Size} ({FileSystem})'.format(pwidth=pwidth, **par)
|
||||
if par['Name'] != '':
|
||||
display_name += '\t"{Name}"'.format(**par)
|
||||
|
||||
# Clear color (if set above)
|
||||
if par_skip:
|
||||
display_name += COLORS['CLEAR']
|
||||
else:
|
||||
display_name += '{YELLOW}\n\t\t\tNo partitions found.{CLEAR}'.format(**COLORS)
|
||||
disk_options.append({'Name': display_name, 'Disk': disk})
|
||||
actions = [
|
||||
{'Name': 'Main Menu', 'Letter': 'M'},
|
||||
]
|
||||
|
||||
# Menu loop
|
||||
selection = menu_select(prompt, disk_options, actions)
|
||||
|
||||
if (selection.isnumeric()):
|
||||
return disk_options[int(selection)-1]['Disk']
|
||||
elif (selection == 'M'):
|
||||
abort_to_main_menu()
|
||||
|
||||
if __name__ == '__main__':
|
||||
print("This file is not meant to be called directly.")
|
||||
|
|
@ -1,9 +1,9 @@
|
|||
# Wizard Kit PE: Functions - Windows Setup
|
||||
|
||||
from functions.common import *
|
||||
from functions.data import *
|
||||
|
||||
# STATIC VARIABLES
|
||||
DISKPART_SCRIPT = '{tmp}\\diskpart.script'.format(tmp=os.environ['TMP'])
|
||||
DISKPART_SCRIPT = r'{}\diskpart.script'.format(global_vars['Env']['TMP'])
|
||||
WINDOWS_VERSIONS = [
|
||||
{'Name': 'Windows 7 Home Basic',
|
||||
'Image File': 'Win7',
|
||||
|
|
@ -43,24 +43,6 @@ WINDOWS_VERSIONS = [
|
|||
'Family': '10'},
|
||||
]
|
||||
|
||||
def is_index_in_image(bin=None, filename=None, imagename=None):
|
||||
# Bail early
|
||||
if bin is None:
|
||||
raise Exception('bin not specified.')
|
||||
if filename is None:
|
||||
raise Exception('Filename not specified.')
|
||||
if imagename is None:
|
||||
raise Exception('Image Name not specified.')
|
||||
|
||||
cmd = '{bin}\\wimlib\\wimlib-imagex info "{filename}" "{imagename}"'.format(bin=bin, filename=filename, imagename=imagename)
|
||||
try:
|
||||
run_program(cmd)
|
||||
except subprocess.CalledProcessError:
|
||||
print_error('Invalid image: {filename}'.format(filename=filename))
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def find_windows_image(bin, windows_version):
|
||||
"""Search for a Windows source image file on local drives and network drives (in that order)"""
|
||||
image = {}
|
||||
|
|
@ -73,7 +55,7 @@ def find_windows_image(bin, windows_version):
|
|||
filename = '{drive}:\\images\\{imagefile}'.format(drive=tmp[0], imagefile=imagefile)
|
||||
filename_ext = '{filename}.{ext}'.format(filename=filename, ext=ext)
|
||||
if os.path.isfile(filename_ext):
|
||||
if is_index_in_image(bin, filename_ext, windows_version['Image Name']):
|
||||
if wim_contains_image(bin, filename_ext, windows_version['Image Name']):
|
||||
image['Ext'] = ext
|
||||
image['File'] = filename
|
||||
image['Glob'] = '--ref="{File}*.swm"'.format(**image) if ext == 'swm' else ''
|
||||
|
|
@ -89,7 +71,7 @@ def find_windows_image(bin, windows_version):
|
|||
filename = '\\\\{IP}\\{Share}\\images\\{imagefile}'.format(imagefile=imagefile, **WINDOWS_SERVER)
|
||||
filename_ext = '{filename}.{ext}'.format(filename=filename, ext=ext)
|
||||
if os.path.isfile(filename_ext):
|
||||
if is_index_in_image(bin, filename_ext, windows_version['Image Name']):
|
||||
if wim_contains_image(bin, filename_ext, windows_version['Image Name']):
|
||||
image['Ext'] = ext
|
||||
image['File'] = filename
|
||||
image['Glob'] = '--ref="{File}*.swm"'.format(**image) if ext == 'swm' else ''
|
||||
|
|
@ -186,18 +168,6 @@ def format_mbr(disk=None, windows_family=None):
|
|||
run_program('diskpart /s {script}'.format(script=DISKPART_SCRIPT))
|
||||
time.sleep(2)
|
||||
|
||||
def get_boot_mode():
|
||||
boot_mode = 'Legacy'
|
||||
try:
|
||||
reg_key = winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, '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
|
||||
|
||||
def mount_windows_share():
|
||||
"""Mount the Windows images share unless labeled as already mounted."""
|
||||
if WINDOWS_SERVER['Mounted']:
|
||||
|
|
@ -253,5 +223,23 @@ def setup_windows_re(windows_version=None, windows_letter='W', tools_letter='T')
|
|||
def update_boot_partition(system_letter='S', windows_letter='W', mode='ALL'):
|
||||
run_program('bcdboot {win}:\\Windows /s {sys}: /f {mode}'.format(win=windows_letter, sys=system_letter, mode=mode))
|
||||
|
||||
def wim_contains_image(bin=None, filename=None, imagename=None):
|
||||
# Bail early
|
||||
if bin is None:
|
||||
raise Exception('bin not specified.')
|
||||
if filename is None:
|
||||
raise Exception('Filename not specified.')
|
||||
if imagename is None:
|
||||
raise Exception('Image Name not specified.')
|
||||
|
||||
cmd = '{bin}\\wimlib\\wimlib-imagex info "{filename}" "{imagename}"'.format(bin=bin, filename=filename, imagename=imagename)
|
||||
try:
|
||||
run_program(cmd)
|
||||
except subprocess.CalledProcessError:
|
||||
print_error('Invalid image: {filename}'.format(filename=filename))
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
if __name__ == '__main__':
|
||||
print("This file is not meant to be called directly.")
|
||||
|
|
|
|||
|
|
@ -1,6 +1,8 @@
|
|||
# Wizard Kit PE: Menus
|
||||
|
||||
from functions.data import *
|
||||
from functions.backup import *
|
||||
from functions.disk import *
|
||||
from functions.windows_setup import *
|
||||
|
||||
# STATIC VARIABLES
|
||||
FAST_COPY_ARGS = [
|
||||
|
|
|
|||
|
|
@ -6,7 +6,6 @@ import sys
|
|||
# Init
|
||||
os.chdir(os.path.dirname(os.path.realpath(__file__)))
|
||||
sys.path.append(os.getcwd())
|
||||
from functions.common import *
|
||||
from functions.winpe_menus import *
|
||||
init_global_vars()
|
||||
os.system('title {}: Root Menu'.format(KIT_NAME_FULL))
|
||||
|
|
|
|||
Loading…
Reference in a new issue