diff --git a/.bin/Scripts/functions/common.py b/.bin/Scripts/functions/common.py index 2c299628..d8ea62e5 100644 --- a/.bin/Scripts/functions/common.py +++ b/.bin/Scripts/functions/common.py @@ -226,7 +226,7 @@ def menu_select(title='~ Untitled Menu ~', # Set title if 'Title' in global_vars: - title = '{}\n\n{}'.format(global_vars['Title']) + title = '{}\n\n{}'.format(global_vars['Title'], title) # Build menu menu_splash = '{}\n\n'.format(title) diff --git a/.bin/Scripts/functions/disk.py b/.bin/Scripts/functions/disk.py index 78d08c3c..c1e88491 100644 --- a/.bin/Scripts/functions/disk.py +++ b/.bin/Scripts/functions/disk.py @@ -1,7 +1,7 @@ # Wizard Kit PE: Functions - Disk from functions.common import * -import partition_uids +from functions import partition_uids # Regex REGEX_BAD_PARTITION = re.compile(r'(RAW|Unknown)', re.IGNORECASE) @@ -197,7 +197,7 @@ def get_table_type(disk): elif REGEX_DISK_RAW.search(output): part_type = 'RAW' else: - part_type = 'Unknown + part_type = 'Unknown' return part_type @@ -288,7 +288,7 @@ def remove_volume_letters(keep=None): with open(DISKPART_SCRIPT, 'w') as script: for vol in get_volumes(): if vol['Letter'].upper() != keep.upper(): - script.write('select volume {}\n'.format(vol['Number'])) + script.write('select volume {}\n'.format(vol['Number'])) script.write('remove noerr\n') # Run script @@ -319,7 +319,7 @@ def scan_disks(): # Done return disks -def select_disk(title='Which disk?', disks): +def select_disk(title='Which disk?', disks=[]): """Select a disk from the attached disks""" # Build menu disk_options = [] @@ -360,7 +360,7 @@ def select_disk(title='Which disk?', disks): if (selection.isnumeric()): return disk_options[int(selection)-1]['Disk'] elif (selection == 'M'): - raise GeneralAbort + raise GenericAbort if __name__ == '__main__': print("This file is not meant to be called directly.") diff --git a/.bin/Scripts/functions/windows_setup.py b/.bin/Scripts/functions/windows_setup.py index 7097a371..0096b116 100644 --- a/.bin/Scripts/functions/windows_setup.py +++ b/.bin/Scripts/functions/windows_setup.py @@ -3,7 +3,6 @@ from functions.data import * # STATIC VARIABLES -DISKPART_SCRIPT = r'{}\diskpart.script'.format(global_vars['Env']['TMP']) WINDOWS_VERSIONS = [ {'Name': 'Windows 7 Home Basic', 'Image File': 'Win7', @@ -57,17 +56,20 @@ def find_windows_image(windows_version): # Check for network source if not image: mount_windows_share() - if not WINDOWS_SERVER['Mounted']: - return None - for ext in ['esd', 'wim', 'swm']: - path = r'\\{}\{}\images\{}.ext'.format( - WINDOWS_SERVER['IP'], WINDOWS_SERVER['Share'], imagefile, ext) - if os.path.isfile(path) and wim_contains_image(path, imagename): - image['Path'] = path - image['Source'] = None - if ext == 'swm': - image['Glob'] = '--ref="{}*.swm"'.format(image['Path'][:-4]) - break + if WINDOWS_SERVER['Mounted']: + for ext in ['esd', 'wim', 'swm']: + path = r'\\{}\{}\images\{}.ext'.format( + WINDOWS_SERVER['IP'], + WINDOWS_SERVER['Share'], + imagefile, + ext) + if os.path.isfile(path) and wim_contains_image(path, imagename): + image['Path'] = path + image['Source'] = None + if ext == 'swm': + image['Glob'] = '--ref="{}*.swm"'.format( + image['Path'][:-4]) + break # Display image to be used (if any) and return if image: @@ -76,7 +78,7 @@ def find_windows_image(windows_version): else: print_error('Failed to find Windows source image for {}'.format( windows_version['Name'])) - raise GeneralAbort + raise GenericAbort def format_disk(disk, use_gpt): """Format disk for use as a Windows OS disk.""" @@ -150,11 +152,8 @@ def format_mbr(disk): def mount_windows_share(): """Mount the Windows images share unless labeled as already mounted.""" - if WINDOWS_SERVER['Mounted']: - # Blindly skip if we mounted earlier - continue - - mount_network_share(WINDOWS_SERVER) + if not WINDOWS_SERVER['Mounted']: + mount_network_share(WINDOWS_SERVER) def select_windows_version(): actions = [ @@ -170,7 +169,7 @@ def select_windows_version(): if selection.isnumeric(): return WINDOWS_VERSIONS[int(selection)-1] elif selection == 'M': - raise GeneralAbort + raise GenericAbort def setup_windows(windows_image, windows_version): cmd = [ diff --git a/.bin/Scripts/functions/winpe_menus.py b/.bin/Scripts/functions/winpe_menus.py index d7673774..dcd0ba18 100644 --- a/.bin/Scripts/functions/winpe_menus.py +++ b/.bin/Scripts/functions/winpe_menus.py @@ -66,6 +66,7 @@ def menu_backup(): # Set ticket Number clear_screen() + print_standard('{}\n'.format(global_vars['Title'])) ticket_number = get_ticket_number() # Mount backup shares @@ -197,8 +198,9 @@ def menu_root(): if (selection.isnumeric()): try: menus[int(selection)-1]['Menu']() - except AbortError: - pass + 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'): @@ -211,10 +213,20 @@ def menu_root(): 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 @@ -262,7 +274,7 @@ def menu_setup(): print_warning(dest_disk['Format Warnings']) if (not ask('Is this correct?')): - raise GeneralAbort + raise GenericAbort # Safety check print_standard('\nSAFETY CHECK') @@ -271,7 +283,7 @@ def menu_setup(): 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 GeneralAbort + raise GenericAbort # Remove volume letters so S, T, & W can be used below remove_volume_letters(keep=windows_image['Source']) @@ -321,7 +333,7 @@ def menu_setup(): pause('\nPress Enter to return to main menu... ') def menu_tools(): - tools = [k for k in sorted(PE_TOOLS.keys())] + tools = [{'Name': k} for k in sorted(PE_TOOLS.keys())] actions = [{'Name': 'Main Menu', 'Letter': 'M'},] set_title(KIT_NAME_FULL) diff --git a/.bin/Scripts/winpe_root_menu.py b/.bin/Scripts/winpe_root_menu.py index 94166876..279fa33d 100644 --- a/.bin/Scripts/winpe_root_menu.py +++ b/.bin/Scripts/winpe_root_menu.py @@ -11,12 +11,12 @@ init_global_vars() set_title('{}: Root Menu'.format(KIT_NAME_FULL)) global_vars['LogFile'] = r'{LogDir}\WinPE.log'.format(**global_vars) +# STATIC VARIABLES +DISKPART_SCRIPT = r'{}\diskpart.script'.format(global_vars['Env']['TMP']) + if __name__ == '__main__': try: menu_root() - except GenericAbort: - # pause('Press Enter to return to main menu... ') - pass except SystemExit: pass except: