From cd3f3ed8d38fdd90ada4a57875373a9a8fb2fb87 Mon Sep 17 00:00:00 2001
From: Alan Mason <1923621+2Shirt@users.noreply.github.com>
Date: Fri, 1 Dec 2017 12:02:05 -0800
Subject: [PATCH 1/9] Always more cleanup

---
 Scripts/functions/disk.py          | 6 +++---
 Scripts/functions/windows_setup.py | 7 ++-----
 2 files changed, 5 insertions(+), 8 deletions(-)

diff --git a/Scripts/functions/disk.py b/Scripts/functions/disk.py
index eb2bda4f..5d6cc2a7 100644
--- a/Scripts/functions/disk.py
+++ b/Scripts/functions/disk.py
@@ -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 = []
diff --git a/Scripts/functions/windows_setup.py b/Scripts/functions/windows_setup.py
index cbbe1525..ae8d401b 100644
--- a/Scripts/functions/windows_setup.py
+++ b/Scripts/functions/windows_setup.py
@@ -150,11 +150,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 = [

From 5aab8a98e371ab5dca128f8567f76de2ae400dcb Mon Sep 17 00:00:00 2001
From: Alan Mason <1923621+2Shirt@users.noreply.github.com>
Date: Fri, 1 Dec 2017 12:24:11 -0800
Subject: [PATCH 2/9] Bugfix: Was trying to use 'Env' before it was set.

---
 Scripts/functions/windows_setup.py | 1 -
 Scripts/winpe_root_menu.py         | 3 +++
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/Scripts/functions/windows_setup.py b/Scripts/functions/windows_setup.py
index ae8d401b..c96ed5a2 100644
--- a/Scripts/functions/windows_setup.py
+++ b/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',
diff --git a/Scripts/winpe_root_menu.py b/Scripts/winpe_root_menu.py
index 4bd39aef..407b3e84 100644
--- a/Scripts/winpe_root_menu.py
+++ b/Scripts/winpe_root_menu.py
@@ -11,6 +11,9 @@ 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()

From ab6fd19c021120988335f6ef859a517bd33eb747 Mon Sep 17 00:00:00 2001
From: Alan Mason <1923621+2Shirt@users.noreply.github.com>
Date: Fri, 1 Dec 2017 12:26:11 -0800
Subject: [PATCH 3/9] Bugfix: menu_select()

---
 Scripts/functions/common.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Scripts/functions/common.py b/Scripts/functions/common.py
index 59826f0f..9a2211a4 100644
--- a/Scripts/functions/common.py
+++ b/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)

From 06affc3bd109dedb875a5cc809030745336b6fcb Mon Sep 17 00:00:00 2001
From: Alan Mason <1923621+2Shirt@users.noreply.github.com>
Date: Fri, 1 Dec 2017 12:26:37 -0800
Subject: [PATCH 4/9] Bugfix: menu_tools()

---
 Scripts/functions/winpe_menus.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Scripts/functions/winpe_menus.py b/Scripts/functions/winpe_menus.py
index c6ba4e8c..f34376ae 100644
--- a/Scripts/functions/winpe_menus.py
+++ b/Scripts/functions/winpe_menus.py
@@ -321,7 +321,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)
 

From 97e93e812be4c7990d1860ff3577fc12e28795a1 Mon Sep 17 00:00:00 2001
From: Alan Mason <1923621+2Shirt@users.noreply.github.com>
Date: Fri, 1 Dec 2017 12:27:26 -0800
Subject: [PATCH 5/9] Bugfix: partition_uids import

---
 Scripts/functions/disk.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Scripts/functions/disk.py b/Scripts/functions/disk.py
index 5d6cc2a7..76601ffb 100644
--- a/Scripts/functions/disk.py
+++ b/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)

From 5d32f3c94a1f3f79d85d850ffbde66a567755463 Mon Sep 17 00:00:00 2001
From: Alan Mason <1923621+2Shirt@users.noreply.github.com>
Date: Fri, 1 Dec 2017 12:27:53 -0800
Subject: [PATCH 6/9] Bugfix: Abort logic

---
 Scripts/functions/disk.py          | 2 +-
 Scripts/functions/windows_setup.py | 4 ++--
 Scripts/functions/winpe_menus.py   | 9 +++++----
 Scripts/winpe_root_menu.py         | 3 ---
 4 files changed, 8 insertions(+), 10 deletions(-)

diff --git a/Scripts/functions/disk.py b/Scripts/functions/disk.py
index 76601ffb..8cff593e 100644
--- a/Scripts/functions/disk.py
+++ b/Scripts/functions/disk.py
@@ -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/Scripts/functions/windows_setup.py b/Scripts/functions/windows_setup.py
index c96ed5a2..90ff549f 100644
--- a/Scripts/functions/windows_setup.py
+++ b/Scripts/functions/windows_setup.py
@@ -75,7 +75,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."""
@@ -166,7 +166,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/Scripts/functions/winpe_menus.py b/Scripts/functions/winpe_menus.py
index f34376ae..64a63e97 100644
--- a/Scripts/functions/winpe_menus.py
+++ b/Scripts/functions/winpe_menus.py
@@ -197,8 +197,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'):
@@ -262,7 +263,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 +272,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'])
diff --git a/Scripts/winpe_root_menu.py b/Scripts/winpe_root_menu.py
index 407b3e84..1ddea8a5 100644
--- a/Scripts/winpe_root_menu.py
+++ b/Scripts/winpe_root_menu.py
@@ -17,9 +17,6 @@ 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:

From a6f9518648887a5c14c2dac477da0647fe9da06d Mon Sep 17 00:00:00 2001
From: Alan Mason <1923621+2Shirt@users.noreply.github.com>
Date: Fri, 1 Dec 2017 12:41:39 -0800
Subject: [PATCH 7/9] Bugfix: find_windows_image()

---
 Scripts/functions/windows_setup.py | 25 ++++++++++++++-----------
 1 file changed, 14 insertions(+), 11 deletions(-)

diff --git a/Scripts/functions/windows_setup.py b/Scripts/functions/windows_setup.py
index 90ff549f..1d960367 100644
--- a/Scripts/functions/windows_setup.py
+++ b/Scripts/functions/windows_setup.py
@@ -56,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:

From 56c396045144afa8b5465fae6752f8418408d0ba Mon Sep 17 00:00:00 2001
From: Alan Mason <1923621+2Shirt@users.noreply.github.com>
Date: Fri, 1 Dec 2017 12:42:23 -0800
Subject: [PATCH 8/9] Bugfix: menu_setup()

---
 Scripts/functions/winpe_menus.py | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/Scripts/functions/winpe_menus.py b/Scripts/functions/winpe_menus.py
index 64a63e97..637f5d0f 100644
--- a/Scripts/functions/winpe_menus.py
+++ b/Scripts/functions/winpe_menus.py
@@ -212,6 +212,15 @@ 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

From f608e9b88c6b368aeda0d124801b0a81d31cd49b Mon Sep 17 00:00:00 2001
From: Alan Mason <1923621+2Shirt@users.noreply.github.com>
Date: Fri, 1 Dec 2017 12:42:46 -0800
Subject: [PATCH 9/9] Show title before getting ticket_number

---
 Scripts/functions/winpe_menus.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/Scripts/functions/winpe_menus.py b/Scripts/functions/winpe_menus.py
index 637f5d0f..b89ce79d 100644
--- a/Scripts/functions/winpe_menus.py
+++ b/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
@@ -225,6 +226,7 @@ def menu_setup():
 
     # Set ticket ID
     clear_screen()
+    print_standard('{}\n'.format(global_vars['Title']))
     ticket_number = get_ticket_number()
 
     # Select the version of Windows to apply