From 125907ed3afd1c5eb44d1088c9e7b4d3a6cffffd Mon Sep 17 00:00:00 2001 From: 2Shirt <2xShirt@gmail.com> Date: Wed, 21 Apr 2021 05:15:12 -0600 Subject: [PATCH] Add Auto Repairs menus --- scripts/auto_repairs.py | 2 +- scripts/wk/repairs/win.py | 265 +++++++++++++++++++++++++++++--------- scripts/wk/std.py | 85 ++++++------ 3 files changed, 250 insertions(+), 102 deletions(-) diff --git a/scripts/auto_repairs.py b/scripts/auto_repairs.py index 4409e3ab..f4c471fc 100644 --- a/scripts/auto_repairs.py +++ b/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 diff --git a/scripts/wk/repairs/win.py b/scripts/wk/repairs/win.py index 6cba5133..324d05a7 100644 --- a/scripts/wk/repairs/win.py +++ b/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,117 +32,242 @@ 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') - reg_delete_value('HKCU', AUTO_REPAIR_KEY, 'SessionStarted') + 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', ] - run_program(cmd) + try: + run_program(cmd) + except CalledProcessError: + LOG.error("Failed to remove scheduled task or it doesn't exist.") # Disable Autologon - run_tool('Sysinternals', 'Autologon') - reg_set_value( - 'HKLM', r'Software\Microsoft\Windows NT\CurrentVersion\Winlogon', - 'AutoAdminLogon', '0', 'SZ', - ) - reg_delete_value( - 'HKLM', r'Software\Microsoft\Windows NT\CurrentVersion\Winlogon', - 'DefaultUserName', - ) - reg_delete_value( - 'HKLM', r'Software\Microsoft\Windows NT\CurrentVersion\Winlogon', - 'DefaultDomainName', - ) + 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', + 'AutoAdminLogon', '0', 'SZ', + ) + reg_delete_value( + 'HKLM', r'Software\Microsoft\Windows NT\CurrentVersion\Winlogon', + 'DefaultUserName', + ) + reg_delete_value( + 'HKLM', r'Software\Microsoft\Windows NT\CurrentVersion\Winlogon', + 'DefaultDomainName', + ) def init(): - """Initialize AutoRepair.""" - session_started = False - try: - session_started = reg_read_value('HKCU', AUTO_REPAIR_KEY, 'SessionStarted') - except FileNotFoundError: - pass + """Initialize Auto Repairs.""" + session_started = is_session_started() # Start or resume a repair session - if not session_started: - init_run() - init_session() - else: + if session_started: 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() + 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) print('') - init_run() # Done return session_started -def init_run(): - """Initialize AutoRepair Run.""" - atexit.register(start_explorer) +def init_run(options): + """Initialize Auto Repairs Run.""" + if options['Kill Explorer']['Selected']: + atexit.register(start_explorer) + kill_explorer() # TODO: Sync Clock - kill_explorer() # TODO: RKill -def init_session(): - """Initialize AutoRepair session.""" +def init_session(options): + """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 - # TODO: Enable and create restore point - run_tool('Sysinternals', 'Autologon') - # TODO: Disable Windows updates - # TODO: Reset Windows updates + if options['Use Autologon']['Selected']: + run_tool('Sysinternals', 'Autologon') reboot() -def run_auto_repair(): - """Run AutoRepair.""" - update_log_path(dest_name='Auto-Repair Tool', timestamp=True) - title = f'{KIT_NAME_FULL}: Auto-Repair Tool' +def is_session_started(): + """Check if session was started, returns bool.""" + session_started = False + 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 - session_started = init() - if not session_started: - # TODO: Show Menu - pass + # Generate menus + menus = build_menus(title) + + # Init + try: + session_started = init() + except KeyboardInterrupt: + # 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.""" diff --git a/scripts/wk/std.py b/scripts/wk/std.py index 2768fea8..f55cd1b3 100644 --- a/scripts/wk/std.py +++ b/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