From b41027562a13d3ee411e82342a51ec24595bc086 Mon Sep 17 00:00:00 2001
From: 2Shirt <2xShirt@gmail.com>
Date: Sun, 15 Sep 2019 18:47:52 -0700
Subject: [PATCH] Added SafeMode enter/exit sections

---
 .../outer_scripts_to_review/safemode_enter.py | 39 -----------------
 .../outer_scripts_to_review/safemode_exit.py  | 39 -----------------
 scripts/safemode_enter.py                     | 37 ++++++++++++++++
 scripts/safemode_exit.py                      | 37 ++++++++++++++++
 scripts/wk/__init__.py                        |  2 +-
 scripts/wk/cfg/__init__.py                    |  2 +-
 scripts/wk/cfg/log.py                         |  2 +-
 scripts/wk/cfg/main.py                        |  4 +-
 scripts/wk/cfg/net.py                         |  2 +-
 scripts/wk/hw/__init__.py                     |  2 +-
 scripts/wk/io.py                              |  2 +-
 scripts/wk/kit/__init__.py                    |  2 +-
 scripts/wk/log.py                             |  2 +-
 scripts/wk/os/__init__.py                     |  8 +++-
 scripts/wk/os/win.py                          | 42 +++++++++++++++++++
 scripts/wk/std.py                             |  4 +-
 scripts/wk/sw/__init__.py                     |  2 +-
 17 files changed, 136 insertions(+), 92 deletions(-)
 delete mode 100644 scripts/outer_scripts_to_review/safemode_enter.py
 delete mode 100644 scripts/outer_scripts_to_review/safemode_exit.py
 create mode 100644 scripts/safemode_enter.py
 create mode 100644 scripts/safemode_exit.py
 create mode 100644 scripts/wk/os/win.py

diff --git a/scripts/outer_scripts_to_review/safemode_enter.py b/scripts/outer_scripts_to_review/safemode_enter.py
deleted file mode 100644
index de9ad119..00000000
--- a/scripts/outer_scripts_to_review/safemode_enter.py
+++ /dev/null
@@ -1,39 +0,0 @@
-# Wizard Kit: Enter SafeMode by editing the BCD
-
-import os
-import sys
-
-# Init
-sys.path.append(os.path.dirname(os.path.realpath(__file__)))
-from functions.safemode import *
-init_global_vars()
-os.system('title {}: SafeMode Tool'.format(KIT_NAME_FULL))
-
-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()
-
-    # 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 as sys_exit:
-    exit_script(sys_exit.code)
-  except:
-    major_exception()
-
-# vim: sts=2 sw=2 ts=2
diff --git a/scripts/outer_scripts_to_review/safemode_exit.py b/scripts/outer_scripts_to_review/safemode_exit.py
deleted file mode 100644
index 6c47b02d..00000000
--- a/scripts/outer_scripts_to_review/safemode_exit.py
+++ /dev/null
@@ -1,39 +0,0 @@
-# Wizard Kit: Exit SafeMode by editing the BCD
-
-import os
-import sys
-
-# Init
-sys.path.append(os.path.dirname(os.path.realpath(__file__)))
-from functions.safemode import *
-init_global_vars()
-os.system('title {}: SafeMode Tool'.format(KIT_NAME_FULL))
-
-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()
-
-    # 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 as sys_exit:
-    exit_script(sys_exit.code)
-  except:
-    major_exception()
-
-# vim: sts=2 sw=2 ts=2
diff --git a/scripts/safemode_enter.py b/scripts/safemode_enter.py
new file mode 100644
index 00000000..87660560
--- /dev/null
+++ b/scripts/safemode_enter.py
@@ -0,0 +1,37 @@
+"""Wizard Kit: Enter SafeMode by editing the BCD"""
+# vim: sts=2 sw=2 ts=2
+
+import wk
+
+
+def main():
+  """Prompt user to enter safe mode."""
+  title = f'{wk.cfg.main.KIT_NAME_FULL}: SafeMode Tool'
+  try_print = wk.std.TryAndPrint()
+  wk.std.clear_screen()
+  wk.std.set_title(title)
+  wk.std.print_info(title)
+  print('')
+
+  # Ask
+  if not wk.std.ask('Enable booting to SafeMode (with Networking)?'):
+    wk.std.abort()
+  print('')
+
+  # Configure SafeMode
+  try_print.run('Set BCD option...', wk.os.win.enable_safemode)
+  try_print.run('Enable MSI in SafeMode...', wk.os.win.enable_safemode_msi)
+
+  # Done
+  print('Done.')
+  wk.std.pause('Press Enter to reboot...')
+  wk.std.run_program('shutdown -r -t 3'.split(), check=False)
+
+
+if __name__ == '__main__':
+  try:
+    main()
+  except SystemExit:
+    raise
+  except: #pylint: disable=bare-except
+    wk.std.major_exception()
diff --git a/scripts/safemode_exit.py b/scripts/safemode_exit.py
new file mode 100644
index 00000000..a1bbfbd1
--- /dev/null
+++ b/scripts/safemode_exit.py
@@ -0,0 +1,37 @@
+"""Wizard Kit: Exit SafeMode by editing the BCD"""
+# vim: sts=2 sw=2 ts=2
+
+import wk
+
+
+def main():
+  """Prompt user to exit safe mode."""
+  title = f'{wk.cfg.main.KIT_NAME_FULL}: SafeMode Tool'
+  try_print = wk.std.TryAndPrint()
+  wk.std.clear_screen()
+  wk.std.set_title(title)
+  wk.std.print_info(title)
+  print('')
+
+  # Ask
+  if not wk.std.ask('Disable booting to SafeMode?'):
+    wk.std.abort()
+  print('')
+
+  # Configure SafeMode
+  try_print.run('Remove BCD option...', wk.os.win.disable_safemode)
+  try_print.run('Disable MSI in SafeMode...', wk.os.win.disable_safemode_msi)
+
+  # Done
+  print('Done.')
+  wk.std.pause('Press Enter to reboot...')
+  wk.std.run_program('shutdown -r -t 3'.split(), check=False)
+
+
+if __name__ == '__main__':
+  try:
+    main()
+  except SystemExit:
+    raise
+  except: #pylint: disable=bare-except
+    wk.std.major_exception()
diff --git a/scripts/wk/__init__.py b/scripts/wk/__init__.py
index d9264d83..959978c7 100644
--- a/scripts/wk/__init__.py
+++ b/scripts/wk/__init__.py
@@ -1,4 +1,4 @@
-'''WizardKit: wk module init'''
+"""WizardKit: wk module init"""
 # vim: sts=2 sw=2 ts=2
 
 from sys import version_info as version
diff --git a/scripts/wk/cfg/__init__.py b/scripts/wk/cfg/__init__.py
index ff52308f..9538f1c8 100644
--- a/scripts/wk/cfg/__init__.py
+++ b/scripts/wk/cfg/__init__.py
@@ -1,4 +1,4 @@
-'''WizardKit: cfg module init'''
+"""WizardKit: cfg module init"""
 
 from wk.cfg import log
 from wk.cfg import main
diff --git a/scripts/wk/cfg/log.py b/scripts/wk/cfg/log.py
index 925be4d3..a9a6e6ca 100644
--- a/scripts/wk/cfg/log.py
+++ b/scripts/wk/cfg/log.py
@@ -1,4 +1,4 @@
-'''WizardKit: Config - Log'''
+"""WizardKit: Config - Log"""
 # vim: sts=2 sw=2 ts=2
 
 
diff --git a/scripts/wk/cfg/main.py b/scripts/wk/cfg/main.py
index 6f5ac11f..2928f54a 100644
--- a/scripts/wk/cfg/main.py
+++ b/scripts/wk/cfg/main.py
@@ -1,6 +1,6 @@
-'''WizardKit: Config - Main
+"""WizardKit: Config - Main
 
-NOTE: A non-standard format is used for BASH/BATCH/PYTHON compatibility'''
+NOTE: A non-standard format is used for BASH/BATCH/PYTHON compatibility"""
 # pylint: disable=bad-whitespace
 # vim: sts=2 sw=2 ts=2
 
diff --git a/scripts/wk/cfg/net.py b/scripts/wk/cfg/net.py
index 8a84620c..344c695b 100644
--- a/scripts/wk/cfg/net.py
+++ b/scripts/wk/cfg/net.py
@@ -1,4 +1,4 @@
-'''WizardKit: Config - Net'''
+"""WizardKit: Config - Net"""
 # vim: sts=2 sw=2 ts=2
 
 
diff --git a/scripts/wk/hw/__init__.py b/scripts/wk/hw/__init__.py
index 6c19c9c7..2b0fbc10 100644
--- a/scripts/wk/hw/__init__.py
+++ b/scripts/wk/hw/__init__.py
@@ -1 +1 @@
-'''WizardKit: hw module init'''
+"""WizardKit: hw module init"""
diff --git a/scripts/wk/io.py b/scripts/wk/io.py
index 5ff75d7b..846db77e 100644
--- a/scripts/wk/io.py
+++ b/scripts/wk/io.py
@@ -1,4 +1,4 @@
-'''WizardKit: I/O Functions'''
+"""WizardKit: I/O Functions"""
 # vim: sts=2 sw=2 ts=2
 
 import os
diff --git a/scripts/wk/kit/__init__.py b/scripts/wk/kit/__init__.py
index 61f06215..2fc43258 100644
--- a/scripts/wk/kit/__init__.py
+++ b/scripts/wk/kit/__init__.py
@@ -1 +1 @@
-'''WizardKit: kit module init'''
+"""WizardKit: kit module init"""
diff --git a/scripts/wk/log.py b/scripts/wk/log.py
index e2ee6336..e930592b 100644
--- a/scripts/wk/log.py
+++ b/scripts/wk/log.py
@@ -1,4 +1,4 @@
-'''WizardKit: Log Functions'''
+"""WizardKit: Log Functions"""
 # vim: sts=2 sw=2 ts=2
 
 import atexit
diff --git a/scripts/wk/os/__init__.py b/scripts/wk/os/__init__.py
index da1056d8..5e58dbc8 100644
--- a/scripts/wk/os/__init__.py
+++ b/scripts/wk/os/__init__.py
@@ -1 +1,7 @@
-'''WizardKit: os module init'''
+"""WizardKit: os module init"""
+# vim: sts=2 sw=2 ts=2
+
+import os
+
+if os.name == 'nt':
+  from wk.os import win
diff --git a/scripts/wk/os/win.py b/scripts/wk/os/win.py
new file mode 100644
index 00000000..85073e46
--- /dev/null
+++ b/scripts/wk/os/win.py
@@ -0,0 +1,42 @@
+"""WizardKit: Windows Functions"""
+# vim: sts=2 sw=2 ts=2
+
+from wk.std import run_program
+
+# STATIC VARIABLES
+REG_MSISERVER = r'HKLM\SYSTEM\CurrentControlSet\Control\SafeBoot\Network\MSIServer'
+
+
+# Functions
+def disable_safemode():
+  """Edit BCD to remove safeboot value."""
+  cmd = ['bcdedit', '/deletevalue', '{default}', 'safeboot']
+  run_program(cmd)
+
+
+def disable_safemode_msi():
+  """Disable MSI access under safemode."""
+  cmd = ['reg', 'delete', REG_MSISERVER, '/f']
+  run_program(cmd)
+
+
+def enable_safemode():
+  """Edit BCD to set safeboot as default."""
+  cmd = ['bcdedit', '/set', '{default}', 'safeboot', 'network']
+  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)
+
+
+if __name__ == '__main__':
+  print("This file is not meant to be called directly.")
diff --git a/scripts/wk/std.py b/scripts/wk/std.py
index 8a5065ca..d1633af6 100644
--- a/scripts/wk/std.py
+++ b/scripts/wk/std.py
@@ -1,4 +1,4 @@
-'''WizardKit: Standard Functions'''
+"""WizardKit: Standard Functions"""
 # pylint: disable=too-many-lines
 # vim: sts=2 sw=2 ts=2
 
@@ -476,7 +476,7 @@ class TryAndPrint():
     if exception_name not in self.list_warnings:
       self.list_warnings.append(exception_name)
 
-  def run_function(
+  def run(
       self, message, function, *args,
       catch_all=True, print_return=False, verbose=False, **kwargs):
     # pylint: disable=catching-non-exception
diff --git a/scripts/wk/sw/__init__.py b/scripts/wk/sw/__init__.py
index f64f5df6..1cc6236f 100644
--- a/scripts/wk/sw/__init__.py
+++ b/scripts/wk/sw/__init__.py
@@ -1 +1 @@
-'''WizardKit: sw module init'''
+"""WizardKit: sw module init"""