Update Auto Repair sections

* Expanded saving/loading settings from registry
* Keep previous selections by default but allow changes
* Print previous session(s) to have the whole session info present
* Changed variable names for clarity (why am I so bad at this?)

scripts/auto_repairs.py

@@ -15,7 +15,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, **kwargs):
     self.name = name
 
     # Color reboot entries
@@ -28,29 +28,29 @@ class MenuEntry():
 
     # Set details
     self.details = {
+      'Function': function,
       'Selected': True,
       **kwargs,
       }
-    if function:
-      self.details['Function'] = function
 
 
 # TODO: Deleteme
 TRY_AND_PRINT = wk.std.TryAndPrint()
+TRY_AND_PRINT.width = 50
 def placeholder_function(group, name):
   TRY_AND_PRINT.run(f'{name}...', time.sleep, random.randint(1, 3))
-  wk.repairs.win.save_settings(group, name, done=True)
+  wk.repairs.win.save_settings(group, name, done=True, result='SUCCESS')
 
 def placeholder_reboot(group, name):
   print('"Rebooting" shortly...')
   time.sleep(random.randint(1, 3))
-  wk.repairs.win.save_settings(group, name, done=True)
+  wk.repairs.win.save_settings(group, name, done=True, result='DONE')
   raise SystemExit
 
 
 # STATIC VARIABLES
 BASE_MENUS = {
-  'Options': {
+  'Groups': {
     'Backup Settings': (
       MenuEntry('Enable RegBack', placeholder_function),
       MenuEntry('Enable System Restore', placeholder_function),
@@ -87,15 +87,15 @@ BASE_MENUS = {
       MenuEntry('Enable Windows Updates', placeholder_function),
       ),
     },
-  'Toggles': (
+  'Options': (
-    MenuEntry('Kill Explorer'),
+    MenuEntry('Kill Explorer', placeholder_function),
-    MenuEntry('Run RKill at startup'),
+    MenuEntry('Run RKill at startup', placeholder_function),
-    MenuEntry('Use Autologon'),
+    MenuEntry('Use Autologon', placeholder_function),
     ),
   'Actions': (
-    MenuEntry('Options'),
+    MenuEntry('Options', placeholder_function),
-    MenuEntry('Start', Separator=True),
+    MenuEntry('Start', placeholder_function, Separator=True),
-    MenuEntry('Quit'),
+    MenuEntry('Quit', placeholder_function),
     ),
   }
 

scripts/wk/repairs/win.py

@@ -34,6 +34,7 @@ from wk.std       import (
   print_standard,
   print_warning,
   set_title,
+  show_data,
   sleep,
   strip_colors,
   )
@@ -45,7 +46,9 @@ AUTO_REPAIR_DELAY_IN_SECONDS = 30
 AUTO_REPAIR_KEY = fr'Software\{KIT_NAME_FULL}\Auto Repairs'
 CONEMU = 'ConEmuPID' in os.environ
 OS_VERSION = float(platform.win32_ver()[0])
+WIDTH = 50
 TRY_PRINT = TryAndPrint()
+TRY_PRINT.width = WIDTH
 TRY_PRINT.verbose = True
 #for error in ('subprocess.CalledProcessError', 'FileNotFoundError'):
 #  TRY_PRINT.add_error(error)
@@ -60,22 +63,30 @@ def build_menus(base_menus, title):
   # Main Menu
   for entry in base_menus['Actions']:
     menus['Main'].add_action(entry.name, entry.details)
-  for group in base_menus['Options']:
+  for group in base_menus['Groups']:
     menus['Main'].add_option(group, {'Selected': True})
 
   # Options
   menus['Options'] = Menu(title=f'{title}\n{color_string("Options", "GREEN")}')
-  for entry in base_menus['Toggles']:
+  for entry in base_menus['Options']:
-    menus['Options'].add_toggle(entry.name, entry.details)
+    menus['Options'].add_option(entry.name, entry.details)
-  menus['Options'].add_action('Main Menu')
+  menus['Options'].add_action('All')
+  menus['Options'].add_action('None')
+  menus['Options'].add_action('Main Menu', {'Separator': True})
+  menus['Options'].add_action('Quit')
 
   # Run groups
-  for group, entries in base_menus['Options'].items():
+  for group, entries in base_menus['Groups'].items():
     menus[group] = Menu(title=f'{title}\n{color_string(group, "GREEN")}')
-    menus[group].disabled_str = 'Done'
+    menus[group].disabled_str = 'Locked'
     for entry in entries:
       menus[group].add_option(entry.name, entry.details)
-    menus[group].add_action('Main Menu')
+    menus[group].add_action('All')
+    menus[group].add_action('None')
+    menus[group].add_action('Select Skipped Entries', {'Separator': True})
+    menus[group].add_action('Unlock All Entries')
+    menus[group].add_action('Main Menu', {'Separator': True})
+    menus[group].add_action('Quit')
 
   # Initialize main menu display names
   menus['Main'].update()
@@ -89,11 +100,16 @@ def end_session():
   print_info('Ending repair session')
   auto_admin_logon = '0'
 
-  # Delete Auto Repairs session key
+  # Delete Auto Repairs keys
   try:
     reg_delete_value('HKCU', AUTO_REPAIR_KEY, 'SessionStarted')
   except FileNotFoundError:
     LOG.error('Ending repair session but session not started.')
+  try:
+    cmd = ['reg', 'delete', fr'HKCU\{AUTO_REPAIR_KEY}', '/f']
+    run_program(cmd)
+  except CalledProcessError:
+    LOG.error('Failed to remote Auto Repairs session settings')
 
   # Remove logon task
   cmd = [
@@ -134,13 +150,17 @@ def get_entry_settings(group, name):
   """Get menu entry settings from the registry, returns dict."""
   key_path = fr'{AUTO_REPAIR_KEY}\{group}\{strip_colors(name)}'
   settings = {}
-  for value in ('Done', 'Failed', 'Result'):
+  for value in ('done', 'failed', 'result', 'selected', 'skipped', 'warning'):
     try:
-      settings[value] = reg_read_value('HKCU', key_path, value)
+      settings[value.title()] = reg_read_value('HKCU', key_path, value)
     except FileNotFoundError:
       # Ignore and use current settings
       pass
 
+  # Disable previously run or skipped entries
+  if settings.get('Done', False) or settings.get('Skipped', False):
+    settings['Disabled'] = True
+
   # Done
   return settings
 
@@ -164,7 +184,7 @@ def init(menus):
 
 def init_run(options):
   """Initialize Auto Repairs Run."""
-  if options.toggles['Kill Explorer']['Selected']:
+  if options['Kill Explorer']['Selected']:
     atexit.register(start_explorer)
     kill_explorer()
   # TODO: Sync Clock
@@ -187,8 +207,9 @@ def init_session(options):
   run_program(cmd)
 
   # One-time tasks
-  if options.toggles['Use Autologon']['Selected']:
+  if options['Use Autologon']['Selected']:
     run_tool('Sysinternals', 'Autologon')
+  # TODO: TDSSKiller?
   # TODO: Re-enable reboot()
 
 
@@ -211,9 +232,6 @@ def load_settings(menus):
       continue
     for name in menu.options:
       menu.options[name].update(get_entry_settings(group, name))
-      if menu.options[name].get('Done', '0') == '1':
-        menu.options[name]['Selected'] = False
-        menu.options[name]['Disabled'] = True
 
 
 def run_auto_repairs(base_menus):
@@ -238,36 +256,29 @@ def run_auto_repairs(base_menus):
 
   # Show Menu
   if session_started is None or not session_started:
-    while True:
+    try:
-      update_main_menu(menus)
+      show_main_menu(base_menus, menus)
-      selection = menus['Main'].simple_select(update=False)
+    except SystemExit:
-      if selection[0] in base_menus['Options'] or selection[0] == 'Options':
+      if ask('End session?'):
-        menus[selection[0]].advanced_select()
+        end_session()
-      elif 'Start' in selection:
+      raise
-        break
-      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()
 
   # Start or resume repairs
-  init_run(menus['Options'])
+  save_selection_settings(menus)
+  init_run(menus['Options'].options)
   if not session_started:
-    init_session(menus['Options'])
+    init_session(menus['Options'].options)
 
   # Run repairs
   clear_screen()
   for group, menu in menus.items():
     if group in ('Main', 'Options'):
       continue
-    for name, details in menu.options.items():
+    run_group(group, menu)
-      if details.get('Done', None) == '1':
-        continue
-      details['Function'](group, name)
 
   # Done
   end_session()
@@ -275,12 +286,119 @@ def run_auto_repairs(base_menus):
   pause('Press Enter to exit...')
 
 
-def save_settings(group, name, done=False, failed=False, result='Unknown'):
+def run_group(group, menu):
-  """Load session settings from the registry."""
+  """Run entries in group if appropriate."""
+  print_info(f'  {group}')
+  for name, details in menu.options.items():
+    name_str = strip_colors(name)
+    skipped = details.get('Skipped', False)
+    done = details.get('Done', False)
+    disabled = details.get('Disabled', False)
+    selected = details.get('Selected', False)
+
+    # Selection changed
+    if (skipped or done) and not disabled and selected:
+      save_settings(group, name, done=False, skipped=False)
+      details['Function'](group, name)
+      continue
+
+    # Previously skipped
+    if skipped:
+      show_data(f'{name_str}...', 'Skipped', 'YELLOW', width=WIDTH)
+      continue
+
+    # Previously ran
+    if done:
+      color = 'GREEN'
+      if details.get('Failed', False):
+        color = 'RED'
+      elif details.get('Warning', False):
+        color = 'YELLOW'
+      show_data(
+        f'{name_str}...', details.get('Result', 'Unknown'), color, width=WIDTH,
+        )
+      continue
+
+    # Not selected
+    if not selected:
+      show_data(f'{name_str}...', 'Skipped', 'YELLOW', width=WIDTH)
+      save_settings(group, name, skipped=True)
+      continue
+
+    # Selected
+    details['Function'](group, name)
+
+
+def save_selection_settings(menus):
+  """Save selections in the registry."""
+  for group, menu in menus.items():
+    if group == 'Main':
+      continue
+    for name, details in menu.options.items():
+      save_settings(
+        group, name,
+        disabled=details.get('Disabled', False),
+        selected=details.get('Selected', False),
+        )
+
+
+def save_settings(group, name, **kwargs):
+  """Save entry settings in the registry."""
   key_path = fr'{AUTO_REPAIR_KEY}\{group}\{strip_colors(name)}'
-  reg_set_value('HKCU', key_path, 'Done', '1' if done else '0', 'SZ')
+  for value_name, data in kwargs.items():
-  reg_set_value('HKCU', key_path, 'Failed', '1' if failed else '0', 'SZ')
+    if isinstance(data, bool):
-  reg_set_value('HKCU', key_path, 'Result', result, 'SZ')
+      data = 1 if data else 0
+    if isinstance(data, int):
+      data_type = 'DWORD'
+    elif isinstance(data, str):
+      data_type = 'SZ'
+    else:
+      raise TypeError(f'Invalid data: "{data}" ({type(data)})')
+    reg_set_value('HKCU', key_path, value_name, data, data_type)
+
+
+def show_main_menu(base_menus, menus):
+  """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]])
+    elif 'Start' in selection:
+      break
+    elif 'Quit' in selection:
+      raise SystemExit
+
+
+def show_sub_menu(menu):
+  """Show sub-menu and handle sub-menu actions."""
+  while True:
+    selection = menu.advanced_select()
+    if 'Main Menu' in selection:
+      break
+    if 'Quit' in selection:
+      raise SystemExit
+
+    # Modify entries
+    key = 'Selected'
+    unlock_all = False
+    unlock_skipped = False
+    if 'Select Skipped Entries' in selection:
+      key = 'Disabled'
+      unlock_skipped = True
+      value = False
+    if 'Unlock All Entries' in selection:
+      key = 'Disabled'
+      unlock_all = True
+      value = False
+    else:
+      value = 'All' in selection
+    for name in menu.options:
+      if (unlock_all
+          or (unlock_skipped and not menu.options[name].get('Selected', False))
+          or not menu.options[name].get('Disabled', False)
+          ):
+        menu.options[name][key] = value
 
 
 def update_main_menu(menus):
@@ -318,6 +436,7 @@ def auto_dism(group, name):
   save_settings(
     group, name, done=True,
     failed=result['Failed'],
+    warning=not result['Failed'] and needs_reboot,
     result=result['Message'],
     )
 

scripts/wk/std.py

@@ -985,11 +985,13 @@ def set_title(title):
     print_error('Setting the title is only supported under Windows.')
 
 
-def show_data(message, data, color=None):
+def show_data(message, data, color=None, indent=None, width=None):
-  """Display info using standard WIDTH and INDENT."""
+  """Display info using default or provided indent and width."""
   colors = (None, color if color else None)
+  indent = INDENT if indent is None else indent
+  width = WIDTH if width is None else width
   print_colored(
-    (f'{" "*INDENT}{message:<{WIDTH}}', data),
+    (f'{" "*indent}{message:<{width}}', data),
     colors,
     log=True,
     sep='',