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

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

.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
+        # Configure SafeMode
-        cmd = ['bcdedit', '/set', '{default}', 'safeboot', 'network']
+        try_and_print(message='Set BCD option...',
-        run_program(cmd, check=False)
+            function=enable_safemode, other_results=other_results)
-        
+        try_and_print(message='Enable MSI in SafeMode...',
-        # Enable MSI access under safemode
+            function=enable_safemode_msi, other_results=other_results)
-        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)
         
         # Done
+        print_standard('\nDone.')
+        pause('Press Enter to reboot...')
+        reboot()
         exit_script()
     except SystemExit:
         pass

.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
+        # Configure SafeMode
-        for boot in ['{current}', '{default}']:
+        try_and_print(message='Remove BCD option...',
-            cmd = ['bcdedit', '/deletevalue', boot, 'safeboot']
+            function=disable_safemode, other_results=other_results)
-            run_program(cmd, check=False)
+        try_and_print(message='Disable MSI in SafeMode...',
-        
+            function=disable_safemode_msi, other_results=other_results)
-        # 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)
         
         # Done
+        print_standard('\nDone.')
+        pause('Press Enter to reboot...')
+        reboot()
         exit_script()
     except SystemExit:
         pass