Updated winpe_menus.py & cleanup

* Removed "Windows Family" logic
  * Win7 will now have recovery tools setup
* Added format_disk()
This commit is contained in:
Alan Mason 2017-12-01 11:05:40 -08:00
parent deb7c76ffb
commit 90c4189942
3 changed files with 125 additions and 127 deletions

View file

@ -360,7 +360,7 @@ def select_disk(title='Which disk?', disks):
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'):
abort_to_main_menu() raise GeneralAbort
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.")

View file

@ -7,46 +7,38 @@ DISKPART_SCRIPT = r'{}\diskpart.script'.format(global_vars['Env']['TMP'])
WINDOWS_VERSIONS = [ WINDOWS_VERSIONS = [
{'Name': 'Windows 7 Home Basic', {'Name': 'Windows 7 Home Basic',
'Image File': 'Win7', 'Image File': 'Win7',
'Image Name': 'Windows 7 HOMEBASIC', 'Image Name': 'Windows 7 HOMEBASIC'},
'Family': '7'},
{'Name': 'Windows 7 Home Premium', {'Name': 'Windows 7 Home Premium',
'Image File': 'Win7', 'Image File': 'Win7',
'Image Name': 'Windows 7 HOMEPREMIUM', 'Image Name': 'Windows 7 HOMEPREMIUM'},
'Family': '7'},
{'Name': 'Windows 7 Professional', {'Name': 'Windows 7 Professional',
'Image File': 'Win7', 'Image File': 'Win7',
'Image Name': 'Windows 7 PROFESSIONAL', 'Image Name': 'Windows 7 PROFESSIONAL'},
'Family': '7'},
{'Name': 'Windows 7 Ultimate', {'Name': 'Windows 7 Ultimate',
'Image File': 'Win7', 'Image File': 'Win7',
'Image Name': 'Windows 7 ULTIMATE', 'Image Name': 'Windows 7 ULTIMATE'},
'Family': '7'},
{'Name': 'Windows 8.1', {'Name': 'Windows 8.1',
'Image File': 'Win8', 'Image File': 'Win8',
'Image Name': 'Windows 8.1', 'Image Name': 'Windows 8.1',
'Family': '8',
'CRLF': True}, 'CRLF': True},
{'Name': 'Windows 8.1 Pro', {'Name': 'Windows 8.1 Pro',
'Image File': 'Win8', 'Image File': 'Win8',
'Image Name': 'Windows 8.1 Pro', 'Image Name': 'Windows 8.1 Pro'},
'Family': '8'},
{'Name': 'Windows 10 Home', {'Name': 'Windows 10 Home',
'Image File': 'Win10', 'Image File': 'Win10',
'Image Name': 'Windows 10 Home', 'Image Name': 'Windows 10 Home',
'Family': '10',
'CRLF': True}, 'CRLF': True},
{'Name': 'Windows 10 Pro', {'Name': 'Windows 10 Pro',
'Image File': 'Win10', 'Image File': 'Win10',
'Image Name': 'Windows 10 Pro', 'Image Name': 'Windows 10 Pro'},
'Family': '10'},
] ]
def find_windows_image(windows_version): def find_windows_image(windows_version):
"""Search for a Windows source image file, returns dict. """Search for a Windows source image file, returns dict.
Searches on local drives and then the WINDOWS_SERVER share.""" Searches on local disks and then the WINDOWS_SERVER share."""
image = {} image = {}
imagefile = windows_version['Image File'] imagefile = windows_version['Image File']
imagename = windows_version['Image Name'] imagename = windows_version['Image Name']
@ -86,8 +78,15 @@ def find_windows_image(windows_version):
windows_version['Name'])) windows_version['Name']))
raise GeneralAbort raise GeneralAbort
def format_gpt(disk, windows_family): def format_disk(disk, use_gpt):
"""Format disk for use as a Windows OS drive using the GPT layout.""" """Format disk for use as a Windows OS disk."""
if use_gpt:
format_gpt(disk)
else:
format_mbr(disk)
def format_gpt(disk):
"""Format disk for use as a Windows OS disk using the GPT layout."""
with open(DISKPART_SCRIPT, 'w') as script: with open(DISKPART_SCRIPT, 'w') as script:
# Partition table # Partition table
script.write('select disk {}\n'.format(disk['Number'])) script.write('select disk {}\n'.format(disk['Number']))
@ -95,7 +94,7 @@ def format_gpt(disk, windows_family):
script.write('convert gpt\n') script.write('convert gpt\n')
# System partition # System partition
# NOTE: ESP needs to be >= 260 for Advanced Format 4K drives # NOTE: ESP needs to be >= 260 for Advanced Format 4K disks
script.write('create partition efi size=500\n') script.write('create partition efi size=500\n')
script.write('format quick fs=fat32 label="System"\n') script.write('format quick fs=fat32 label="System"\n')
script.write('assign letter="S"\n') script.write('assign letter="S"\n')
@ -108,21 +107,20 @@ def format_gpt(disk, windows_family):
script.write('format quick fs=ntfs label="Windows"\n') script.write('format quick fs=ntfs label="Windows"\n')
script.write('assign letter="W"\n') script.write('assign letter="W"\n')
# Recovery Tools partition (Windows 8+) # Recovery Tools partition
if re.search(r'^(8|10)', windows_family): script.write('shrink minimum=500\n')
script.write('shrink minimum=500\n') script.write('create partition primary\n')
script.write('create partition primary\n') script.write('format quick fs=ntfs label="Recovery Tools"\n')
script.write('format quick fs=ntfs label="Recovery Tools"\n') script.write('assign letter="T"\n')
script.write('assign letter="T"\n') script.write('set id="de94bba4-06d1-4d40-a16a-bfd50179d6ac"\n')
script.write('set id="de94bba4-06d1-4d40-a16a-bfd50179d6ac"\n') script.write('gpt attributes=0x8000000000000001\n')
script.write('gpt attributes=0x8000000000000001\n')
# Run script # Run script
run_program(['diskpart', '/s', DISKPART_SCRIPT]) run_program(['diskpart', '/s', DISKPART_SCRIPT])
time.sleep(2) time.sleep(2)
def format_mbr(disk, windows_family): def format_mbr(disk):
"""Format disk for use as a Windows OS drive using the MBR layout.""" """Format disk for use as a Windows OS disk using the MBR layout."""
with open(DISKPART_SCRIPT, 'w') as script: with open(DISKPART_SCRIPT, 'w') as script:
# Partition table # Partition table
script.write('select disk {}\n'.format(disk['Number'])) script.write('select disk {}\n'.format(disk['Number']))
@ -139,13 +137,12 @@ def format_mbr(disk, windows_family):
script.write('format fs=ntfs quick label="Windows"\n') script.write('format fs=ntfs quick label="Windows"\n')
script.write('assign letter="W"\n') script.write('assign letter="W"\n')
# Recovery Tools partition (Windows 8+) # Recovery Tools partition
if re.search(r'^(8|10)', windows_family): script.write('shrink minimum=500\n')
script.write('shrink minimum=500\n') script.write('create partition primary\n')
script.write('create partition primary\n') script.write('format quick fs=ntfs label="Recovery"\n')
script.write('format quick fs=ntfs label="Recovery"\n') script.write('assign letter="T"\n')
script.write('assign letter="T"\n') script.write('set id=27\n')
script.write('set id=27\n')
# Run script # Run script
run_program(['diskpart', '/s', DISKPART_SCRIPT]) run_program(['diskpart', '/s', DISKPART_SCRIPT])

View file

@ -5,7 +5,7 @@ from functions.disk import *
from functions.windows_setup import * from functions.windows_setup import *
# STATIC VARIABLES # STATIC VARIABLES
FAST_COPY_ARGS = [ FAST_COPY_PE_ARGS = [
'/cmd=noexist_only', '/cmd=noexist_only',
'/utf8', '/utf8',
'/skip_empty_dir', '/skip_empty_dir',
@ -20,7 +20,7 @@ PE_TOOLS = {
}, },
'FastCopy': { 'FastCopy': {
'Path': r'FastCopy\FastCopy.exe', 'Path': r'FastCopy\FastCopy.exe',
'Args': FAST_COPY_ARGS, 'Args': FAST_COPY_PE_ARGS,
}, },
'HWiNFO': { 'HWiNFO': {
'Path': r'HWiNFO\HWiNFO.exe', 'Path': r'HWiNFO\HWiNFO.exe',
@ -51,7 +51,7 @@ PE_TOOLS = {
} }
def menu_backup(): def menu_backup():
"""Take backup images of partition(s) in the WIM format and save them to a backup share""" """Take backup images of partition(s) in the WIM format."""
errors = False errors = False
other_results = { other_results = {
'Error': { 'Error': {
@ -65,7 +65,7 @@ def menu_backup():
set_title('{}: Backup Menu'.format(KIT_NAME_FULL)) set_title('{}: Backup Menu'.format(KIT_NAME_FULL))
# Set ticket Number # Set ticket Number
os.system('cls') clear_screen()
ticket_number = get_ticket_number() ticket_number = get_ticket_number()
# Mount backup shares # Mount backup shares
@ -75,8 +75,14 @@ def menu_backup():
destination = select_backup_destination() destination = select_backup_destination()
# Scan disks # Scan disks
try_and_print(message='Assigning letters...', function=assign_volume_letters, other_results=other_results) try_and_print(
result = try_and_print(message='Getting drive info...', function=scan_disks, other_results=other_results) message = 'Assigning letters...',
function = assign_volume_letters,
other_results = other_results)
result = try_and_print(
message = 'Getting disk info...',
function = scan_disks,
other_results = other_results)
if result['CS']: if result['CS']:
disks = result['Out'] disks = result['Out']
else: else:
@ -84,7 +90,7 @@ def menu_backup():
raise GenericAbort raise GenericAbort
# Select disk to backup # Select disk to backup
disk = select_disk('For which drive are we creating backups?', disks) disk = select_disk('For which disk are we creating backups?', disks)
if not disk: if not disk:
raise GenericAbort raise GenericAbort
@ -92,9 +98,8 @@ def menu_backup():
prep_disk_for_backup(destination, disk, ticket_number) prep_disk_for_backup(destination, disk, ticket_number)
# Display details for backup task # Display details for backup task
os.system('cls') clear_screen()
print_info('Create Backup - Details:\n') print_info('Create Backup - Details:\n')
# def show_info(message='~Some message~', info='~Some info~', indent=8, width=32):
show_info(message='Ticket:', info=ticket_number) show_info(message='Ticket:', info=ticket_number)
show_info( show_info(
message = 'Source:', message = 'Source:',
@ -115,9 +120,12 @@ def menu_backup():
# Backup partition(s) # Backup partition(s)
print_info('\n\nStarting task.\n') print_info('\n\nStarting task.\n')
for par in disk['Partitions']: for par in disk['Partitions']:
message = 'Partition {} Backup...'.format(par['Number']) result = try_and_print(
result = try_and_print(message=message, function=backup_partition, message = 'Partition {} Backup...'.format(par['Number']),
other_results=other_results, disk=disk, partition=par) function = backup_partition,
other_results = other_results,
disk = disk,
partition = par)
if not result['CS']: if not result['CS']:
errors = True errors = True
par['Error'] = result['Error'] par['Error'] = result['Error']
@ -128,9 +136,11 @@ def menu_backup():
for par in disk['Partitions']: for par in disk['Partitions']:
if par['Number'] in disk['Bad Partitions']: if par['Number'] in disk['Bad Partitions']:
continue # Skip verification continue # Skip verification
message = 'Partition {} Image...'.format(par['Number']) result = try_and_print(
result = try_and_print(message=message, function=verify_wim_backup, message = 'Partition {} Image...'.format(par['Number']),
other_results=other_results, partition=par) function = verify_wim_backup,
other_results = other_results,
partition = par)
if not result['CS']: if not result['CS']:
errors = True errors = True
par['Error'] = result['Error'] par['Error'] = result['Error']
@ -199,12 +209,12 @@ def menu_root():
exit_script() exit_script()
def menu_setup(): def menu_setup():
"""Format a drive, partition for MBR or GPT, apply a Windows image, and rebuild the boot files""" """Format a disk (MBR/GPT), apply a Windows image, and setup boot files."""
errors = False errors = False
set_title('{}: Setup Menu'.format(KIT_NAME_FULL)) set_title('{}: Setup Menu'.format(KIT_NAME_FULL))
# Set ticket ID # Set ticket ID
os.system('cls') clear_screen()
ticket_number = get_ticket_number() ticket_number = get_ticket_number()
# Select the version of Windows to apply # Select the version of Windows to apply
@ -214,16 +224,22 @@ def menu_setup():
windows_image = find_windows_image(windows_version) windows_image = find_windows_image(windows_version)
# Scan disks # Scan disks
try_and_print(message='Assigning letters...', function=assign_volume_letters, other_results=other_results) try_and_print(
result = try_and_print(message='Getting drive info...', function=scan_disks, other_results=other_results) message = 'Assigning letters...',
function = assign_volume_letters,
other_results = other_results)
result = try_and_print(
message = 'Getting disk info...',
function = scan_disks,
other_results = other_results)
if result['CS']: if result['CS']:
disks = result['Out'] disks = result['Out']
else: else:
print_error('ERROR: No disks found.') print_error('ERROR: No disks found.')
raise GenericAbort raise GenericAbort
# Select drive to use as the OS drive # Select disk to use as the OS disk
dest_disk = select_disk('To which drive are we installing Windows?', disks) dest_disk = select_disk('To which disk are we installing Windows?', disks)
if not disk: if not disk:
raise GenericAbort raise GenericAbort
@ -231,92 +247,77 @@ def menu_setup():
prep_disk_for_formatting(dest_disk) prep_disk_for_formatting(dest_disk)
# Display details for setup task # Display details for setup task
os.system('cls') clear_screen()
print('Setup Windows - Details:\n') print_info('Setup Windows - Details:\n')
print(' Ticket: \t{ticket_number}'.format(ticket_number=ticket_number)) show_info(message='Ticket:', info=ticket_number)
print(' Installing: \t{winver}'.format(winver=windows_version['Name'])) show_info(message='Installing:', info=windows_version['Name'])
print(' Boot Method:\t{_type}'.format( show_info(
_type='UEFI (GPT)' if dest_disk['Use GPT'] else 'Legacy (MBR)')) message = 'Boot Method:',
print(' Using Image:\t{}'.format(windows_image['Path'])) info = 'UEFI (GPT)' if dest_disk['Use GPT'] else 'Legacy (MBR)')
print_warning(' ERASING: \t[{Table}] ({Type}) {Name} {Size}\n'.format(**dest_disk)) show_info(message='Using Image:', info=windows_version['Path'])
print_warning(' ERASING: \t[{Table}] ({Type}) {Name} {Size}\n'.format(
**dest_disk))
for par in dest_disk['Partitions']: for par in dest_disk['Partitions']:
print_warning(par['Display String']) print_warning(par['Display String'])
print_warning(dest_disk['Format Warnings']) print_warning(dest_disk['Format Warnings'])
if (not ask('Is this correct?')): if (not ask('Is this correct?')):
abort_to_main_menu('Aborting Windows setup') raise GeneralAbort
# Safety check # Safety check
print('\nSAFETY CHECK') print_standard('\nSAFETY CHECK')
print_warning('All data will be DELETED from the drive and partition(s) listed above.') print_warning('All data will be DELETED from the '
print_warning('This is irreversible and will lead to {CLEAR}{RED}DATA LOSS.'.format(**COLORS)) 'disk & partition(s) listed above.')
print_warning('This is irreversible and will lead '
'to {CLEAR}{RED}DATA LOSS.'.format(**COLORS))
if (not ask('Asking again to confirm, is this correct?')): if (not ask('Asking again to confirm, is this correct?')):
abort_to_main_menu('Aborting Windows setup') raise GeneralAbort
# Release currently used volume letters (ensures that the drives will get S, T, & W as needed below) # Remove volume letters so S, T, & W can be used below
remove_volume_letters(keep=windows_image['Source']) remove_volume_letters(keep=windows_image['Source'])
new_letter = reassign_volume_letter(letter=windows_image['Source']) new_letter = reassign_volume_letter(letter=windows_image['Source'])
if new_letter: if new_letter:
windows_image['Source'] = new_letter windows_image['Source'] = new_letter
# Format and partition drive # Format and partition disk
print('\n Formatting Drive... \t\t', end='', flush=True) result = try_and_print(
try: message = 'Formatting disk...',
if (dest_disk['Use GPT']): function = format_disk,
format_gpt(dest_disk, windows_version['Family']) other_results = other_results,
else: disk = dest_disk,
format_mbr(dest_disk, windows_version['Family']) use_gpt = dest_disk['Use GPT'])
print_success('Complete.') if not result['CS']:
except: # We need to crash as the disk is in an unknown state
# We need to crash as the drive is in an unknown state print_error('ERROR: Failed to format disk.')
print_error('Failed.') raise GenericAbort
raise
# Apply Image # Apply Image
print(' Applying Image... \t\t', end='', flush=True) result = try_and_print(
try: message = 'Applying image...',
setup_windows(windows_image, windows_version) function = setup_windows,
print_success('Complete.') other_results = other_results,
except subprocess.CalledProcessError: windows_image = windows_image,
print_error('Failed.') windows_version = windows_version)
errors = True if not result['CS']:
except: # We need to crash as the disk is in an unknown state
# We need to crash as the OS is in an unknown state print_error('ERROR: Failed to apply image.')
print_error('Failed.') raise GenericAbort
raise
# Create Boot files # Create Boot files
print(' Update Boot Partition...\t\t', end='', flush=True) try_and_print(
try: message = 'Updating boot files...',
update_boot_partition() function = update_boot_partition,
print_success('Complete.') other_results = other_results)
except subprocess.CalledProcessError:
# Don't need to crash as this is (potentially) recoverable
print_error('Failed.')
errors = True
except:
print_error('Failed.')
raise
# Setup WinRE # Setup WinRE
print(' Update Recovery Tools...\t\t', end='', flush=True) try_and_print(
try: message = 'Updating recovery tools...',
setup_windows_re(windows_version) function = setup_windows_re,
print_success('Complete.') other_results = other_results,
except SetupError: windows_version = windows_version)
print('Skipped.')
except:
# Don't need to crash as this is (potentially) recoverable
print_error('Failed.')
errors = True
# Print summary # Print summary
if errors: print_standard('\nDone.')
print_warning('\nErrors were encountered during setup.')
time.sleep(30)
else:
print_success('\nNo errors were encountered during setup.')
time.sleep(5)
pause('\nPress Enter to return to main menu... ') pause('\nPress Enter to return to main menu... ')
def menu_tools(): def menu_tools():
@ -359,9 +360,9 @@ def select_minidump_path():
# Remove RAMDisk letter # Remove RAMDisk letter
if 'X' in tmp: if 'X' in tmp:
tmp.remove('X') tmp.remove('X')
for drive in tmp: for disk in tmp:
if os.path.exists('{drive}:\\Windows\\MiniDump'.format(drive=drive)): if os.path.exists('{}:\\Windows\\MiniDump'.format(disk)):
dumps.append({'Name': '{drive}:\\Windows\\MiniDump'.format(drive=drive)}) dumps.append({'Name': '{}:\\Windows\\MiniDump'.format(disk)})
# Check results before showing menu # Check results before showing menu
if len(dumps) == 0: if len(dumps) == 0: