diff --git a/scripts/auto_repairs.py b/scripts/auto_repairs.py new file mode 100644 index 00000000..4409e3ab --- /dev/null +++ b/scripts/auto_repairs.py @@ -0,0 +1,18 @@ +"""Wizard Kit: Auto-Repair Tool""" +# vim: sts=2 sw=2 ts=2 + +import os +import sys + +os.chdir(os.path.dirname(os.path.realpath(__file__))) +sys.path.append(os.getcwd()) +import wk + + +if __name__ == '__main__': + try: + wk.repairs.win.run_auto_repair() + except SystemExit: + raise + except: #pylint: disable=bare-except + wk.std.major_exception() diff --git a/scripts/wk/__init__.py b/scripts/wk/__init__.py index b6a11b56..25620f35 100644 --- a/scripts/wk/__init__.py +++ b/scripts/wk/__init__.py @@ -13,6 +13,7 @@ from wk import kit from wk import log from wk import net from wk import os +from wk import repairs from wk import std from wk import sw from wk import tmux diff --git a/scripts/wk/os/win.py b/scripts/wk/os/win.py index f4a03f30..20ee4917 100644 --- a/scripts/wk/os/win.py +++ b/scripts/wk/os/win.py @@ -16,8 +16,7 @@ except ImportError as err: raise err from wk.borrowed import acpi -from wk.exe import get_procs, run_program, wait_for_procs -from wk.log import format_log_path +from wk.exe import run_program from wk.std import GenericError, GenericWarning, sleep @@ -314,119 +313,6 @@ def reg_set_value(hive, key, name, data, data_type, option=None): winreg.SetValue(hive, key, data_type, data) -# Repair Functions -def run_chkdsk_offline(): - """Set filesystem 'dirty bit' to force a CHKDSK during startup.""" - cmd = ['fsutil', 'dirty', 'set', os.environ.get('SYSTEMDRIVE', 'C:')] - proc = run_program(cmd, check=False) - - # Check result - if proc.returncode > 0: - raise GenericError('Failed to set dirty bit.') - - -def run_chkdsk_online(): - """Run CHKDSK. - - NOTE: If run on Windows 8+ online repairs are attempted. - """ - cmd = ['CHKDSK', os.environ.get('SYSTEMDRIVE', 'C:')] - if OS_VERSION >= 8: - cmd.extend(['/scan', '/perf']) - if CONEMU: - cmd.extend(['-new_console:n', '-new_console:s33V']) - retried = False - - # Run scan - run_program(cmd, check=False) - try: - proc = get_procs('chkdsk.exe')[0] - return_code = proc.wait() - except IndexError: - # Failed to get CHKDSK process, set return_code to force a retry - return_code = 255 - if return_code > 1: - # Try again - retried = True - run_program(cmd, check=False) - try: - proc = get_procs('chkdsk.exe')[0] - return_code = proc.wait() - except IndexError: - # Failed to get CHKDSK process - return_code = -1 - - # Check result - if return_code == -1: - raise GenericError('Failed to find CHKDSK process.') - if (return_code == 0 and retried) or return_code == 1: - raise GenericWarning('Repaired (or manually aborted)') - if return_code > 1: - raise GenericError('Issue(s) detected') - - -def run_dism(repair=True): - """Run DISM to either scan or repair component store health.""" - conemu_args = ['-new_console:n', '-new_console:s33V'] if CONEMU else [] - - # Bail early - if OS_VERSION < 8: - raise GenericWarning('Unsupported OS') - - # Run (repair) scan - log_path = format_log_path( - log_name=f'DISM_{"Restore" if repair else "Scan"}Health', tool=True, - ) - cmd = [ - 'DISM', '/Online', '/Cleanup-Image', - '/RestoreHealth' if repair else '/ScanHealth', - f'/LogPath:{log_path}', - *conemu_args, - ] - run_program(cmd, check=False, pipe=False) - wait_for_procs('dism.exe') - - # Run check health - log_path = format_log_path(log_name='DISM_CheckHealth.log', tool=True) - cmd = [ - 'DISM', '/Online', '/Cleanup-Image', - '/CheckHealth', - f'/LogPath:{log_path}', - ] - proc = run_program(cmd, check=False) - - # Check for errors - if 'no component store corruption detected' not in proc.stdout.lower(): - raise GenericError('Issue(s) detected') - - -def run_sfc_scan(): - """Run SFC and save results.""" - cmd = ['sfc', '/scannow'] - log_path = format_log_path(log_name='SFC', tool=True) - err_path = log_path.with_suffix('.err') - - # Run SFC - proc = run_program(cmd, check=False, encoding='utf-16le') - - # Save output - os.makedirs(log_path.parent, exist_ok=True) - with open(log_path, 'a') as _f: - _f.write(proc.stdout) - with open(err_path, 'a') as _f: - _f.write(proc.stderr) - - # Check result - if 'did not find any integrity violations' in proc.stdout: - pass - elif 'successfully repaired' in proc.stdout: - raise GenericWarning('Repaired') - elif 'found corrupt files' in proc.stdout: - raise GenericError('Corruption detected') - else: - raise OSError - - # Safe Mode Functions def disable_safemode(): """Edit BCD to remove safeboot value.""" diff --git a/scripts/wk/repairs/__init__.py b/scripts/wk/repairs/__init__.py new file mode 100644 index 00000000..cbdf27e4 --- /dev/null +++ b/scripts/wk/repairs/__init__.py @@ -0,0 +1,7 @@ +"""WizardKit: repairs module init""" +# vim: sts=2 sw=2 ts=2 + +import platform + +if platform.system() == 'Windows': + from wk.repairs import win diff --git a/scripts/wk/repairs/win.py b/scripts/wk/repairs/win.py new file mode 100644 index 00000000..8abf080a --- /dev/null +++ b/scripts/wk/repairs/win.py @@ -0,0 +1,277 @@ +"""WizardKit: Repairs - Windows""" +# vim: sts=2 sw=2 ts=2 + +import atexit +import logging +import os +import platform +import sys + +from wk.cfg.main import KIT_NAME_FULL +from wk.exe import get_procs, run_program, popen_program, wait_for_procs +from wk.log import format_log_path, update_log_path +from wk.os.win import * # pylint: disable=wildcard-import +from wk.std import ( + GenericError, GenericWarning, TryAndPrint, + abort, + clear_screen, print_info, + pause, sleep, + set_title, + ) + + +# STATIC VARIABLES +LOG = logging.getLogger(__name__) +AUTO_REPAIR_DELAY_IN_SECONDS = 30 +AUTO_REPAIR_KEY = fr'Software\{KIT_NAME_FULL}\AutoRepair' +CONEMU = 'ConEmuPID' in os.environ +OS_VERSION = float(platform.win32_ver()[0]) +TRY_PRINT = TryAndPrint() + + +# AutoRepair Functions +def end_session(): + """End AutoRepair session.""" + print_info('Ending repair session') + reg_delete_value('HKCU', AUTO_REPAIR_KEY, 'SessionStarted') + + # Remove logon task + cmd = [ + 'schtasks', '/delete', '/f', + '/tn', f'{KIT_NAME_FULL}-AutoRepair', + ] + run_program(cmd) + + # Disable Autologon + # TODO: Run 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 + + # Start or resume a repair session + if not session_started: + init_run() + init_session() + else: + 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() + print('') + init_run() + + # Done + return session_started + + +def init_run(): + """Initialize AutoRepair Run.""" + atexit.register(start_explorer) + # TODO: Sync Clock + kill_explorer() + # TODO: RKill + + +def init_session(): + """Initialize AutoRepair session.""" + print_info('Starting repair session') + reg_set_value('HKCU', AUTO_REPAIR_KEY, 'SessionStarted', 1, 'DWORD') + + # Create logon task for AutoRepair + cmd = [ + 'schtasks', '/create', '/f', + '/sc', 'ONLOGON', + '/tn', f'{KIT_NAME_FULL}-AutoRepair', + '/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 + # TODO: Run Autologon + # TODO: Disable Windows updates + # TODO: Reset Windows updates + 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' + 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 + + # Run repairs + # TODO: Run repairs + end_session() + + # Done + print_info('Done') + pause('Press Enter to exit...') + + +# OS Built-in Functions +def kill_explorer(): + """Kill all Explorer processes.""" + cmd = ['taskkill', '/im', 'explorer.exe', '/f'] + run_program(cmd, check=False) + + +def reboot(): + """Reboot the system.""" + atexit.unregister(start_explorer) + cmd = ['shutdown', '-r', '-t', '0'] + run_program(cmd, check=False) + raise SystemExit + + +def run_chkdsk_offline(): + """Set filesystem 'dirty bit' to force a CHKDSK during startup.""" + cmd = ['fsutil', 'dirty', 'set', os.environ.get('SYSTEMDRIVE', 'C:')] + proc = run_program(cmd, check=False) + + # Check result + if proc.returncode > 0: + raise GenericError('Failed to set dirty bit.') + + +def run_chkdsk_online(): + """Run CHKDSK. + + NOTE: If run on Windows 8+ online repairs are attempted. + """ + cmd = ['CHKDSK', os.environ.get('SYSTEMDRIVE', 'C:')] + if OS_VERSION >= 8: + cmd.extend(['/scan', '/perf']) + if CONEMU: + cmd.extend(['-new_console:n', '-new_console:s33V']) + retried = False + + # Run scan + run_program(cmd, check=False) + try: + proc = get_procs('chkdsk.exe')[0] + return_code = proc.wait() + except IndexError: + # Failed to get CHKDSK process, set return_code to force a retry + return_code = 255 + if return_code > 1: + # Try again + retried = True + run_program(cmd, check=False) + try: + proc = get_procs('chkdsk.exe')[0] + return_code = proc.wait() + except IndexError: + # Failed to get CHKDSK process + return_code = -1 + + # Check result + if return_code == -1: + raise GenericError('Failed to find CHKDSK process.') + if (return_code == 0 and retried) or return_code == 1: + raise GenericWarning('Repaired (or manually aborted)') + if return_code > 1: + raise GenericError('Issue(s) detected') + + +def run_dism(repair=True): + """Run DISM to either scan or repair component store health.""" + conemu_args = ['-new_console:n', '-new_console:s33V'] if CONEMU else [] + + # Bail early + if OS_VERSION < 8: + raise GenericWarning('Unsupported OS') + + # Run (repair) scan + log_path = format_log_path( + log_name=f'DISM_{"Restore" if repair else "Scan"}Health', tool=True, + ) + cmd = [ + 'DISM', '/Online', '/Cleanup-Image', + '/RestoreHealth' if repair else '/ScanHealth', + f'/LogPath:{log_path}', + *conemu_args, + ] + run_program(cmd, check=False, pipe=False) + wait_for_procs('dism.exe') + + # Run check health + log_path = format_log_path(log_name='DISM_CheckHealth.log', tool=True) + cmd = [ + 'DISM', '/Online', '/Cleanup-Image', + '/CheckHealth', + f'/LogPath:{log_path}', + ] + proc = run_program(cmd, check=False) + + # Check for errors + if 'no component store corruption detected' not in proc.stdout.lower(): + raise GenericError('Issue(s) detected') + + +def run_sfc_scan(): + """Run SFC and save results.""" + cmd = ['sfc', '/scannow'] + log_path = format_log_path(log_name='SFC', tool=True) + err_path = log_path.with_suffix('.err') + + # Run SFC + proc = run_program(cmd, check=False, encoding='utf-16le') + + # Save output + os.makedirs(log_path.parent, exist_ok=True) + with open(log_path, 'a') as _f: + _f.write(proc.stdout) + with open(err_path, 'a') as _f: + _f.write(proc.stderr) + + # Check result + if 'did not find any integrity violations' in proc.stdout: + pass + elif 'successfully repaired' in proc.stdout: + raise GenericWarning('Repaired') + elif 'found corrupt files' in proc.stdout: + raise GenericError('Corruption detected') + else: + raise OSError + + +def start_explorer(): + """Start Explorer.""" + popen_program(['explorer.exe']) + + +if __name__ == '__main__': + print("This file is not meant to be called directly.")