Add presets to AutoRepairs

scripts/auto_repairs.py

@@ -13,7 +13,7 @@ import wk # pylint: disable=wrong-import-position
 class MenuEntry():
   # pylint: disable=too-few-public-methods
   """Simple class to allow cleaner code below."""
-  def __init__(self, name, function=None, **kwargs):
+  def __init__(self, name, function=None, selected=True, **kwargs):
     self.name = name
 
     # Color reboot entries
@@ -27,7 +27,7 @@ class MenuEntry():
     # Set details
     self.details = {
       'Function': function,
-      'Selected': True,
+      'Selected': selected,
       **kwargs,
       }
 
@@ -74,23 +74,33 @@ BASE_MENUS = {
       ),
     },
   'Options': (
-    MenuEntry('Kill Explorer'),
+    MenuEntry('Kill Explorer', selected=False),
     MenuEntry('Run RKill'),
     MenuEntry('Run TDSSKiller (once)'),
     MenuEntry('Sync Clock'),
-    MenuEntry('Use Autologon'),
+    MenuEntry('Use Autologon', selected=False),
     ),
   'Actions': (
+    MenuEntry('Load Preset'),
     MenuEntry('Options'),
     MenuEntry('Start', Separator=True),
     MenuEntry('Quit'),
     ),
   }
-
+PRESETS = {
+  'Default': { # Will be expanded at runtime using BASE_MENUS
+    'Options': (
+      'Run RKill',
+      'Run TDSSKiller (once)',
+      'Sync Clock',
+      ),
+    },
+  'Custom': {}, # Will remain empty at runtime
+  }
 
 if __name__ == '__main__':
   try:
-    wk.repairs.win.run_auto_repairs(BASE_MENUS)
+    wk.repairs.win.run_auto_repairs(BASE_MENUS, PRESETS)
   except KeyboardInterrupt:
     wk.std.abort()
   except SystemExit:

scripts/wk/repairs/win.py

@@ -87,6 +87,7 @@ GPUPDATE_SUCCESS_STRINGS = (
   'User Policy update has completed successfully.',
   )
 IN_CONEMU = 'ConEmuPID' in os.environ
+MENU_PRESETS = Menu()
 PROGRAMFILES_32 = os.environ.get(
   'PROGRAMFILES(X86)', os.environ.get(
     'PROGRAMFILES', r'C:\Program Files (x86)',
@@ -111,7 +112,7 @@ for error in ('CalledProcessError', 'FileNotFoundError'):
 
 
 # Auto Repairs
-def build_menus(base_menus, title):
+def build_menus(base_menus, title, presets):
   """Build menus, returns dict."""
   menus = {}
   menus['Main'] = Menu(title=f'{title}\n{color_string("Main Menu", "GREEN")}')
@@ -158,6 +159,22 @@ def build_menus(base_menus, title):
           sys.modules[__name__], _function,
           )
 
+  # Update presets
+  for group, entries in base_menus['Groups'].items():
+    presets['Default'][group] = tuple(
+      entry.name for entry in entries if entry.details['Selected']
+      )
+
+  # Update presets Menu
+  MENU_PRESETS.title = f'{title}\n{color_string("Load Preset", "GREEN")}'
+  MENU_PRESETS.add_option('Default')
+  for name in presets:
+    MENU_PRESETS.add_option(name)
+  MENU_PRESETS.add_option('Custom')
+  MENU_PRESETS.add_action('Main Menu')
+  MENU_PRESETS.add_action('Quit')
+  MENU_PRESETS.update()
+
   # Done
   return menus
 
@@ -238,13 +255,6 @@ def init(menus):
   """Initialize Auto Repairs."""
   session_started = is_session_started()
 
-  # Run AVRemover
-  if not session_started:
-    TRY_PRINT.run(
-      'Run AVRemover...', run_tool, 'AVRemover', 'AVRemover',
-      download=True, msg_good='DONE',
-      )
-
   # Check if autologon is needed
   if not session_started and is_autologon_enabled():
     LOG.warning('Skipping Autologon to preserve current settings.')
@@ -308,6 +318,11 @@ def init_session(options):
     set_timezone(WINDOWS_TIME_ZONE)
 
   # One-time tasks
+  if not session_started:
+    TRY_PRINT.run(
+      'Run AVRemover...', run_tool, 'AVRemover', 'AVRemover',
+      download=True, msg_good='DONE',
+      )
   if options['Run TDSSKiller (once)']['Selected']:
     TRY_PRINT.run('Running TDSSKiller...', run_tdsskiller, msg_good='DONE')
   print('')
@@ -344,6 +359,32 @@ def is_session_started():
   return session_started
 
 
+def load_preset(menus, presets, title, enable_menu_exit=True):
+  """Load menu settings from preset and ask selection question(s)."""
+  if not enable_menu_exit:
+    MENU_PRESETS.actions['Main Menu'].update({'Disabled':True, 'Hidden':True})
+
+  # Get selection
+  selection = MENU_PRESETS.simple_select()
+
+  # Exit early
+  if 'Main Menu' in selection:
+    return
+  if 'Quit' in selection:
+    raise SystemExit
+
+  # Load preset
+  preset = presets[selection[0]]
+  for group, menu in menus.items():
+    group_enabled = group in preset
+    for name in menu.options:
+      value = group_enabled and name in preset[group]
+      menu.options[name]['Selected'] = value
+
+  # Re-enable Main Menu action if disabled
+  MENU_PRESETS.actions['Main Menu'].update({'Disabled':False, 'Hidden':False})
+
+
 def load_settings(menus):
   """Load session settings from the registry."""
   for group, menu in menus.items():
@@ -353,7 +394,7 @@ def load_settings(menus):
       menu.options[name].update(get_entry_settings(group, name))
 
 
-def run_auto_repairs(base_menus):
+def run_auto_repairs(base_menus, presets):
   """Run Auto Repairs."""
   set_log_path()
   title = f'{KIT_NAME_FULL}: Auto Repairs'
@@ -364,7 +405,10 @@ def run_auto_repairs(base_menus):
 
   # Generate menus
   print_standard('Initializing...')
-  menus = build_menus(base_menus, title)
+  menus = build_menus(base_menus, title, presets)
+
+  # Get repair preset
+  load_preset(menus, presets, title, enable_menu_exit=False)
 
   # Init
   try:
@@ -376,7 +420,7 @@ def run_auto_repairs(base_menus):
   # Show Menu
   if session_started is None or not session_started:
     try:
-      show_main_menu(base_menus, menus)
+      show_main_menu(base_menus, menus, presets, title)
     except SystemExit:
       if ask('End session?'):
         end_session()
@@ -508,13 +552,15 @@ def set_log_path():
     update_log_path(dest_name=log_path, keep_history=False, timestamp=False, append=True)
 
 
-def show_main_menu(base_menus, menus):
+def show_main_menu(base_menus, menus, presets, title):
   """Show main menu and handle actions."""
   while True:
     update_main_menu(menus)
     selection = menus['Main'].simple_select(update=False)
     if selection[0] in base_menus['Groups'] or selection[0] == 'Options':
       show_sub_menu(menus[selection[0]])
+    if selection[0] == 'Load Preset':
+      load_preset(menus, presets, title)
     elif 'Start' in selection:
       break
     elif 'Quit' in selection: