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?)
This commit is contained in:
parent
10b443f0f5
commit
9a77a5cb9b
3 changed files with 174 additions and 53 deletions
|
|
@ -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': (
|
||||
MenuEntry('Kill Explorer'),
|
||||
MenuEntry('Run RKill at startup'),
|
||||
MenuEntry('Use Autologon'),
|
||||
'Options': (
|
||||
MenuEntry('Kill Explorer', placeholder_function),
|
||||
MenuEntry('Run RKill at startup', placeholder_function),
|
||||
MenuEntry('Use Autologon', placeholder_function),
|
||||
),
|
||||
'Actions': (
|
||||
MenuEntry('Options'),
|
||||
MenuEntry('Start', Separator=True),
|
||||
MenuEntry('Quit'),
|
||||
MenuEntry('Options', placeholder_function),
|
||||
MenuEntry('Start', placeholder_function, Separator=True),
|
||||
MenuEntry('Quit', placeholder_function),
|
||||
),
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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']:
|
||||
menus['Options'].add_toggle(entry.name, entry.details)
|
||||
menus['Options'].add_action('Main Menu')
|
||||
for entry in base_menus['Options']:
|
||||
menus['Options'].add_option(entry.name, entry.details)
|
||||
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:
|
||||
update_main_menu(menus)
|
||||
selection = menus['Main'].simple_select(update=False)
|
||||
if selection[0] in base_menus['Options'] or selection[0] == 'Options':
|
||||
menus[selection[0]].advanced_select()
|
||||
elif 'Start' in selection:
|
||||
break
|
||||
elif 'Quit' in selection:
|
||||
if ask('End session?'):
|
||||
end_session()
|
||||
raise SystemExit
|
||||
try:
|
||||
show_main_menu(base_menus, menus)
|
||||
except SystemExit:
|
||||
if ask('End session?'):
|
||||
end_session()
|
||||
raise
|
||||
|
||||
# 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():
|
||||
if details.get('Done', None) == '1':
|
||||
continue
|
||||
details['Function'](group, name)
|
||||
run_group(group, menu)
|
||||
|
||||
# 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'):
|
||||
"""Load session settings from the registry."""
|
||||
def run_group(group, menu):
|
||||
"""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')
|
||||
reg_set_value('HKCU', key_path, 'Failed', '1' if failed else '0', 'SZ')
|
||||
reg_set_value('HKCU', key_path, 'Result', result, 'SZ')
|
||||
for value_name, data in kwargs.items():
|
||||
if isinstance(data, bool):
|
||||
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'],
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
"""Display info using standard WIDTH and INDENT."""
|
||||
def show_data(message, data, color=None, indent=None, width=None):
|
||||
"""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='',
|
||||
|
|
|
|||
Loading…
Reference in a new issue