# Wizard Kit: WinPE Menus from functions.backup import * from functions.disk import * from functions.windows_setup import * # STATIC VARIABLES FAST_COPY_PE_ARGS = [ '/cmd=noexist_only', '/utf8', '/skip_empty_dir', '/linkdest', '/no_ui', '/auto_close', '/exclude={}'.format(';'.join(FAST_COPY_EXCLUDES)), ] PE_TOOLS = { 'BlueScreenView': { 'Path': r'BlueScreenView\BlueScreenView.exe', }, 'FastCopy': { 'Path': r'FastCopy\FastCopy.exe', 'Args': FAST_COPY_PE_ARGS, }, 'HWiNFO': { 'Path': r'HWiNFO\HWiNFO.exe', }, 'NT Password Editor': { 'Path': r'NT Password Editor\ntpwedit.exe', }, 'Notepad++': { 'Path': r'NotepadPlusPlus\NotepadPlusPlus.exe', }, 'PhotoRec': { 'Path': r'TestDisk\photorec_win.exe', 'Args': ['-new_console:n'], }, 'Prime95': { 'Path': r'Prime95\prime95.exe', }, 'ProduKey': { 'Path': r'ProduKey\ProduKey.exe', }, 'Q-Dir': { 'Path': r'Q-Dir\Q-Dir.exe', }, 'TestDisk': { 'Path': r'TestDisk\testdisk_win.exe', 'Args': ['-new_console:n'], }, } def check_pe_tools(): for k in PE_TOOLS.keys(): PE_TOOLS[k]['Path'] = r'{}\{}'.format( global_vars['BinDir'], PE_TOOLS[k]['Path']) global_vars['Tools']['wimlib-imagex'] = re.sub( r'\\x(32|64)', r'', global_vars['Tools']['wimlib-imagex'], re.IGNORECASE) def menu_backup(): """Take backup images of partition(s) in the WIM format.""" errors = False other_results = { 'Error': { 'CalledProcessError': 'Unknown Error', 'PathNotFoundError': 'Missing', }, 'Warning': { 'GenericAbort': 'Skipped', 'GenericRepair': 'Repaired', }} set_title('{}: Backup Menu'.format(KIT_NAME_FULL)) # Set ticket Number clear_screen() print_standard('{}\n'.format(global_vars['Title'])) ticket_number = get_ticket_number() # Mount backup shares mount_backup_shares() # Select destination destination = select_backup_destination(auto_select=False) # Scan disks try_and_print( 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']: disks = result['Out'] else: print_error('ERROR: No disks found.') raise GenericAbort # Select disk to backup disk = select_disk('For which disk are we creating backups?', disks) if not disk: raise GenericAbort # "Prep" disk prep_disk_for_backup(destination, disk, ticket_number) # Display details for backup task clear_screen() print_info('Create Backup - Details:\n') show_data(message='Ticket:', data=ticket_number) show_data( message = 'Source:', data = '[{Table}] ({Type}) {Name} {Size}'.format(**disk), ) show_data( message = 'Destination:', data = destination.get('Display Name', destination['Name']), ) for par in disk['Partitions']: message = 'Partition {}:'.format(par['Number']) data = par['Display String'] if par['Number'] in disk['Bad Partitions']: show_data(message=message, data=data, width=30, warning=True) elif par['Image Exists']: show_data(message=message, data=data, width=30, info=True) else: show_data(message=message, data=data, width=30) print_standard(disk['Backup Warnings']) # Ask to proceed if (not ask('Proceed with backup?')): raise GenericAbort # Backup partition(s) print_info('\n\nStarting task.\n') for par in disk['Partitions']: result = try_and_print( message = 'Partition {} Backup...'.format(par['Number']), function = backup_partition, other_results = other_results, disk = disk, par = par) if not result['CS'] and not isinstance(result['Error'], GenericAbort): errors = True par['Error'] = result['Error'] # Verify backup(s) if disk['Valid Partitions']: print_info('\n\n Verifying backup images(s)\n') for par in disk['Partitions']: if par['Number'] in disk['Bad Partitions']: continue # Skip verification result = try_and_print( message = 'Partition {} Image...'.format(par['Number']), function = verify_wim_backup, other_results = other_results, partition = par) if not result['CS']: errors = True par['Error'] = result['Error'] # Print summary if errors: print_warning('\nErrors were encountered and are detailed below.') for par in [p for p in disk['Partitions'] if 'Error' in p]: print_standard(' Partition {} Error:'.format(par['Number'])) if hasattr(par['Error'], 'stderr'): try: par['Error'] = par['Error'].stderr.decode() except: # Deal with badly formatted error message pass try: par['Error'] = par['Error'].splitlines() par['Error'] = [line.strip() for line in par['Error']] par['Error'] = [line for line in par['Error'] if line] for line in par['Error']: print_error('\t{}'.format(line)) except: print_error('\t{}'.format(par['Error'])) else: print_success('\nNo errors were encountered during imaging.') if 'LogFile' in global_vars: cmd = [ global_vars['Tools']['NotepadPlusPlus'], global_vars['LogFile']] try: popen_program(cmd) except Exception: print_error('ERROR: Failed to open log.') sleep(30) pause('\nPress Enter to return to main menu... ') def menu_root(): check_pe_tools() menus = [ {'Name': 'Create Backups', 'Menu': menu_backup}, {'Name': 'Setup Windows', 'Menu': menu_setup}, {'Name': 'Misc Tools', 'Menu': menu_tools}, ] actions = [ {'Name': 'Command Prompt', 'Letter': 'C'}, {'Name': 'Reboot', 'Letter': 'R'}, {'Name': 'Shutdown', 'Letter': 'S'}, ] # Main loop while True: set_title(KIT_NAME_FULL) selection = menu_select( title = 'Main Menu', main_entries = menus, action_entries = actions, secret_exit = True) if (selection.isnumeric()): try: menus[int(selection)-1]['Menu']() except GenericAbort: print_warning('\nAborted\n') pause('Press Enter to return to main menu... ') elif (selection == 'C'): run_program(['cmd', '-new_console:n'], check=False) elif (selection == 'R'): run_program(['wpeutil', 'reboot']) elif (selection == 'S'): run_program(['wpeutil', 'shutdown']) else: sys.exit() def menu_setup(): """Format a disk (MBR/GPT), apply a Windows image, and setup boot files.""" errors = False other_results = { 'Error': { 'CalledProcessError': 'Unknown Error', 'PathNotFoundError': 'Missing', }, 'Warning': { 'GenericAbort': 'Skipped', 'GenericRepair': 'Repaired', }} set_title('{}: Setup Menu'.format(KIT_NAME_FULL)) # Set ticket ID clear_screen() print_standard('{}\n'.format(global_vars['Title'])) ticket_number = get_ticket_number() # Select the version of Windows to apply windows_version = select_windows_version() # Find Windows image windows_image = find_windows_image(windows_version) # Scan disks try_and_print( 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']: disks = result['Out'] else: print_error('ERROR: No disks found.') raise GenericAbort # Select disk to use as the OS disk dest_disk = select_disk('To which disk are we installing Windows?', disks) if not dest_disk: raise GenericAbort # "Prep" disk prep_disk_for_formatting(dest_disk) # Display details for setup task clear_screen() print_info('Setup Windows - Details:\n') show_data(message='Ticket:', data=ticket_number) show_data(message='Installing:', data=windows_version['Name']) show_data( message = 'Boot Method:', data = 'UEFI (GPT)' if dest_disk['Use GPT'] else 'Legacy (MBR)') show_data(message='Using Image:', data=windows_image['Path']) show_data( message = 'ERASING:', data = '[{Table}] ({Type}) {Name} {Size}\n'.format(**dest_disk), warning = True) for par in dest_disk['Partitions']: show_data( message = 'Partition {}:'.format(par['Number']), data = par['Display String'], warning = True) print_warning(dest_disk['Format Warnings']) if (not ask('Is this correct?')): raise GenericAbort # Safety check print_standard('\nSAFETY CHECK') print_warning('All data will be DELETED from the ' 'disk and 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?')): raise GenericAbort # Remove volume letters so S, T, & W can be used below remove_volume_letters(keep=windows_image['Source']) new_letter = reassign_volume_letter(letter=windows_image['Source']) if new_letter: windows_image['Source'] = new_letter # Format and partition disk result = try_and_print( message = 'Formatting disk...', function = format_disk, other_results = other_results, disk = dest_disk, use_gpt = dest_disk['Use GPT']) if not result['CS']: # We need to crash as the disk is in an unknown state print_error('ERROR: Failed to format disk.') raise GenericAbort # Apply Image result = try_and_print( message = 'Applying image...', function = setup_windows, other_results = other_results, windows_image = windows_image, windows_version = windows_version) if not result['CS']: # We need to crash as the disk is in an unknown state print_error('ERROR: Failed to apply image.') raise GenericAbort # Create Boot files try_and_print( message = 'Updating boot files...', function = update_boot_partition, other_results = other_results) # Setup WinRE try_and_print( message = 'Updating recovery tools...', function = setup_windows_re, other_results = other_results, windows_version = windows_version) # Print summary print_standard('\nDone.') if 'LogFile' in global_vars: cmd = [ global_vars['Tools']['NotepadPlusPlus'], global_vars['LogFile']] try: popen_program(cmd) except Exception: print_error('ERROR: Failed to open log.') sleep(30) pause('\nPress Enter to return to main menu... ') def menu_tools(): tools = [{'Name': k} for k in sorted(PE_TOOLS.keys())] actions = [{'Name': 'Main Menu', 'Letter': 'M'},] set_title(KIT_NAME_FULL) # Menu loop while True: selection = menu_select( title = 'Tools Menu', main_entries = tools, action_entries = actions) if (selection.isnumeric()): name = tools[int(selection)-1]['Name'] cmd = [PE_TOOLS[name]['Path']] + PE_TOOLS[name].get('Args', []) if name == 'Blue Screen View': # Select path to scan minidump_path = select_minidump_path() if minidump_path: cmd.extend(['/MiniDumpFolder', minidump_path]) try: popen_program(cmd) except Exception: print_error('Failed to run {}'.format(name)) sleep(2) pause() elif (selection == 'M'): break def select_minidump_path(): dumps = [] # Assign volume letters first assign_volume_letters() # Search for minidumps tmp = run_program('mountvol') tmp = [d for d in re.findall(r'.*([A-Za-z]):\\', tmp.stdout.decode())] # Remove RAMDisk letter if 'X' in tmp: tmp.remove('X') for disk in tmp: if os.path.exists('{}:\\Windows\\MiniDump'.format(disk)): dumps.append({'Name': '{}:\\Windows\\MiniDump'.format(disk)}) # Check results before showing menu if len(dumps) == 0: print_error(' No BSoD / MiniDump paths found') sleep(2) return None # Menu selection = menu_select( title = 'Which BSoD / MiniDump path are we scanning?', main_entries = dumps) return dumps[int(selection) - 1]['Name'] if __name__ == '__main__': print("This file is not meant to be called directly.")