From d51bf9fe631bfd773bd5acf3ef950d0a2200f48a Mon Sep 17 00:00:00 2001 From: 2Shirt <2xShirt@gmail.com> Date: Wed, 28 Sep 2022 00:02:18 -0700 Subject: [PATCH] Add presets to AutoRepairs --- scripts/auto_repairs.py | 22 ++++++++---- scripts/wk/repairs/win.py | 70 ++++++++++++++++++++++++++++++++------- 2 files changed, 74 insertions(+), 18 deletions(-) diff --git a/scripts/auto_repairs.py b/scripts/auto_repairs.py index 8d572d07..2d033533 100644 --- a/scripts/auto_repairs.py +++ b/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: diff --git a/scripts/wk/repairs/win.py b/scripts/wk/repairs/win.py index 6cdc0bd5..2fb22db3 100644 --- a/scripts/wk/repairs/win.py +++ b/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: