From 0902a7997f30eac0e5d994e08213bb4b24532e1f Mon Sep 17 00:00:00 2001
From: 2Shirt <1923621+2Shirt@users.noreply.github.com>
Date: Mon, 14 May 2018 12:58:06 -0600
Subject: [PATCH] Moved SafeMode sections into functions/safemode.py

* Allows for better descriptions of the actions being taken
* Allowd for better error handling
* Fixes issue #34
---
 .bin/Scripts/functions/safemode.py | 36 ++++++++++++++++++++++++++++++
 .bin/Scripts/safemode_enter.py     | 32 +++++++++++---------------
 .bin/Scripts/safemode_exit.py      | 30 +++++++++++--------------
 3 files changed, 62 insertions(+), 36 deletions(-)
 create mode 100644 .bin/Scripts/functions/safemode.py

diff --git a/.bin/Scripts/functions/safemode.py b/.bin/Scripts/functions/safemode.py
new file mode 100644
index 00000000..9f44aa04
--- /dev/null
+++ b/.bin/Scripts/functions/safemode.py
@@ -0,0 +1,36 @@
+# Wizard Kit: Functions - SafeMode
+
+from functions.common import *
+
+# STATIC VARIABLES
+REG_MSISERVER = r'HKLM\SYSTEM\CurrentControlSet\Control\SafeBoot\Network\MSIServer'
+
+def disable_safemode_msi():
+    """Disable MSI access under safemode."""
+    cmd = ['reg', 'delete', REG_MSISERVER, '/f']
+    run_program(cmd)
+
+def disable_safemode():
+    """Edit BCD to remove safeboot value."""
+    cmd = ['bcdedit', '/deletevalue', '{default}', 'safeboot']
+    run_program(cmd)
+
+def enable_safemode_msi():
+    """Enable MSI access under safemode."""
+    cmd = ['reg', 'add', REG_MSISERVER, '/f']
+    run_program(cmd)
+    cmd = ['reg', 'add', REG_MSISERVER, '/ve',
+            '/t', 'REG_SZ', '/d', 'Service', '/f']
+    run_program(cmd)
+
+def enable_safemode():
+    """Edit BCD to set safeboot as default."""
+    cmd = ['bcdedit', '/set', '{default}', 'safeboot', 'network']
+    run_program(cmd)
+
+def reboot(delay=3):
+    cmd = ['shutdown', '-r', '-t', '{}'.format(delay)]
+    run_program(cmd, check=False)
+
+if __name__ == '__main__':
+    print("This file is not meant to be called directly.")
diff --git a/.bin/Scripts/safemode_enter.py b/.bin/Scripts/safemode_enter.py
index 825e7352..c3213cb9 100644
--- a/.bin/Scripts/safemode_enter.py
+++ b/.bin/Scripts/safemode_enter.py
@@ -3,13 +3,10 @@
 import os
 import sys
 
-# STATIC VARIABLES
-REG_MSISERVER = r'HKLM\SYSTEM\CurrentControlSet\Control\SafeBoot\Network\MSIServer'
-
 # Init
 os.chdir(os.path.dirname(os.path.realpath(__file__)))
 sys.path.append(os.getcwd())
-from functions.common import *
+from functions.safemode import *
 init_global_vars()
 os.system('title {}: SafeMode Tool'.format(KIT_NAME_FULL))
 
@@ -17,26 +14,23 @@ if __name__ == '__main__':
     try:
         clear_screen()
         print_info('{}: SafeMode Tool\n'.format(KIT_NAME_FULL))
+        other_results = {
+            'Error': {'CalledProcessError':   'Unknown Error'},
+            'Warning': {}}
+        
         if not ask('Enable booting to SafeMode (with Networking)?'):
             abort()
         
-        # Edit BCD to set safeboot as default
-        cmd = ['bcdedit', '/set', '{default}', 'safeboot', 'network']
-        run_program(cmd, check=False)
-        
-        # Enable MSI access under safemode
-        cmd = ['reg', 'add', REG_MSISERVER, '/f']
-        run_program(cmd, check=False)
-        cmd = ['reg', 'add', REG_MSISERVER, '/ve',
-                '/t', 'REG_SZ', '/d', 'Service', '/f']
-        run_program(cmd, check=False)
-    
-        ## Done ##
-        pause('Press Enter to reboot...')
-        cmd = ['shutdown', '-r', '-t', '3']
-        run_program(cmd, check=False)
+        # Configure SafeMode
+        try_and_print(message='Set BCD option...',
+            function=enable_safemode, other_results=other_results)
+        try_and_print(message='Enable MSI in SafeMode...',
+            function=enable_safemode_msi, other_results=other_results)
         
         # Done
+        print_standard('\nDone.')
+        pause('Press Enter to reboot...')
+        reboot()
         exit_script()
     except SystemExit:
         pass
diff --git a/.bin/Scripts/safemode_exit.py b/.bin/Scripts/safemode_exit.py
index 267aac60..1449cff5 100644
--- a/.bin/Scripts/safemode_exit.py
+++ b/.bin/Scripts/safemode_exit.py
@@ -3,13 +3,10 @@
 import os
 import sys
 
-# STATIC VARIABLES
-REG_MSISERVER = r'HKLM\SYSTEM\CurrentControlSet\Control\SafeBoot\Network\MSIServer'
-
 # Init
 os.chdir(os.path.dirname(os.path.realpath(__file__)))
 sys.path.append(os.getcwd())
-from functions.common import *
+from functions.safemode import *
 init_global_vars()
 os.system('title {}: SafeMode Tool'.format(KIT_NAME_FULL))
 
@@ -17,24 +14,23 @@ if __name__ == '__main__':
     try:
         clear_screen()
         print_info('{}: SafeMode Tool\n'.format(KIT_NAME_FULL))
+        other_results = {
+            'Error': {'CalledProcessError':   'Unknown Error'},
+            'Warning': {}}
+        
         if not ask('Disable booting to SafeMode?'):
             abort()
         
-        # Edit BCD to remove safeboot value
-        for boot in ['{current}', '{default}']:
-            cmd = ['bcdedit', '/deletevalue', boot, 'safeboot']
-            run_program(cmd, check=False)
-        
-        # Disable MSI access under safemode
-        cmd = ['reg', 'delete', REG_MSISERVER, '/f']
-        run_program(cmd, check=False)
-        
-        ## Done ##
-        pause('Press Enter to reboot...')
-        cmd = ['shutdown', '-r', '-t', '3']
-        run_program(cmd, check=False)
+        # Configure SafeMode
+        try_and_print(message='Remove BCD option...',
+            function=disable_safemode, other_results=other_results)
+        try_and_print(message='Disable MSI in SafeMode...',
+            function=disable_safemode_msi, other_results=other_results)
         
         # Done
+        print_standard('\nDone.')
+        pause('Press Enter to reboot...')
+        reboot()
         exit_script()
     except SystemExit:
         pass