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