Merge branch 'python-overhaul' into dev

# Conflicts:
#	.bin/Scripts/functions/common.py
#	.bin/Scripts/functions/disk.py
#	.bin/Scripts/functions/windows_setup.py
#	.bin/Scripts/functions/winpe_menus.py
#	.bin/Scripts/winpe_root_menu.py

.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)

.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.")

.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']:
+        if WINDOWS_SERVER['Mounted']:
-            return None
+            for ext in ['esd', 'wim', 'swm']:
-        for ext in ['esd', 'wim', 'swm']:
+                path = r'\\{}\{}\images\{}.ext'.format(
-            path = r'\\{}\{}\images\{}.ext'.format(
+                    WINDOWS_SERVER['IP'],
-                WINDOWS_SERVER['IP'], WINDOWS_SERVER['Share'], imagefile, ext)
+                    WINDOWS_SERVER['Share'],
-            if os.path.isfile(path) and wim_contains_image(path, imagename):
+                    imagefile,
-                image['Path'] = path
+                    ext)
-                image['Source'] = None
+                if os.path.isfile(path) and wim_contains_image(path, imagename):
-                if ext == 'swm':
+                    image['Path'] = path
-                    image['Glob'] = '--ref="{}*.swm"'.format(image['Path'][:-4])
+                    image['Source'] = None
-                break
+                    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']:
+    if not WINDOWS_SERVER['Mounted']:
-        # Blindly skip if we mounted earlier
+        mount_network_share(WINDOWS_SERVER)
-        continue
-    
-    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 = [

.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:
+            except GenericAbort:
-                pass
+                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)
 

.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: