Add Auto Repairs menus

scripts/auto_repairs.py

@@ -11,7 +11,7 @@ import wk
 
 if __name__ == '__main__':
   try:
-    wk.repairs.win.run_auto_repair()
+    wk.repairs.win.run_auto_repairs()
   except SystemExit:
     raise
   except: #pylint: disable=bare-except

scripts/wk/repairs/win.py

@@ -7,6 +7,8 @@ import os
 import platform
 import sys
 
+from subprocess import CalledProcessError
+
 from wk.cfg.main  import KIT_NAME_FULL
 from wk.exe       import get_procs, run_program, popen_program, wait_for_procs
 from wk.kit.tools import run_tool
@@ -15,9 +17,11 @@ from wk.os.win    import reg_delete_value, reg_read_value, reg_set_value
 from wk.std       import (
   GenericError,
   GenericWarning,
+  Menu,
   TryAndPrint,
-  abort,
+  ask,
   clear_screen,
+  color_string,
   pause,
   print_info,
   set_title,
@@ -28,26 +32,125 @@ from wk.std       import (
 # STATIC VARIABLES
 LOG = logging.getLogger(__name__)
 AUTO_REPAIR_DELAY_IN_SECONDS = 30
-AUTO_REPAIR_KEY = fr'Software\{KIT_NAME_FULL}\AutoRepair'
+AUTO_REPAIR_KEY = fr'Software\{KIT_NAME_FULL}\Auto Repairs'
 CONEMU = 'ConEmuPID' in os.environ
+MENUS = {
+  'Options': {
+    'Backup Settings': (
+      ('Enable System Restore', None),
+      ('Create System Restore', None),
+      ('Backup Browsers', None),
+      ('Backup Power Plans', None),
+      ),
+    'Windows Repairs': (
+      ('Disable Windows Updates', None),
+      ('Reset Windows Updates', None),
+      ('-Reboot-', None),
+      ('CHKDSK', None),
+      ('DISM RestoreHealth', None),
+      ('SFC Scan', None),
+      ('Fix File Associations', None),
+      ('Clear Proxy Settings', None),
+      ('Disable Pending Renames', None),
+      ('Registry Repairs', None),
+      ('Repair Safe Mode', None),
+      ('Reset Windows Policies', None),
+      ),
+    'Malware Cleanup': (
+      ('BleachBit', None),
+      ('HitmanPro', None),
+      ('KVRT', None),
+      ('Windows Defender', None),
+      ('-Reboot-', None),
+      ),
+    'Manual Steps': (
+      ('AdwCleaner', None),
+      ('IO Bit Uninstaller', None),
+      ('Enable Windows Updates', None),
+      ),
+    },
+  'Toggles': (
+    'Kill Explorer',
+    'Run RKill at startup',
+    'Use Autologon',
+    ),
+  'Actions': (
+    'Options',
+    'Start',
+    'Quit',
+    ),
+  }
 OS_VERSION = float(platform.win32_ver()[0])
 TRY_PRINT = TryAndPrint()
 
 
-# AutoRepair Functions
+# Auto Repairs Functions
+def build_menus(title):
+  """Build menus, returns dict."""
+  menus = {}
+  menus['Main'] = Menu(title=f'{title}\n{color_string("Main Menu", "GREEN")}')
+
+  # Run groups
+  for group, items in MENUS['Options'].items():
+    menus[group] = Menu(title=f'{title}\n{color_string(group, "GREEN")}')
+    menus[group].disabled_str = color_string('Forced', 'ORANGE')
+    for item in items:
+      menus[group].add_option(item[0], {'Display Name': item[0], 'Selected': True})
+    if '-Reboot-' in menus[group].options:
+      menus[group].options['-Reboot-']['Disabled'] = True
+    menus[group].add_action('Main Menu')
+  for option in MENUS['Options']:
+    menus['Main'].add_option(option, {'Display Name': option, 'Selected': True})
+
+  # Actions
+  for action in MENUS['Actions']:
+    menus['Main'].add_action(action, {'Display Name': action, 'Selected': True})
+  menus['Main'].actions['Start']['Separator'] = True
+
+  # Options
+  menus['Options'] = Menu(title=f'{title}\n{color_string("Options", "GREEN")}')
+  for toggle in MENUS['Toggles']:
+    menus['Options'].add_toggle(toggle, {'Display Name': toggle, 'Selected': True})
+  menus['Options'].add_action('Main Menu')
+
+  # Initialize main menu display names
+  menus['Main'].update()
+
+  # Done
+  return menus
+
+
 def end_session():
-  """End AutoRepair session."""
+  """End Auto Repairs session."""
   print_info('Ending repair session')
+  auto_admin_logon = '0'
+
+  # Delete Auto Repairs session key
+  try:
     reg_delete_value('HKCU', AUTO_REPAIR_KEY, 'SessionStarted')
+  except FileNotFoundError:
+    LOG.error('Ending repair session but session not started.')
 
   # Remove logon task
   cmd = [
     'schtasks', '/delete', '/f',
-    '/tn', f'{KIT_NAME_FULL}-AutoRepair',
+    '/tn', f'{KIT_NAME_FULL}-AutoRepairs',
     ]
+  try:
     run_program(cmd)
+  except CalledProcessError:
+    LOG.error("Failed to remove scheduled task or it doesn't exist.")
 
   # Disable Autologon
+  try:
+    auto_admin_logon = reg_read_value(
+      'HKLM', r'Software\Microsoft\Windows NT\CurrentVersion\Winlogon',
+      'AutoAdminLogon',
+      )
+  except FileNotFoundError:
+    # Ignore and assume it's disabled
+    return
+  if auto_admin_logon != '0':
     run_tool('Sysinternals', 'Autologon')
     reg_set_value(
       'HKLM', r'Software\Microsoft\Windows NT\CurrentVersion\Winlogon',
@@ -64,81 +167,107 @@ def end_session():
 
 
 def init():
-  """Initialize AutoRepair."""
+  """Initialize Auto Repairs."""
-  session_started = False
+  session_started = is_session_started()
-  try:
-    session_started = reg_read_value('HKCU', AUTO_REPAIR_KEY, 'SessionStarted')
-  except FileNotFoundError:
-    pass
 
   # Start or resume a repair session
-  if not session_started:
+  if session_started:
-    init_run()
-    init_session()
-  else:
     print_info('Resuming session, press CTRL+c to cancel')
-    try:
     for _x in range(AUTO_REPAIR_DELAY_IN_SECONDS, 0, -1):
       print(f'  {_x} second{"" if _x==1 else "s"} remaining...  \r', end='')
       sleep(1)
-    except KeyboardInterrupt:
-      abort()
     print('')
-    init_run()
 
   # Done
   return session_started
 
 
-def init_run():
+def init_run(options):
-  """Initialize AutoRepair Run."""
+  """Initialize Auto Repairs Run."""
+  if options['Kill Explorer']['Selected']:
     atexit.register(start_explorer)
-  # TODO: Sync Clock
     kill_explorer()
+  # TODO: Sync Clock
   # TODO: RKill
 
 
-def init_session():
+def init_session(options):
-  """Initialize AutoRepair session."""
+  """Initialize Auto Repairs session."""
   print_info('Starting repair session')
   reg_set_value('HKCU', AUTO_REPAIR_KEY, 'SessionStarted', 1, 'DWORD')
 
-  # Create logon task for AutoRepair
+  # Create logon task for Auto Repairs
   cmd = [
     'schtasks', '/create', '/f',
     '/sc', 'ONLOGON',
-    '/tn', f'{KIT_NAME_FULL}-AutoRepair',
+    '/tn', f'{KIT_NAME_FULL}-AutoRepairs',
     '/rl', 'HIGHEST',
     '/tr', fr'C:\Windows\System32\cmd.exe "/C {sys.executable} {sys.argv[0]}"',
     ]
   run_program(cmd)
 
   # One-time tasks
-  # TODO: Backup Registry
+  if options['Use Autologon']['Selected']:
-  # TODO: Enable and create restore point
     run_tool('Sysinternals', 'Autologon')
-  # TODO: Disable Windows updates
-  # TODO: Reset Windows updates
   reboot()
 
 
-def run_auto_repair():
+def is_session_started():
-  """Run AutoRepair."""
+  """Check if session was started, returns bool."""
-  update_log_path(dest_name='Auto-Repair Tool', timestamp=True)
+  session_started = False
-  title = f'{KIT_NAME_FULL}: Auto-Repair Tool'
+  try:
+    session_started = reg_read_value('HKCU', AUTO_REPAIR_KEY, 'SessionStarted')
+  except FileNotFoundError:
+    pass
+
+  # Done
+  return session_started
+
+
+def run_auto_repairs():
+  """Run Auto Repairs."""
+  update_log_path(dest_name='Auto Repairs', timestamp=True)
+  title = f'{KIT_NAME_FULL}: Auto Repairs'
   clear_screen()
   set_title(title)
   print_info(title)
   print('')
 
-  # Show Menu on first run
+  # Generate menus
+  menus = build_menus(title)
+
+  # Init
+  try:
     session_started = init()
-  if not session_started:
+  except KeyboardInterrupt:
-    # TODO: Show Menu
+    # Assuming session was started and resume countdown was interrupted
+    session_started = None
+
+  # Show Menu
+  if session_started is None or not session_started:
+    while True:
+      update_main_menu(menus)
+      selection = menus['Main'].simple_select(update=False)
+      if selection[0] in MENUS['Options'] or selection[0] == 'Options':
+        menus[selection[0]].advanced_select()
+      elif 'Start' in selection:
+        # TODO: Run repairs
         pass
+      elif 'Quit' in selection:
+        if ask('End session?'):
+          end_session()
+        raise SystemExit
+
+    # Re-check if a repair session was started
+    if session_started is None:
+      session_started = is_session_started()
 
   # Run repairs
   # TODO: Run repairs
+  init_run(menus['Options'])
+  if not session_started:
+    init_session(menus['Options'])
+  atexit.unregister(start_explorer)
   end_session()
 
   # Done
@@ -146,6 +275,24 @@ def run_auto_repair():
   pause('Press Enter to exit...')
 
 
+def update_main_menu(menus):
+  """Update main menu based on current selections."""
+  index = 1
+  skip = '-Reboot-'
+  for name in menus['Main'].options:
+    checkmark = ' '
+    selected = [
+      _v['Selected'] for _k, _v in menus[name].options.items() if _k != skip
+      ]
+    if all(selected):
+      checkmark = '✓'
+    elif any(selected):
+      checkmark = '-'
+    display_name = f' {index}: [{checkmark}] {name}'
+    index += 1
+    menus['Main'].options[name]['Display Name'] = display_name
+
+
 # OS Built-in Functions
 def kill_explorer():
   """Kill all Explorer processes."""

scripts/wk/std.py

@@ -228,44 +228,6 @@ class Menu():
     # Done
     return resolved_selection
 
-  def _update(self, single_selection=True, settings_mode=False):
-    """Update menu items in preparation for printing to screen."""
-    index = 0
-
-    # Fix selection status for sets
-    for set_details in self.sets.values():
-      set_selected = True
-      set_targets = set_details['Targets']
-      for option, option_details in self.options.items():
-        if option in set_targets and not option_details['Selected']:
-          set_selected = False
-        elif option not in set_targets and option_details['Selected']:
-          set_selected = False
-      set_details['Selected'] = set_selected
-
-    # Numbered sections
-    for section in (self.sets, self.toggles, self.options):
-      for name, details in section.items():
-        if details.get('Hidden', False):
-          # Skip hidden lines and don't increment index
-          continue
-        index += 1
-        details['Display Name'] = self._get_display_name(
-          name,
-          details,
-          index=index,
-          no_checkboxes=single_selection,
-          setting_item=settings_mode,
-          )
-
-    # Actions
-    for name, details in self.actions.items():
-      details['Display Name'] = self._get_display_name(
-        name,
-        details,
-        no_checkboxes=True,
-        )
-
   def _update_entry_selection_status(self, entry, toggle=True, status=None):
     """Update entry selection status either directly or by toggling."""
     if entry in self.sets:
@@ -340,7 +302,7 @@ class Menu():
     NOTE: Menu is displayed until an action entry is selected.
     """
     while True:
-      self._update(single_selection=False)
+      self.update(single_selection=False)
       user_selection = self._user_select(prompt)
       selected_entry = self._resolve_selection(user_selection)
       if user_selection.isnumeric():
@@ -364,7 +326,7 @@ class Menu():
       }
 
     while True:
-      self._update(single_selection=True, settings_mode=True)
+      self.update(single_selection=True, settings_mode=True)
       user_selection = self._user_select(prompt)
       selected_entry = self._resolve_selection(user_selection)
       if user_selection.isnumeric():
@@ -381,12 +343,51 @@ class Menu():
     # Done
     return selected_entry
 
-  def simple_select(self, prompt='Please make a selection: '):
+  def simple_select(self, prompt='Please make a selection: ', update=True):
     """Display menu and make a single selection, returns tuple."""
-    self._update()
+    if update:
+      self.update()
     user_selection = self._user_select(prompt)
     return self._resolve_selection(user_selection)
 
+  def update(self, single_selection=True, settings_mode=False):
+    """Update menu items in preparation for printing to screen."""
+    index = 0
+
+    # Fix selection status for sets
+    for set_details in self.sets.values():
+      set_selected = True
+      set_targets = set_details['Targets']
+      for option, option_details in self.options.items():
+        if option in set_targets and not option_details['Selected']:
+          set_selected = False
+        elif option not in set_targets and option_details['Selected']:
+          set_selected = False
+      set_details['Selected'] = set_selected
+
+    # Numbered sections
+    for section in (self.sets, self.toggles, self.options):
+      for name, details in section.items():
+        if details.get('Hidden', False):
+          # Skip hidden lines and don't increment index
+          continue
+        index += 1
+        details['Display Name'] = self._get_display_name(
+          name,
+          details,
+          index=index,
+          no_checkboxes=single_selection,
+          setting_item=settings_mode,
+          )
+
+    # Actions
+    for name, details in self.actions.items():
+      details['Display Name'] = self._get_display_name(
+        name,
+        details,
+        no_checkboxes=True,
+        )
+
 
 class TryAndPrint():
   # pylint: disable=too-many-instance-attributes