diff --git a/scripts/auto_repairs.py b/scripts/auto_repairs.py index 0f80b4f3..5bb19f17 100644 --- a/scripts/auto_repairs.py +++ b/scripts/auto_repairs.py @@ -1,16 +1,23 @@ """WizardKit: Auto Repair Tool""" # vim: sts=2 sw=2 ts=2 +from typing import Any + import wk # Classes -REBOOT_STR = wk.std.color_string('Reboot', 'YELLOW') +REBOOT_STR = wk.ui.ansi.color_string('Reboot', 'YELLOW') class MenuEntry(): """Simple class to allow cleaner code below.""" - def __init__(self, name, function=None, selected=True, **kwargs): - self.name = name - self.details = { + def __init__( + self, + name: str, + function: str | None = None, + selected: bool = True, + **kwargs): + self.name: str = name + self.details: dict[str, Any] = { 'Function': function, 'Selected': selected, **kwargs, @@ -175,8 +182,8 @@ if __name__ == '__main__': try: wk.repairs.win.run_auto_repairs(BASE_MENUS, PRESETS) except KeyboardInterrupt: - wk.std.abort() + wk.ui.cli.abort() except SystemExit: raise except: # noqa: E722 - wk.std.major_exception() + wk.ui.cli.major_exception() diff --git a/scripts/auto_setup.py b/scripts/auto_setup.py index c8f28acf..9af22ddf 100644 --- a/scripts/auto_setup.py +++ b/scripts/auto_setup.py @@ -1,15 +1,22 @@ """WizardKit: Auto System Setup Tool""" # vim: sts=2 sw=2 ts=2 +from typing import Any + import wk # Classes class MenuEntry(): """Simple class to allow cleaner code below.""" - def __init__(self, name, function=None, selected=True, **kwargs): - self.name = name - self.details = { + def __init__( + self, + name: str, + function: str | None = None, + selected: bool = True, + **kwargs): + self.name: str = name + self.details: dict[str, Any] = { 'Function': function, 'Selected': selected, **kwargs, @@ -161,8 +168,8 @@ if __name__ == '__main__': try: wk.setup.win.run_auto_setup(BASE_MENUS, PRESETS) except KeyboardInterrupt: - wk.std.abort() + wk.ui.cli.abort() except SystemExit: raise except: # noqa: E722 - wk.std.major_exception() + wk.ui.cli.major_exception() diff --git a/scripts/build-ufd b/scripts/build-ufd index 8bb04f63..3dacb7f3 100755 --- a/scripts/build-ufd +++ b/scripts/build-ufd @@ -11,4 +11,4 @@ if __name__ == '__main__': except SystemExit: raise except: # noqa: E722 - wk.std.major_exception() + wk.ui.cli.major_exception() diff --git a/scripts/build_kit_windows.py b/scripts/build_kit_windows.py index e739014b..a9323fc1 100644 --- a/scripts/build_kit_windows.py +++ b/scripts/build_kit_windows.py @@ -8,8 +8,8 @@ if __name__ == '__main__': try: wk.kit.build.build_kit() except KeyboardInterrupt: - wk.std.abort() + wk.ui.cli.abort() except SystemExit: raise except: # noqa: E722 - wk.std.major_exception() + wk.ui.cli.major_exception() diff --git a/scripts/ddrescue-tui.py b/scripts/ddrescue-tui.py index 183595a3..8e273dda 100755 --- a/scripts/ddrescue-tui.py +++ b/scripts/ddrescue-tui.py @@ -12,7 +12,7 @@ if __name__ == '__main__': docopt(wk.clone.ddrescue.DOCSTRING) except SystemExit: print('') - wk.std.pause('Press Enter to exit...') + wk.ui.cli.pause('Press Enter to exit...') raise try: @@ -20,4 +20,4 @@ if __name__ == '__main__': except SystemExit: raise except: # noqa: E722 - wk.std.major_exception() + wk.ui.cli.major_exception() diff --git a/scripts/embedded_python_env.py b/scripts/embedded_python_env.py index 7ebf7b98..1c63099b 100644 --- a/scripts/embedded_python_env.py +++ b/scripts/embedded_python_env.py @@ -8,7 +8,7 @@ python.exe -i embedded_python_env.py import wk -wk.std.print_colored( +wk.ui.cli.print_colored( (wk.cfg.main.KIT_NAME_FULL, ': ', 'Debug Console'), ('GREEN', None, 'YELLOW'), sep='', diff --git a/scripts/hw-diags.py b/scripts/hw-diags.py index 13c456a3..e2d00a41 100755 --- a/scripts/hw-diags.py +++ b/scripts/hw-diags.py @@ -12,7 +12,7 @@ if __name__ == '__main__': docopt(wk.hw.diags.DOCSTRING) except SystemExit: print('') - wk.std.pause('Press Enter to exit...') + wk.ui.cli.pause('Press Enter to exit...') raise try: @@ -20,4 +20,4 @@ if __name__ == '__main__': except SystemExit: raise except: # noqa: E722 - wk.std.major_exception() + wk.ui.cli.major_exception() diff --git a/scripts/hw-sensors b/scripts/hw-sensors index 1f1393c9..5fedcf26 100755 --- a/scripts/hw-sensors +++ b/scripts/hw-sensors @@ -7,15 +7,15 @@ import platform import wk -def main(): +def main() -> None: """Show sensor data on screen.""" sensors = wk.hw.sensors.Sensors() if platform.system() == 'Darwin': - wk.std.clear_screen() + wk.ui.cli.clear_screen() while True: print('\033[100A', end='') sensors.update_sensor_data() - wk.std.print_report(sensors.generate_report('Current', 'Max')) + wk.ui.cli.print_report(sensors.generate_report('Current', 'Max')) wk.std.sleep(1) elif platform.system() == 'Linux': proc = wk.exe.run_program(cmd=['mktemp']) @@ -43,4 +43,4 @@ if __name__ == '__main__': except SystemExit: raise except: # noqa: E722 - wk.std.major_exception() + wk.ui.cli.major_exception() diff --git a/scripts/journal-datarec-monitor b/scripts/journal-datarec-monitor new file mode 100755 index 00000000..4a8f7e0f --- /dev/null +++ b/scripts/journal-datarec-monitor @@ -0,0 +1,7 @@ +#!/bin/bash +# +## Monitor journal log for data recovery related events + +echo -e 'Monitoring journal output...\n' +journalctl -kf \ + | grep -Ei --color=always 'ata|nvme|scsi|sd[a..z]+|usb|comreset|critical|error' diff --git a/scripts/launch_sdio.py b/scripts/launch_sdio.py index 03d4070c..fbeaa4c0 100644 --- a/scripts/launch_sdio.py +++ b/scripts/launch_sdio.py @@ -1,6 +1,8 @@ """WizardKit: Launch Snappy Driver Installer Origin""" # vim: sts=2 sw=2 ts=2 +from subprocess import CompletedProcess + import wk from wk.cfg.net import SDIO_SERVER @@ -20,19 +22,19 @@ SDIO_REMOTE_PATH = wk.io.get_path_obj( ) # Functions -def try_again(): +def try_again() -> bool: """Ask to try again or quit.""" - if wk.std.ask(' Try again?'): + if wk.ui.cli.ask(' Try again?'): return True - if not wk.std.ask(' Use local version?'): - wk.std.abort() + if not wk.ui.cli.ask(' Use local version?'): + wk.ui.cli.abort() return False -def use_network_sdio(): +def use_network_sdio() -> bool: """Try to mount SDIO server.""" use_network = False - def _mount_server(): + def _mount_server() -> CompletedProcess: print('Connecting to server... (Press CTRL+c to use local copy)') return wk.net.mount_network_share(SDIO_SERVER, read_write=False) @@ -47,7 +49,7 @@ def use_network_sdio(): except KeyboardInterrupt: break except MOUNT_EXCEPTIONS as err: - wk.std.print_error(f' {err}') + wk.ui.cli.print_error(f' {err}') if not try_again(): break else: @@ -57,7 +59,7 @@ def use_network_sdio(): break # Failed to mount - wk.std.print_error(' Failed to mount server') + wk.ui.cli.print_error(' Failed to mount server') if not try_again(): break @@ -66,7 +68,7 @@ def use_network_sdio(): if __name__ == '__main__': - wk.std.set_title( + wk.ui.cli.set_title( f'{wk.cfg.main.KIT_NAME_FULL}: Snappy Driver Installer Origin Launcher', ) log_dir = wk.log.format_log_path(tool=True).parent @@ -76,7 +78,7 @@ if __name__ == '__main__': try: USE_NETWORK = use_network_sdio() except KeyboardInterrupt: - wk.std.abort() + wk.ui.cli.abort() # Run SDIO EXE_PATH = SDIO_LOCAL_PATH diff --git a/scripts/max-cpu-temp b/scripts/max-cpu-temp index 54a4ac3b..9b842538 100755 --- a/scripts/max-cpu-temp +++ b/scripts/max-cpu-temp @@ -5,10 +5,12 @@ import json import re import subprocess +from typing import Any + CPU_REGEX = re.compile(r'(core|k\d+)temp', re.IGNORECASE) NON_TEMP_REGEX = re.compile(r'^(fan|in|curr)', re.IGNORECASE) -def get_data(): +def get_data() -> dict[Any, Any]: cmd = ('sensors', '-j') data = {} raw_data = [] @@ -38,7 +40,7 @@ def get_data(): return data -def get_max_temp(data): +def get_max_temp(data) -> str: cpu_temps = [] max_cpu_temp = '??° C' for adapter, sources in data.items(): diff --git a/scripts/mount-all-volumes b/scripts/mount-all-volumes index 3adde487..becd7475 100755 --- a/scripts/mount-all-volumes +++ b/scripts/mount-all-volumes @@ -8,34 +8,34 @@ import wk # Functions -def main(): +def main() -> None: """Mount all volumes and show results.""" - wk.std.print_standard(f'{wk.cfg.main.KIT_NAME_FULL}: Volume mount tool') - wk.std.print_standard(' ') + wk.ui.cli.print_standard(f'{wk.cfg.main.KIT_NAME_FULL}: Volume mount tool') + wk.ui.cli.print_standard(' ') # Mount volumes and get report - wk.std.print_standard('Mounting volumes...') + wk.ui.cli.print_standard('Mounting volumes...') wk.os.linux.mount_volumes() report = wk.os.linux.build_volume_report() # Show results - wk.std.print_info('Results') - wk.std.print_report(report) + wk.ui.cli.print_info('Results') + wk.ui.cli.print_report(report) # GUI mode if 'gui' in sys.argv: - wk.std.pause('Press Enter to exit...') + wk.ui.cli.pause('Press Enter to exit...') wk.exe.popen_program(['nohup', 'thunar', '/media']) if __name__ == '__main__': if wk.std.PLATFORM != 'Linux': os_name = wk.std.PLATFORM.replace('Darwin', 'macOS') - wk.std.print_error(f'This script is not supported under {os_name}.') - wk.std.abort() + wk.ui.cli.print_error(f'This script is not supported under {os_name}.') + wk.ui.cli.abort() try: main() except SystemExit: raise except: # noqa: E722 - wk.std.major_exception() + wk.ui.cli.major_exception() diff --git a/scripts/mount-backup-shares b/scripts/mount-backup-shares index d19ffaf6..2b3dcb92 100755 --- a/scripts/mount-backup-shares +++ b/scripts/mount-backup-shares @@ -6,9 +6,9 @@ import wk # Functions -def main(): +def main() -> None: """Attempt to mount backup shares and print report.""" - wk.std.print_info('Mounting Backup Shares') + wk.ui.cli.print_info('Mounting Backup Shares') report = wk.net.mount_backup_shares() for line in report: color = 'GREEN' @@ -17,7 +17,7 @@ def main(): color = 'RED' elif 'Already' in line: color = 'YELLOW' - print(wk.std.color_string(line, color)) + print(wk.ansi.color_string(line, color)) if __name__ == '__main__': @@ -26,4 +26,4 @@ if __name__ == '__main__': except SystemExit: raise except: # noqa: E722 - wk.std.major_exception() + wk.ui.cli.major_exception() diff --git a/scripts/msword-search b/scripts/msword-search index a5aac0b6..14062725 100755 --- a/scripts/msword-search +++ b/scripts/msword-search @@ -46,13 +46,13 @@ def scan_file(file_path, search): if __name__ == '__main__': try: # Prep - wk.std.clear_screen() + wk.ui.cli.clear_screen() terms = [re.sub(r'\s+', r'\s*', t) for t in sys.argv[1:]] search = '({})'.format('|'.join(terms)) if len(sys.argv) == 1: # Print usage - wk.std.print_standard(USAGE) + wk.ui.cli.print_standard(USAGE) else: matches = [] for entry in scan_for_docs(SCANDIR): @@ -60,20 +60,20 @@ if __name__ == '__main__': # Strip None values (i.e. non-matching entries) matches = [m for m in matches if m] if matches: - wk.std.print_success('Found {} {}:'.format( + wk.ui.cli.print_success('Found {} {}:'.format( len(matches), 'Matches' if len(matches) > 1 else 'Match')) for match in matches: - wk.std.print_standard(match) + wk.ui.cli.print_standard(match) else: - wk.std.print_error('No matches found.') + wk.ui.cli.print_error('No matches found.') # Done - wk.std.print_standard('\nDone.') + wk.ui.cli.print_standard('\nDone.') #pause("Press Enter to exit...") except SystemExit: raise except: # noqa: E722 - wk.std.major_exception() + wk.ui.cli.major_exception() # vim: sts=2 sw=2 ts=2 diff --git a/scripts/pyproject.toml b/scripts/pyproject.toml index 8ad87a8d..e12c2d8f 100644 --- a/scripts/pyproject.toml +++ b/scripts/pyproject.toml @@ -8,6 +8,7 @@ "wk/os/__init__.py" = ["F401"] "wk/repairs/__init__.py" = ["F401"] "wk/setup/__init__.py" = ["F401"] +"wk/ui/__init__.py" = ["F401"] # Long lines "auto_setup.py" = ["E501"] diff --git a/scripts/unmount-backup-shares b/scripts/unmount-backup-shares index e7375c08..d7e3b24e 100755 --- a/scripts/unmount-backup-shares +++ b/scripts/unmount-backup-shares @@ -6,16 +6,16 @@ import wk # Functions -def main(): +def main() -> None: """Attempt to mount backup shares and print report.""" - wk.std.print_info('Unmounting Backup Shares') + wk.ui.cli.print_info('Unmounting Backup Shares') report = wk.net.unmount_backup_shares() for line in report: color = 'GREEN' line = f' {line}' if 'Not mounted' in line: color = 'YELLOW' - print(wk.std.color_string(line, color)) + print(wk.ui.ansi.color_string(line, color)) if __name__ == '__main__': @@ -24,4 +24,4 @@ if __name__ == '__main__': except SystemExit: raise except: # noqa: E722 - wk.std.major_exception() + wk.ui.cli.major_exception() diff --git a/scripts/upload-logs b/scripts/upload-logs index 98aa51a4..1d07d8a0 100755 --- a/scripts/upload-logs +++ b/scripts/upload-logs @@ -72,24 +72,24 @@ def compress_and_upload(reason='Testing'): raise wk.std.GenericError('Failed to upload logs') -def main(): +def main() -> None: """Upload logs for review.""" lines = [] - try_and_print = wk.std.TryAndPrint() + try_and_print = wk.ui.cli.TryAndPrint() # Set log wk.log.update_log_path(dest_name='Upload-Logs', timestamp=True) # Instructions - wk.std.print_success(f'{wk.cfg.main.KIT_NAME_FULL}: Upload Logs') - wk.std.print_standard('') - wk.std.print_standard('Please state the reason for the review.') - wk.std.print_info(' End note with an empty line.') - wk.std.print_standard('') + wk.ui.cli.print_success(f'{wk.cfg.main.KIT_NAME_FULL}: Upload Logs') + wk.ui.cli.print_standard('') + wk.ui.cli.print_standard('Please state the reason for the review.') + wk.ui.cli.print_info(' End note with an empty line.') + wk.ui.cli.print_standard('') # Get reason note while True: - text = wk.std.input_text('> ') + text = wk.ui.cli.input_text('> ') if not text: lines.append('') break @@ -109,5 +109,30 @@ def main(): raise SystemExit(1) +def upload_log_dir(reason='Testing') -> None: + """Upload compressed log_dir to the crash server.""" + server = wk.cfg.net.CRASH_SERVER + dest = pathlib.Path(f'~/{reason}_{NOW.strftime("%Y-%m-%dT%H%M%S%z")}.txz') + dest = dest.expanduser().resolve() + + # Compress LOG_DIR (relative to parent dir) + os.chdir(LOG_DIR.parent) + cmd = ['tar', 'caf', dest.name, LOG_DIR.name] + wk.exe.run_program(cmd, check=False) + + # Upload compressed data + url = f'{server["Url"]}/{dest.name}' + result = requests.put( + url, + data=dest.read_bytes(), + headers=server['Headers'], + auth=(server['User'], server['Pass']), + ) + + # Check result + if not result.ok: + raise wk.std.GenericError('Failed to upload logs') + + if __name__ == '__main__': main() diff --git a/scripts/wk/__init__.py b/scripts/wk/__init__.py index b99efc28..61cd542a 100644 --- a/scripts/wk/__init__.py +++ b/scripts/wk/__init__.py @@ -2,7 +2,7 @@ # vim: sts=2 sw=2 ts=2 import platform -from sys import version_info as version +from sys import stderr, version_info from . import cfg from . import clone @@ -17,7 +17,7 @@ from . import os from . import repairs from . import setup from . import std -from . import tmux +from . import ui if platform.system() != 'Windows': from wk import graph @@ -25,17 +25,18 @@ if platform.system() != 'Windows': # Check env -if version < (3, 7): +if version_info < (3, 10): # Unsupported raise RuntimeError( - f'This package is unsupported on Python {version.major}.{version.minor}' + 'This package is unsupported on Python ' + f'{version_info.major}.{version_info.minor}' ) # Init try: log.start() except UserWarning as err: - std.print_warning(err) + print(err, file=stderr) if __name__ == '__main__': diff --git a/scripts/wk/cfg/ddrescue.py b/scripts/wk/cfg/ddrescue.py index 851602fa..bee62458 100644 --- a/scripts/wk/cfg/ddrescue.py +++ b/scripts/wk/cfg/ddrescue.py @@ -1,17 +1,15 @@ """WizardKit: Config - ddrescue""" # vim: sts=2 sw=2 ts=2 -from collections import OrderedDict - # Layout TMUX_SIDE_WIDTH = 21 -TMUX_LAYOUT = OrderedDict({ +TMUX_LAYOUT = { 'Source': {'height': 2, 'Check': True}, 'Ticket': {'height': 2, 'Check': True}, 'Started': {'width': TMUX_SIDE_WIDTH, 'Check': True}, 'Progress': {'width': TMUX_SIDE_WIDTH, 'Check': True}, -}) +} # ddrescue AUTO_PASS_THRESHOLDS = { diff --git a/scripts/wk/cfg/hw.py b/scripts/wk/cfg/hw.py index 438c7e1d..6915b23b 100644 --- a/scripts/wk/cfg/hw.py +++ b/scripts/wk/cfg/hw.py @@ -3,8 +3,6 @@ import re -from collections import OrderedDict - # STATIC VARIABLES ATTRIBUTE_COLORS = ( @@ -24,7 +22,7 @@ BADBLOCKS_RESULTS_REGEX = re.compile( r'^(Checking for bad blocks .read-only test.: ).*\x08+(done|\s+).*?(\x08+)?' ) BADBLOCKS_SKIP_REGEX = re.compile(r'^(Checking|\[)', re.IGNORECASE) -CPU_CRITICAL_TEMP = 99 +CPU_CRITICAL_TEMP = 100 CPU_FAILURE_TEMP = 90 CPU_TEST_MINUTES = 7 IO_GRAPH_WIDTH = 40 @@ -173,18 +171,6 @@ THRESH_HDD_AVG_LOW = 65 * 1024**2 THRESH_SSD_MIN = 90 * 1024**2 THRESH_SSD_AVG_HIGH = 135 * 1024**2 THRESH_SSD_AVG_LOW = 100 * 1024**2 -TMUX_SIDE_WIDTH = 20 -TMUX_LAYOUT = OrderedDict({ - 'Top': {'height': 2, 'Check': True}, - 'Started': {'width': TMUX_SIDE_WIDTH, 'Check': True}, - 'Progress': {'width': TMUX_SIDE_WIDTH, 'Check': True}, - # Testing panes - 'Temps': {'height': 1000, 'Check': False}, - 'Prime95': {'height': 11, 'Check': False}, - 'SMART': {'height': 4, 'Check': True}, - 'badblocks': {'height': 5, 'Check': True}, - 'I/O Benchmark': {'height': 1000, 'Check': False}, - }) # VOLUME THRESHOLDS in percent VOLUME_WARNING_THRESHOLD = 70 VOLUME_FAILURE_THRESHOLD = 85 diff --git a/scripts/wk/cfg/ufd.py b/scripts/wk/cfg/ufd.py index 3c5e1d7a..c811d9b9 100644 --- a/scripts/wk/cfg/ufd.py +++ b/scripts/wk/cfg/ufd.py @@ -1,19 +1,17 @@ """WizardKit: Config - UFD""" # vim: sts=2 sw=2 ts=2 -from collections import OrderedDict - from wk.cfg.main import KIT_NAME_FULL # General -SOURCES = OrderedDict({ +SOURCES = { 'Linux': {'Arg': '--linux', 'Type': 'ISO'}, 'WinPE': {'Arg': '--winpe', 'Type': 'ISO'}, 'ESET SysRescue': {'Arg': '--eset', 'Type': 'IMG'}, 'Main Kit': {'Arg': '--main-kit', 'Type': 'KIT'}, 'Extra Dir': {'Arg': '--extra-dir', 'Type': 'DIR'}, - }) + } # Definitions: Boot entries BOOT_ENTRIES = { diff --git a/scripts/wk/clone/__init__.py b/scripts/wk/clone/__init__.py index 43ffabbc..f6282d3a 100644 --- a/scripts/wk/clone/__init__.py +++ b/scripts/wk/clone/__init__.py @@ -1,3 +1,4 @@ """WizardKit: ddrescue-tui module init""" from . import ddrescue +from . import menus diff --git a/scripts/wk/clone/ddrescue.py b/scripts/wk/clone/ddrescue.py index e5077d7c..aafa95c9 100644 --- a/scripts/wk/clone/ddrescue.py +++ b/scripts/wk/clone/ddrescue.py @@ -14,18 +14,19 @@ import shutil import subprocess import time -from collections import OrderedDict -from docopt import docopt +from typing import Any import psutil import pytz -from wk import cfg, debug, exe, io, log, net, osticket, std, tmux +from docopt import docopt + +from wk import cfg, debug, exe, io, log, net, osticket, std from wk.cfg.ddrescue import ( DDRESCUE_MAP_TEMPLATE, - DDRESCUE_SETTINGS, DDRESCUE_SPECIFIC_PASS_SETTINGS, ) +from wk.clone import menus from wk.hw import disk as hw_disk from wk.hw.smart import ( check_attributes, @@ -33,9 +34,11 @@ from wk.hw.smart import ( smart_status_ok, update_smart_details, ) +from wk.ui import ansi, cli, tmux, tui # STATIC VARIABLES +LOG = logging.getLogger(__name__) DOCSTRING = f'''{cfg.main.KIT_NAME_FULL}: ddrescue TUI Usage: @@ -66,9 +69,6 @@ CLONE_SETTINGS = { # (5, 1) ## Clone source partition #5 to destination partition #1 ], } -if std.PLATFORM == 'Darwin': - DDRESCUE_SETTINGS['Default']['--idirect'] = {'Selected': False, 'Hidden': True} - DDRESCUE_SETTINGS['Default']['--odirect'] = {'Selected': False, 'Hidden': True} DDRESCUE_LOG_REGEX = re.compile( r'^\s*(?P\S+):\s+' r'(?P\d+)\s+' @@ -76,6 +76,7 @@ DDRESCUE_LOG_REGEX = re.compile( r'.*\(\s*(?P\d+\.?\d*)%\)$', re.IGNORECASE, ) +DDRESCUE_OUTPUT_HEIGHT = 14 INITIAL_SKIP_MIN = 64 * 1024 # This is ddrescue's minimum accepted value REGEX_REMAINING_TIME = re.compile( r'remaining time:' @@ -86,18 +87,6 @@ REGEX_REMAINING_TIME = re.compile( r'\s*(?Pn/a)?', re.IGNORECASE ) -LOG = logging.getLogger(__name__) -MENU_ACTIONS = ( - 'Start', - 'Add tech note', - f'Change settings {std.color_string("(experts only)", "YELLOW")}', - f'Detect drives {std.color_string("(experts only)", "YELLOW")}', - f'Fresh start {std.color_string("(experts only)", "YELLOW")}', - 'Quit') -MENU_TOGGLES = { - 'Auto continue (if recovery % over threshold)': True, - 'Retry (mark non-rescued sectors "non-tried")': False, - } PANE_RATIOS = ( 12, # SMART 22, # ddrescue progress @@ -110,11 +99,6 @@ if PLATFORM == 'Darwin': RECOMMENDED_MAP_FSTYPES = re.compile( r'^(apfs|cifs|ext[234]|hfs.?|ntfs|smbfs|vfat|xfs)$' ) -SETTING_PRESETS = ( - 'Default', - 'Fast', - 'Safe', - ) STATUS_COLORS = { 'Passed': 'GREEN', 'Aborted': 'YELLOW', @@ -128,42 +112,42 @@ TIMEZONE = pytz.timezone(cfg.main.LINUX_TIME_ZONE) # Classes class BlockPair(): """Object for tracking source to dest recovery data.""" - def __init__(self, source, destination, model, working_dir): - """Initialize BlockPair() - - NOTE: source should be a wk.hw.obj.Disk() object - and destination should be a pathlib.Path() object. - """ - self.sector_size = source.phy_sec - self.source = source.path - self.destination = destination - self.map_data = {} - self.map_path = None - self.size = source.size - self.status = OrderedDict({ + def __init__( + self, + source_dev: hw_disk.Disk, + destination: pathlib.Path, + working_dir: pathlib.Path, + ): + self.sector_size: int = source_dev.phy_sec + self.source: pathlib.Path = pathlib.Path(source_dev.path) + self.destination: pathlib.Path = destination + self.map_data: dict[str, bool | int] = {} + self.map_path: pathlib.Path = pathlib.Path() + self.size: int = source_dev.size + self.status: dict[str, float | int | str] = { 'read-skip': 'Pending', 'read-full': 'Pending', 'trim': 'Pending', 'scrape': 'Pending', - }) - self.view_map = 'DISPLAY' in os.environ or 'WAYLAND_DISPLAY' in os.environ - self.view_proc = None + } + self.view_map: bool = 'DISPLAY' in os.environ or 'WAYLAND_DISPLAY' in os.environ + self.view_proc: subprocess.Popen | None = None # Set map path # e.g. '(Clone|Image)_Model[_p#]_Size[_Label].map' - map_name = model if model else 'None' - if source.bus == 'Image': + map_name = source_dev.model + if source_dev.bus == 'Image': map_name = 'Image' - if source.parent: - part_num = re.sub(r"^.*?(\d+)$", r"\1", source.path.name) + if source_dev.parent: + part_num = re.sub(r"^.*?(\d+)$", r"\1", self.source.name) map_name += f'_p{part_num}' size_str = std.bytes_to_string( size=self.size, use_binary=False, ) map_name += f'_{size_str.replace(" ", "")}' - if source.raw_details.get('label', ''): - map_name += f'_{source.raw_details["label"]}' + if source_dev.raw_details.get('label', ''): + map_name += f'_{source_dev.raw_details["label"]}' map_name = map_name.replace(' ', '_') map_name = map_name.replace('/', '_') map_name = map_name.replace('\\', '_') @@ -190,15 +174,15 @@ class BlockPair(): # Set initial status self.set_initial_status() - def get_error_size(self): + def get_error_size(self) -> int: """Get error size in bytes, returns int.""" return self.size - self.get_rescued_size() - def get_percent_recovered(self): + def get_percent_recovered(self) -> float: """Get percent rescued from map_data, returns float.""" return 100 * self.map_data.get('rescued', 0) / self.size - def get_rescued_size(self): + def get_rescued_size(self) -> int: """Get rescued size using map data. NOTE: Returns 0 if no map data is available. @@ -206,13 +190,13 @@ class BlockPair(): self.load_map_data() return self.map_data.get('rescued', 0) - def load_map_data(self): + def load_map_data(self) -> None: """Load map data from file. NOTE: If the file is missing it is assumed that recovery hasn't started yet so default values will be returned instead. """ - data = {'full recovery': False, 'pass completed': False} + data: dict[str, bool | int] = {'full recovery': False, 'pass completed': False} # Get output from ddrescuelog cmd = [ @@ -252,7 +236,7 @@ class BlockPair(): # Done self.map_data.update(data) - def pass_complete(self, pass_name): + def pass_complete(self, pass_name) -> bool: """Check if pass_name is complete based on map data, returns bool.""" pending_size = self.map_data['non-tried'] @@ -285,11 +269,11 @@ class BlockPair(): def reset_progress(self): """Reset progress to start fresh recovery.""" self.map_data = {} - self.status = OrderedDict({ + self.status = { 'read': 'Pending', 'trim': 'Pending', 'scrape': 'Pending', - }) + } self.map_path.write_text( data=DDRESCUE_MAP_TEMPLATE.format( name=cfg.main.KIT_NAME_FULL, @@ -301,7 +285,7 @@ class BlockPair(): # Set initial status self.set_initial_status() - def safety_check(self): + def safety_check(self) -> None: """Run safety check and abort if necessary.""" # TODO: Expand section to support non-Linux systems dest_size = -1 @@ -317,14 +301,14 @@ class BlockPair(): # Check destination size if cloning if not self.destination.is_file() and dest_size < self.size: - std.print_error(f'Invalid destination: {self.destination}') + cli.print_error(f'Invalid destination: {self.destination}') raise std.GenericAbort() - def set_initial_status(self): + def set_initial_status(self) -> None: """Read map data and set initial statuses.""" self.load_map_data() percent = self.get_percent_recovered() - for name in self.status.keys(): + for name in self.status: if self.pass_complete(name): self.status[name] = percent else: @@ -333,12 +317,12 @@ class BlockPair(): self.status[name] = percent break - def skip_pass(self, pass_name): + def skip_pass(self, pass_name) -> None: """Mark pass as skipped if applicable.""" if self.status[pass_name] == 'Pending': self.status[pass_name] = 'Skipped' - def update_progress(self, pass_name): + def update_progress(self, pass_name) -> None: """Update progress via map data.""" self.load_map_data() @@ -357,102 +341,36 @@ class BlockPair(): class State(): """Object for tracking hardware diagnostic data.""" def __init__(self): - self.block_pairs = [] - self.destination = None - self.log_dir = None - self.mode = None + self.block_pairs: list[BlockPair] = [] + self.destination: hw_disk.Disk | pathlib.Path = pathlib.Path('/dev/null') + self.log_dir: pathlib.Path = log.format_log_path() + self.log_dir = self.log_dir.parent.joinpath( + f'ddrescue-TUI_{time.strftime("%Y-%m-%d_%H%M%S%z")}/', + ) self.ost = osticket.osTicket() - self.panes = {} - self.source = None - self.working_dir = None + self.progress_out: pathlib.Path = self.log_dir.joinpath('progress.out') + self.mode: str = '?' + self.source: hw_disk.Disk | None = None + self.working_dir: pathlib.Path | None = None + self.ui: tui.TUI = tui.TUI('Source') - # Start a background process to maintain layout - self._init_tmux() - exe.start_thread(self._fix_tmux_layout_loop) - - def _add_block_pair(self, source, destination): + def _add_block_pair(self, source: hw_disk.Disk, destination: pathlib.Path) -> None: """Add BlockPair object and run safety checks.""" self.block_pairs.append( BlockPair( - source=source, + source_dev=source, destination=destination, - model=self.source.model, working_dir=self.working_dir, )) - def _get_clone_settings_path(self): + def _get_clone_settings_path(self) -> pathlib.Path: """get Clone settings file path, returns pathlib.Path obj.""" description = self.source.model if not description: description = self.source.path.name return pathlib.Path(f'{self.working_dir}/Clone_{description}.json') - def _fix_tmux_layout(self, forced=True): - """Fix tmux layout based on cfg.ddrescue.TMUX_LAYOUT.""" - layout = cfg.ddrescue.TMUX_LAYOUT - needs_fixed = tmux.layout_needs_fixed(self.panes, layout) - - # Main layout fix - try: - tmux.fix_layout(self.panes, layout, forced=forced) - except RuntimeError: - # Assuming self.panes changed while running - pass - - # Source/Destination - if forced or needs_fixed: - self.update_top_panes() - - # Return if Progress pane not present - if 'Progress' not in self.panes: - return - - # SMART/Journal - if forced or needs_fixed: - height = tmux.get_pane_size(self.panes['Progress'])[1] - 2 - p_ratios = [int((x/sum(PANE_RATIOS)) * height) for x in PANE_RATIOS] - if 'SMART' in self.panes: - tmux.resize_pane(self.panes['SMART'], height=p_ratios[0]) - tmux.resize_pane(height=p_ratios[1]) - if 'Journal' in self.panes: - tmux.resize_pane(self.panes['Journal'], height=p_ratios[2]) - - def _fix_tmux_layout_loop(self): - """Fix tmux layout on a loop. - - NOTE: This should be called as a thread. - """ - while True: - self._fix_tmux_layout(forced=False) - std.sleep(1) - - def _init_tmux(self): - """Initialize tmux layout.""" - tmux.kill_all_panes() - - # Source (placeholder) - self.panes['Source'] = tmux.split_window( - behind=True, - lines=2, - text=' ', - vertical=True, - ) - - # Started - self.panes['Started'] = tmux.split_window( - lines=cfg.ddrescue.TMUX_SIDE_WIDTH, - target_id=self.panes['Source'], - text=std.color_string( - ['Started', time.strftime("%Y-%m-%d %H:%M %Z")], - ['BLUE', None], - sep='\n', - ), - ) - - # Source / Dest - self.update_top_panes() - - def _load_settings(self, discard_unused_settings=False): + def _load_settings(self, discard_unused_settings: bool = False) -> dict[Any, Any]: """Load settings from previous run, returns dict.""" settings = {} settings_file = self._get_clone_settings_path() @@ -464,7 +382,7 @@ class State(): settings = json.loads(_f.read()) except (OSError, json.JSONDecodeError) as err: LOG.error('Failed to load clone settings') - std.print_error('Invalid clone settings detected.') + cli.print_error('Invalid clone settings detected.') raise std.GenericAbort() from err # Check settings @@ -476,10 +394,10 @@ class State(): bail = False for key in ('model', 'serial'): if settings['Source'][key] != getattr(self.source, key): - std.print_error(f"Clone settings don't match source {key}") + cli.print_error(f"Clone settings don't match source {key}") bail = True if settings['Destination'][key] != getattr(self.destination, key): - std.print_error(f"Clone settings don't match destination {key}") + cli.print_error(f"Clone settings don't match destination {key}") bail = True if bail: raise std.GenericAbort() @@ -501,7 +419,7 @@ class State(): # Done return settings - def _save_settings(self, settings): + def _save_settings(self, settings: dict[Any, Any]) -> None: """Save settings for future runs.""" settings_file = self._get_clone_settings_path() @@ -510,10 +428,10 @@ class State(): with open(settings_file, 'w', encoding='utf-8') as _f: json.dump(settings, _f) except OSError as err: - std.print_error('Failed to save clone settings') + cli.print_error('Failed to save clone settings') raise std.GenericAbort() from err - def add_clone_block_pairs(self): + def add_clone_block_pairs(self) -> list[hw_disk.Disk]: """Add device to device block pairs and set settings if necessary.""" source_sep = get_partition_separator(self.source.path.name) dest_sep = get_partition_separator(self.destination.path.name) @@ -535,7 +453,7 @@ class State(): ) self._add_block_pair(bp_source, bp_dest) else: - source_parts = select_disk_parts('Clone', self.source) + source_parts = menus.select_disk_parts('Clone', self.source) if self.source.path.samefile(source_parts[0].path): # Whole disk (or single partition via args), skip settings bp_dest = self.destination.path @@ -544,9 +462,9 @@ class State(): # New run, use new settings file settings['Needs Format'] = True offset = 0 - user_choice = std.choice( - ['G', 'M', 'S'], + user_choice = cli.choice( 'Format clone using GPT, MBR, or match Source type?', + ['G', 'M', 'S'], ) if user_choice == 'G': settings['Table Type'] = 'GPT' @@ -555,7 +473,7 @@ class State(): else: # Match source type settings['Table Type'] = get_table_type(self.source.path) - if std.ask('Create an empty Windows boot partition on the clone?'): + if cli.ask('Create an empty Windows boot partition on the clone?'): settings['Create Boot Partition'] = True offset = 2 if settings['Table Type'] == 'GPT' else 1 @@ -577,25 +495,28 @@ class State(): # Done return source_parts - def add_image_block_pairs(self, source_parts): + def add_image_block_pairs(self, source_parts: list[hw_disk.Disk]) -> None: """Add device to image file block pairs.""" for part in source_parts: - bp_dest = self.destination - self._add_block_pair(part, bp_dest) + self._add_block_pair(part, self.destination) - def confirm_selections(self, prompt, source_parts=None): + def confirm_selections( + self, + prompt_msg: str, + source_parts: list[hw_disk.Disk], + ) -> None: """Show selection details and prompt for confirmation.""" report = [] # Source - report.append(std.color_string('Source', 'GREEN')) + report.append(ansi.color_string('Source', 'GREEN')) report.extend(build_object_report(self.source)) report.append(' ') # Destination - report.append(std.color_string('Destination', 'GREEN')) + report.append(ansi.color_string('Destination', 'GREEN')) if self.mode == 'Clone': - report[-1] += std.color_string(' (ALL DATA WILL BE DELETED)', 'RED') + report[-1] += ansi.color_string(' (ALL DATA WILL BE DELETED)', 'RED') report.extend(build_object_report(self.destination)) report.append(' ') @@ -603,12 +524,12 @@ class State(): # NOTE: The check for block_pairs is to limit this section # to the second confirmation if self.mode == 'Clone' and self.block_pairs: - report.append(std.color_string('WARNING', 'YELLOW')) + report.append(ansi.color_string('WARNING', 'YELLOW')) report.append( 'All data will be deleted from the destination listed above.', ) report.append( - std.color_string( + ansi.color_string( ['This is irreversible and will lead to', 'DATA LOSS.'], ['YELLOW', 'RED'], ), @@ -627,18 +548,18 @@ class State(): # Map dir if self.working_dir: - report.append(std.color_string('Map Save Directory', 'GREEN')) + report.append(ansi.color_string('Map Save Directory', 'GREEN')) report.append(f'{self.working_dir}/') report.append(' ') if not fstype_is_ok(self.working_dir, map_dir=True): report.append( - std.color_string( + ansi.color_string( 'Map file(s) are being saved to a non-recommended filesystem.', 'YELLOW', ), ) report.append( - std.color_string( + ansi.color_string( ['This is strongly discouraged and may lead to', 'DATA LOSS'], [None, 'RED'], ), @@ -647,11 +568,11 @@ class State(): # Source part(s) selected if source_parts: - report.append(std.color_string('Source Part(s) selected', 'GREEN')) + report.append(ansi.color_string('Source Part(s) selected', 'GREEN')) if self.source.path.samefile(source_parts[0].path): report.append('Whole Disk') else: - report.append(std.color_string(f'{"NAME":<9} SIZE', 'BLUE')) + report.append(ansi.color_string(f'{"NAME":<9} SIZE', 'BLUE')) for part in source_parts: report.append( f'{part.path.name:<9} ' @@ -660,12 +581,12 @@ class State(): report.append(' ') # Prompt user - std.clear_screen() - std.print_report(report) - if not std.ask(prompt): + cli.clear_screen() + cli.print_report(report) + if not cli.ask(prompt_msg): raise std.GenericAbort() - def generate_report(self): + def generate_report(self) -> list[str]: """Generate report of overall and per block_pair results, returns list.""" report = [] @@ -682,7 +603,7 @@ class State(): error_size = self.get_error_size() error_size_str = std.bytes_to_string(error_size, decimals=2) if error_size > 0: - error_size_str = std.color_string(error_size_str, 'YELLOW') + error_size_str = ansi.color_string(error_size_str, 'YELLOW') percent = self.get_percent_recovered() percent = format_status_string(percent, width=0) report.append(f'Overall rescued: {percent}, error size: {error_size_str}') @@ -694,7 +615,7 @@ class State(): error_size = pair.get_error_size() error_size_str = std.bytes_to_string(error_size, decimals=2) if error_size > 0: - error_size_str = std.color_string(error_size_str, 'YELLOW') + error_size_str = ansi.color_string(error_size_str, 'YELLOW') pair_size = std.bytes_to_string(pair.size, decimals=2) percent = pair.get_percent_recovered() percent = format_status_string(percent, width=0) @@ -707,39 +628,36 @@ class State(): # Done return report - def get_error_size(self): + def get_error_size(self) -> int: """Get total error size from block_pairs in bytes, returns int.""" return self.get_total_size() - self.get_rescued_size() - def get_percent_recovered(self): + def get_percent_recovered(self) -> float: """Get total percent rescued from block_pairs, returns float.""" return 100 * self.get_rescued_size() / self.get_total_size() - def get_rescued_size(self): + def get_rescued_size(self) -> int: """Get total rescued size from all block pairs, returns int.""" return sum(pair.get_rescued_size() for pair in self.block_pairs) - def get_total_size(self): + def get_total_size(self) -> int: """Get total size of all block_pairs in bytes, returns int.""" return sum(pair.size for pair in self.block_pairs) - def init_recovery(self, docopt_args): + def init_recovery(self, docopt_args: dict[str, Any]) -> None: """Select source/dest and set env.""" - std.clear_screen() + cli.clear_screen() + disk_menu = menus.disks() source_parts = [] # Set log - self.log_dir = log.format_log_path() - self.log_dir = pathlib.Path( - f'{self.log_dir.parent}/' - f'ddrescue-TUI_{time.strftime("%Y-%m-%d_%H%M%S%z")}/' - ) log.update_log_path( dest_dir=self.log_dir, dest_name='main', keep_history=True, timestamp=False, ) + self.ui.set_progress_file(str(self.progress_out)) # Set mode self.mode = set_mode(docopt_args) @@ -751,33 +669,38 @@ class State(): ) # Select source - self.source = get_object(docopt_args['']) - if not self.source: - self.source = select_disk('Source') - self.update_top_panes() + self.source = select_disk_obj('source', disk_menu, docopt_args['']) + self.ui.set_title('Source', self.source.name) # Select destination - self.destination = get_object(docopt_args['']) - if not self.destination: - if self.mode == 'Clone': - self.destination = select_disk('Destination', self.source) - elif self.mode == 'Image': - self.destination = select_path('Destination') - self.update_top_panes() + if self.mode == 'Clone': + self.destination = select_disk_obj( + 'destination', + disk_menu, + docopt_args[''], + ) + self.ui.add_title_pane('Destination', self.destination.name) + elif self.mode == 'Image': + if docopt_args['']: + self.destination = pathlib.Path(docopt_args['']).resolve() + else: + self.destination = menus.select_path('Destination') + self.ui.add_title_pane('Destination', self.destination) # Update details self.source.update_details(skip_children=False) - self.destination.update_details(skip_children=False) + if self.mode == 'Clone': + self.destination.update_details(skip_children=False) # Confirmation #1 advanced_selection = False try: self.confirm_selections( - prompt='Are these selections correct? (use "No" for advanced selection)', + prompt_msg='Are these selections correct? (use "No" for advanced selection)', source_parts=source_parts, ) except std.GenericAbort: - if std.ask('Proceed to advanced partition selection?'): + if cli.ask('Proceed to advanced partition selection?'): advanced_selection = True else: raise @@ -803,7 +726,7 @@ class State(): if self.mode == 'Clone': source_parts = self.add_clone_block_pairs() else: - source_parts = select_disk_parts(self.mode, self.source) + source_parts = menus.select_disk_parts(self.mode, self.source) self.add_image_block_pairs(source_parts) else: if self.mode == 'Clone': @@ -815,6 +738,8 @@ class State(): # Update SMART data ## TODO: Verify if needed for dev in (self.source, self.destination): + if not isinstance(dev, hw_disk.Disk): + continue enable_smart(dev) update_smart_details(dev) @@ -826,17 +751,19 @@ class State(): # Confirmation #2 self.update_progress_pane('Idle') if advanced_selection: - self.confirm_selections('Start recovery?') + self.confirm_selections('Start recovery?', source_parts) # Unmount source and/or destination under macOS if PLATFORM == 'Darwin': - for disk in (self.source, self.destination): - cmd = ['diskutil', 'unmountDisk', disk.path] + for dev in (self.source, self.destination): + if not isinstance(dev, hw_disk.Disk): + continue + cmd = ['diskutil', 'unmountDisk', dev.path] try: exe.run_program(cmd) except subprocess.CalledProcessError: - std.print_error('Failed to unmount source and/or destination') - std.abort() + cli.print_error('Failed to unmount source and/or destination') + cli.abort() # Prep destination if self.mode == 'Clone': @@ -847,7 +774,7 @@ class State(): for pair in self.block_pairs: pair.safety_check() - def mark_started(self): + def mark_started(self) -> None: """Edit clone settings, if applicable, to mark recovery as started.""" # Skip if not cloning if self.mode != 'Clone': @@ -864,14 +791,14 @@ class State(): settings['First Run'] = False self._save_settings(settings) - def pass_above_threshold(self, pass_name): + def pass_above_threshold(self, pass_name: str) -> bool: """Check if all block_pairs meet the pass threshold, returns bool.""" threshold = cfg.ddrescue.AUTO_PASS_THRESHOLDS[pass_name] return all( p.get_percent_recovered() >= threshold for p in self.block_pairs ) - def pass_complete(self, pass_name): + def pass_complete(self, pass_name: str) -> bool: """Check if all block_pairs completed pass_name, returns bool.""" return all(p.pass_complete(pass_name) for p in self.block_pairs) @@ -902,7 +829,11 @@ class State(): # Post report self.ost.post_response('\n'.join(report), color=color) - def prep_destination(self, source_parts, dry_run=True): + def prep_destination( + self, + source_parts: list[hw_disk.Disk], + dry_run: bool = True, + ) -> None: """Prep destination as necessary.""" # TODO: Split into Linux and macOS # logical sector size is not easily found under macOS @@ -995,14 +926,14 @@ class State(): check=False, ) if proc.returncode != 0: - std.print_error('Error(s) encoundtered while formatting destination') + cli.print_error('Error(s) encoundtered while formatting destination') raise std.GenericAbort() # Update settings settings['Needs Format'] = False self._save_settings(settings) - def retry_all_passes(self): + def retry_all_passes(self) -> None: """Prep block_pairs for a retry recovery attempt.""" bad_statuses = ('*', '/', '-') LOG.warning('Updating block_pairs for retry') @@ -1030,19 +961,19 @@ class State(): # Reinitialize status pair.set_initial_status() - def safety_check_destination(self): + def safety_check_destination(self) -> None: """Run safety checks for destination and abort if necessary.""" errors_detected = False # Check for critical errors if not smart_status_ok(self.destination): - std.print_error( + cli.print_error( f'Critical error(s) detected for: {self.destination.path}', ) # Check for minor errors if not check_attributes(self.destination, only_blocking=True): - std.print_warning( + cli.print_warning( f'Attribute error(s) detected for: {self.destination.path}', ) @@ -1050,7 +981,7 @@ class State(): if errors_detected: raise std.GenericAbort() - def safety_check_size(self): + def safety_check_size(self) -> None: """Run size safety check and abort if necessary.""" required_size = sum(pair.size for pair in self.block_pairs) settings = self._load_settings() if self.mode == 'Clone' else {} @@ -1093,10 +1024,10 @@ class State(): destination_size *= 1.05 error_msg = 'Not enough free space on the destination' if required_size > destination_size: - std.print_error(error_msg) + cli.print_error(error_msg) raise std.GenericAbort() - def save_debug_reports(self): + def save_debug_reports(self) -> None: """Save debug reports to disk.""" LOG.info('Saving debug reports') debug_dir = pathlib.Path(f'{self.log_dir}/debug') @@ -1104,7 +1035,7 @@ class State(): debug_dir.mkdir() # State (self) - std.save_pickles({'state': self}, debug_dir) + debug.save_pickles({'state': self}, debug_dir) with open(f'{debug_dir}/state.report', 'a', encoding='utf-8') as _f: _f.write('[Debug report]\n') _f.write('\n'.join(debug.generate_object_report(self))) @@ -1118,23 +1049,23 @@ class State(): _f.write('\n'.join(debug.generate_object_report(_bp))) _f.write('\n') - def skip_pass(self, pass_name): + def skip_pass(self, pass_name: str) -> None: """Mark block_pairs as skipped if applicable.""" for pair in self.block_pairs: if pair.status[pass_name] == 'Pending': pair.status[pass_name] = 'Skipped' - def update_progress_pane(self, overall_status): + def update_progress_pane(self, overall_status: str) -> None: """Update progress pane.""" report = [] separator = '─────────────────────' width = cfg.ddrescue.TMUX_SIDE_WIDTH # Status - report.append(std.color_string(f'{"Status":^{width}}', 'BLUE')) + report.append(ansi.color_string(f'{"Status":^{width}}', 'BLUE')) if 'NEEDS ATTENTION' in overall_status: report.append( - std.color_string(f'{overall_status:^{width}}', 'YELLOW_BLINK'), + ansi.color_string(f'{overall_status:^{width}}', 'YELLOW_BLINK'), ) else: report.append(f'{overall_status:^{width}}') @@ -1144,12 +1075,12 @@ class State(): if self.block_pairs: total_rescued = self.get_rescued_size() percent = self.get_percent_recovered() - report.append(std.color_string('Overall Progress', 'BLUE')) + report.append(ansi.color_string('Overall Progress', 'BLUE')) report.append( f'Rescued: {format_status_string(percent, width=width-9)}', ) report.append( - std.color_string( + ansi.color_string( [f'{std.bytes_to_string(total_rescued, decimals=2):>{width}}'], [get_percent_color(percent)], ), @@ -1158,7 +1089,7 @@ class State(): # Block pair progress for pair in self.block_pairs: - report.append(std.color_string(pair.source, 'BLUE')) + report.append(ansi.color_string(pair.source, 'BLUE')) for name, status in pair.status.items(): name = name.title() report.append( @@ -1170,25 +1101,25 @@ class State(): if overall_status in ('Active', 'NEEDS ATTENTION'): etoc = get_etoc() report.append(separator) - report.append(std.color_string('Estimated Pass Finish', 'BLUE')) + report.append(ansi.color_string('Estimated Pass Finish', 'BLUE')) if overall_status == 'NEEDS ATTENTION' or etoc == 'N/A': - report.append(std.color_string('N/A', 'YELLOW')) + report.append(ansi.color_string('N/A', 'YELLOW')) else: report.append(etoc) # Write to progress file - out_path = pathlib.Path(f'{self.log_dir}/progress.out') - with open(out_path, 'w', encoding='utf-8') as _f: - _f.write('\n'.join(report)) + self.progress_out.write_text('\n'.join(report), encoding='utf-8', errors='ignore') - def update_top_panes(self): + def update_top_panes(self) -> None: """(Re)create top source/destination panes.""" source_exists = True + source_str = '' dest_exists = True + dest_str = '' width = tmux.get_pane_size()[0] width = int(width / 2) - 1 - def _format_string(obj, width): + def _format_string(obj, width) -> str: """Format source/dest string using obj and width, returns str.""" string = '' @@ -1223,37 +1154,30 @@ class State(): else: dest_exists = self.destination.exists() - # Kill destination pane - if 'Destination' in self.panes: - tmux.kill_pane(self.panes.pop('Destination')) - # Source - source_str = ' ' if self.source: source_str = _format_string(self.source, width) - tmux.respawn_pane( - self.panes['Source'], - text=std.color_string( - ['Source', '' if source_exists else ' (Missing)', '\n', source_str], - ['BLUE', 'RED', None, None], - sep='', - ), - ) # Destination - dest_str = '' if self.destination: dest_str = _format_string(self.destination, width) - self.panes['Destination'] = tmux.split_window( - percent=50, - vertical=False, - target_id=self.panes['Source'], - text=std.color_string( - ['Destination', '' if dest_exists else ' (Missing)', '\n', dest_str], - ['BLUE', 'RED', None, None], - sep='', + + # Reset title panes + self.ui.reset_title_pane( + ansi.color_string( + ['Source', '' if source_exists else ' (Missing)'], + ['BLUE', 'RED'], ), + source_str, ) + if dest_str: + self.ui.add_title_pane( + ansi.color_string( + ['Destination', '' if dest_exists else ' (Missing)'], + ['BLUE', 'RED'], + ), + dest_str, + ) # Ticket Details if self.ost and self.ost.ticket_id: @@ -1279,12 +1203,12 @@ class State(): # Functions -def build_block_pair_report(block_pairs, settings): +def build_block_pair_report(block_pairs, settings) -> list: """Build block pair report, returns list.""" report = [] notes = [] if block_pairs: - report.append(std.color_string('Block Pairs', 'GREEN')) + report.append(ansi.color_string('Block Pairs', 'GREEN')) else: # Bail early return report @@ -1303,7 +1227,7 @@ def build_block_pair_report(block_pairs, settings): if settings: if not settings['First Run']: notes.append( - std.color_string( + ansi.color_string( ['NOTE:', 'Clone settings loaded from previous run.'], ['BLUE', None], ), @@ -1311,14 +1235,14 @@ def build_block_pair_report(block_pairs, settings): if settings['Needs Format'] and settings['Table Type']: msg = f'Destination will be formatted using {settings["Table Type"]}' notes.append( - std.color_string( + ansi.color_string( ['NOTE:', msg], ['BLUE', None], ), ) if any(pair.get_rescued_size() > 0 for pair in block_pairs): notes.append( - std.color_string( + ansi.color_string( ['NOTE:', 'Resume data loaded from map file(s).'], ['BLUE', None], ), @@ -1333,7 +1257,7 @@ def build_block_pair_report(block_pairs, settings): return report -def build_ddrescue_cmd(block_pair, pass_name, settings_menu): +def build_ddrescue_cmd(block_pair, pass_name, settings_menu) -> list[str]: """Build ddrescue cmd using passed details, returns list.""" cmd = ['sudo', 'ddrescue'] if (block_pair.destination.is_block_device() @@ -1383,9 +1307,9 @@ def build_ddrescue_cmd(block_pair, pass_name, settings_menu): return cmd -def build_directory_report(path): +def build_directory_report(path: pathlib.Path) -> list[str]: """Build directory report, returns list.""" - path = f'{path}/' + path_str = f'{path}/' report = [] # Get details @@ -1393,26 +1317,26 @@ def build_directory_report(path): cmd = [ 'findmnt', '--output', 'SIZE,AVAIL,USED,FSTYPE,OPTIONS', - '--target', path, + '--target', path_str, ] proc = exe.run_program(cmd) - width = len(path) + 1 + width = len(path_str) + 1 for line in proc.stdout.splitlines(): line = line.replace('\n', '') if 'FSTYPE' in line: - line = std.color_string(f'{"PATH":<{width}}{line}', 'BLUE') + line = ansi.color_string(f'{"path_str":<{width}}{line}', 'BLUE') else: - line = f'{path:<{width}}{line}' + line = f'{path_str:<{width}}{line}' report.append(line) else: - report.append(std.color_string('PATH', 'BLUE')) - report.append(str(path)) + report.append(ansi.color_string('path_str', 'BLUE')) + report.append(str(path_str)) # Done return report -def build_disk_report(dev): +def build_disk_report(dev: hw_disk.Disk) -> list[str]: """Build device report, returns list.""" report = [] @@ -1441,7 +1365,7 @@ def build_disk_report(dev): # Partition details report.append( - std.color_string( + ansi.color_string( ( f'{"NAME":<{widths["name"]}}' f'{" " if dev.children else ""}' @@ -1485,23 +1409,7 @@ def build_disk_report(dev): return report -def build_main_menu(): - """Build main menu, returns wk.std.Menu.""" - menu = std.Menu(title=std.color_string('ddrescue TUI: Main Menu', 'GREEN')) - menu.separator = ' ' - - # Add actions, options, etc - for action in MENU_ACTIONS: - if not (PLATFORM == 'Darwin' and 'Detect drives' in action): - menu.add_action(action) - for toggle, selected in MENU_TOGGLES.items(): - menu.add_toggle(toggle, {'Selected': selected}) - - # Done - return menu - - -def build_object_report(obj): +def build_object_report(obj) -> list[str]: """Build object report, returns list.""" report = [] @@ -1517,46 +1425,7 @@ def build_object_report(obj): return report -def build_settings_menu(silent=True): - """Build settings menu, returns wk.std.Menu.""" - title_text = [ - std.color_string('ddrescue TUI: Expert Settings', 'GREEN'), - ' ', - std.color_string( - ['These settings can cause', 'MAJOR DAMAGE', 'to drives'], - ['YELLOW', 'RED', 'YELLOW'], - ), - 'Please read the manual before making changes', - ] - menu = std.Menu(title='\n'.join(title_text)) - menu.separator = ' ' - preset = 'Default' - if not silent: - # Ask which preset to use - print(f'Available ddrescue presets: {" / ".join(SETTING_PRESETS)}') - preset = std.choice(SETTING_PRESETS, 'Please select a preset:') - - # Fix selection - for _p in SETTING_PRESETS: - if _p.startswith(preset): - preset = _p - - # Add default settings - menu.add_action('Load Preset') - menu.add_action('Main Menu') - for name, details in DDRESCUE_SETTINGS['Default'].items(): - menu.add_option(name, details.copy()) - - # Update settings using preset - if preset != 'Default': - for name, details in DDRESCUE_SETTINGS[preset].items(): - menu.options[name].update(details.copy()) - - # Done - return menu - - -def build_sfdisk_partition_line(table_type, dev_path, size, details): +def build_sfdisk_partition_line(table_type, dev_path, size, details) -> str: """Build sfdisk partition line using passed details, returns str.""" line = f'{dev_path} : size={size}' dest_type = '' @@ -1583,7 +1452,7 @@ def build_sfdisk_partition_line(table_type, dev_path, size, details): # Safety Check if not dest_type: - std.print_error(f'Failed to determine partition type for: {dev_path}') + cli.print_error(f'Failed to determine partition type for: {dev_path}') raise std.GenericAbort() # Add extra details @@ -1597,7 +1466,7 @@ def build_sfdisk_partition_line(table_type, dev_path, size, details): return line -def check_destination_health(destination): +def check_destination_health(destination) -> str: """Check destination health, returns str.""" result = '' @@ -1618,7 +1487,7 @@ def check_destination_health(destination): return result -def clean_working_dir(state, confirm=False): +def clean_working_dir(state, confirm=False) -> None: """Clean working directory to ensure a fresh recovery session. NOTE: Data from previous sessions will be preserved @@ -1629,16 +1498,16 @@ def clean_working_dir(state, confirm=False): # Confirm if confirm: - std.print_error( + cli.print_error( 'This will reset all progress and create new map file(s).', ) - std.print_warning( + cli.print_warning( "Please only proceed if you understand what you're doing!", ) - std.print_warning( + cli.print_warning( 'NOTE: This will keep the current partition selection(s).', ) - if not std.ask('Continue?'): + if not cli.ask('Continue?'): return # Move settings, maps, etc to backup_dir @@ -1661,17 +1530,17 @@ def clean_working_dir(state, confirm=False): def detect_drives(state): """Detect connected drives and check source/dest selections.""" - std.clear_screen() - std.print_warning(DETECT_DRIVES_NOTICE) - if std.ask('Are you sure you proceed?'): - std.print_standard('Forcing controllers to rescan for devices...') + cli.clear_screen() + cli.print_warning(DETECT_DRIVES_NOTICE) + if cli.ask('Are you sure you proceed?'): + cli.print_standard('Forcing controllers to rescan for devices...') cmd = 'echo "- - -" | sudo tee /sys/class/scsi_host/host*/scan' exe.run_program(cmd, check=False, shell=True) if source_or_destination_changed(state): - std.abort() + cli.abort() -def format_status_string(status, width): +def format_status_string(status, width) -> str: """Format colored status string, returns str.""" color = None percent = -1 @@ -1700,13 +1569,13 @@ def format_status_string(status, width): # Add color if necessary if color: - status_str = std.color_string(status_str, color) + status_str = ansi.color_string(status_str, color) # Done return status_str -def fstype_is_ok(path, map_dir=False): +def fstype_is_ok(path, map_dir=False) -> bool: """Check if filesystem type is acceptable, returns bool.""" is_ok = False fstype = None @@ -1742,7 +1611,7 @@ def fstype_is_ok(path, map_dir=False): return is_ok -def get_ddrescue_settings(settings_menu): +def get_ddrescue_settings(settings_menu) -> list: """Get ddrescue settings from menu selections, returns list.""" settings = [] @@ -1760,7 +1629,7 @@ def get_ddrescue_settings(settings_menu): return settings -def get_etoc(): +def get_etoc() -> str: """Get EToC from ddrescue output, returns str.""" delta = None delta_dict = {} @@ -1789,7 +1658,7 @@ def get_etoc(): return etoc -def get_fstype_macos(path): +def get_fstype_macos(path) -> str: """Get fstype for path under macOS, returns str.""" fstype = 'UNKNOWN' proc = exe.run_program(['mount'], check=False) @@ -1807,41 +1676,43 @@ def get_fstype_macos(path): return fstype -def get_object(path): - """Get object based on path, returns obj.""" - obj = None +def select_disk_obj(label:str, disk_menu: cli.Menu, disk_path: str) -> hw_disk.Disk: + """Get disk based on path or menu selection, returns Disk.""" + if not disk_path: + return menus.select_disk(label.capitalize(), disk_menu) + + # Source was provided, parse and run safety checks + path = pathlib.Path(disk_path).resolve() # Bail early - if not path: - return obj + if not path.exists(): + raise FileNotFoundError(f'Path provided does not exist: {path}') - # Check path - path = pathlib.Path(path).resolve() + # Disk objects if path.is_block_device() or path.is_char_device(): obj = hw_disk.Disk(path) # Child/Parent check if obj.parent: - std.print_warning(f'"{obj.path}" is a child device') - if std.ask(f'Use parent device "{obj.parent}" instead?'): + cli.print_warning(f'"{obj.path}" is a child device') + if cli.ask(f'Use parent device "{obj.parent}" instead?'): obj = hw_disk.Disk(obj.parent) - elif path.is_dir(): - obj = path - elif path.is_file(): - # Assuming file is a raw image, mounting + + # Done + return obj + + # Raw image objects + if path.is_file(): loop_path = mount_raw_image(path) - obj = hw_disk.Disk(loop_path) + return hw_disk.Disk(loop_path) - # Abort if obj not set - if not obj: - std.print_error(f'Invalid source/dest path: {path}') - raise std.GenericAbort() - - # Done - return obj + # Abort if object type couldn't be determined + # NOTE: This shouldn't every be reached? + cli.print_error(f'Invalid {label} path: {disk_path}') + raise std.GenericAbort() -def get_partition_separator(name): +def get_partition_separator(name) -> str: """Get partition separator based on device name, returns str.""" separator = '' if re.search(r'(loop|mmc|nvme)', name, re.IGNORECASE): @@ -1850,7 +1721,7 @@ def get_partition_separator(name): return separator -def get_percent_color(percent): +def get_percent_color(percent) -> str: """Get color based on percentage, returns str.""" color = None if percent > 100: @@ -1866,7 +1737,7 @@ def get_percent_color(percent): return color -def get_table_type(disk_path): +def get_table_type(disk_path) -> str: """Get disk partition table type, returns str. NOTE: If resulting table type is not GPT or MBR @@ -1899,7 +1770,7 @@ def get_table_type(disk_path): # Check type if table_type not in ('GPT', 'MBR'): - std.print_error(f'Unsupported partition table type: {table_type}') + cli.print_error(f'Unsupported partition table type: {table_type}') raise std.GenericAbort() # Done @@ -1907,7 +1778,12 @@ def get_table_type(disk_path): def get_working_dir( - mode, destination, force_local=False, ticket_id=None, ticket_name=None): + mode, + destination, + force_local=False, + ticket_id=None, + ticket_name=None, + ) -> pathlib.Path: """Get working directory using mode and destination, returns path.""" working_dir = None @@ -1922,12 +1798,12 @@ def get_working_dir( try: path = pathlib.Path(destination).resolve() except TypeError as err: - std.print_error(f'Invalid destination: {destination}') + cli.print_error(f'Invalid destination: {destination}') raise std.GenericAbort() from err if path.exists() and fstype_is_ok(path, map_dir=False): working_dir = path elif mode == 'Clone' and not force_local: - std.print_info('Mounting backup shares...') + cli.print_info('Mounting backup shares...') net.mount_backup_shares(read_write=True) for server in cfg.net.BACKUP_SERVERS: path = pathlib.Path( @@ -1962,7 +1838,7 @@ def get_working_dir( return working_dir -def is_missing_source_or_destination(state): +def is_missing_source_or_destination(state) -> bool: """Check if source or destination dissapeared, returns bool.""" missing = False items = { @@ -1977,11 +1853,11 @@ def is_missing_source_or_destination(state): if hasattr(item, 'path'): if not item.path.exists(): missing = True - std.print_error(f'{name} disappeared') + cli.print_error(f'{name} disappeared') elif hasattr(item, 'exists'): if not item.exists(): missing = True - std.print_error(f'{name} disappeared') + cli.print_error(f'{name} disappeared') else: LOG.error('Unknown %s type: %s', name, item) @@ -1992,7 +1868,7 @@ def is_missing_source_or_destination(state): return missing -def source_or_destination_changed(state): +def source_or_destination_changed(state) -> bool: """Verify the source and destination objects are still valid.""" changed = False @@ -2013,11 +1889,11 @@ def source_or_destination_changed(state): # Done if changed: - std.print_error('Source and/or Destination changed') + cli.print_error('Source and/or Destination changed') return changed -def main(): +def main() -> None: """Main function for ddrescue TUI.""" args = docopt(DOCSTRING) log.update_log_path(dest_name='ddrescue-TUI', timestamp=True) @@ -2028,25 +1904,25 @@ def main(): raise RuntimeError('tmux session not found') # Init - atexit.register(tmux.kill_all_panes) - main_menu = build_main_menu() - settings_menu = build_settings_menu() + main_menu = menus.main() state = State() if not args['--force-local-map']: state.ost.select_ticket() if state.ost.disabled: main_menu.actions['Add tech note']['Disabled'] = True main_menu.actions['Add tech note']['Hidden'] = True - main_menu.actions[MENU_ACTIONS[2]]['Separator'] = True + # TODO: Remove this ugly call + main_menu.actions[menus.MENU_ACTIONS[2]]['Separator'] = True else: main_menu.actions['Add tech note']['Separator'] = True try: state.init_recovery(args) except (FileNotFoundError, std.GenericAbort): is_missing_source_or_destination(state) - std.abort() + cli.abort() # Show menu + settings_menu = menus.settings(state.mode) while True: selection = main_menu.advanced_select() @@ -2056,7 +1932,7 @@ def main(): selection = settings_menu.settings_select() if 'Load Preset' in selection: # Rebuild settings menu using preset - settings_menu = build_settings_menu(silent=False) + settings_menu = menus.settings(state.mode, silent=False) else: break @@ -2077,7 +1953,7 @@ def main(): # Start recovery if 'Start' in selection: - std.clear_screen() + cli.clear_screen() run_recovery(state, main_menu, settings_menu, dry_run=args['--dry-run']) # Quit @@ -2087,8 +1963,8 @@ def main(): break # Recovey < 100% - std.print_warning('Recovery is less than 100%') - if std.ask('Are you sure you want to quit?'): + cli.print_warning('Recovery is less than 100%') + if cli.ask('Are you sure you want to quit?'): break # osTicket @@ -2097,10 +1973,10 @@ def main(): # Save results to log LOG.info('') for line in state.generate_report(): - LOG.info(' %s', std.strip_colors(line)) + LOG.info(' %s', ansi.strip_colors(line)) -def mount_raw_image(path): +def mount_raw_image(path) -> pathlib.Path: """Mount raw image using OS specific methods, returns pathlib.Path.""" loopback_path = None @@ -2111,7 +1987,7 @@ def mount_raw_image(path): # Check if not loopback_path: - std.print_error(f'Failed to mount image: {path}') + cli.print_error(f'Failed to mount image: {path}') # Register unmount atexit atexit.register(unmount_loopback_device, loopback_path) @@ -2120,7 +1996,7 @@ def mount_raw_image(path): return loopback_path -def mount_raw_image_linux(path): +def mount_raw_image_linux(path) -> pathlib.Path: """Mount raw image using losetup, returns pathlib.Path.""" loopback_path = None @@ -2142,7 +2018,7 @@ def mount_raw_image_linux(path): # Done return loopback_path -def mount_raw_image_macos(path): +def mount_raw_image_macos(path) -> pathlib.Path: """Mount raw image using hdiutil, returns pathlib.Path.""" loopback_path = None plist_data = {} @@ -2173,15 +2049,16 @@ def mount_raw_image_macos(path): return loopback_path -def run_ddrescue(state, block_pair, pass_name, settings, dry_run=True): +def run_ddrescue(state, block_pair, pass_name, settings, dry_run=True) -> None: """Run ddrescue using passed settings.""" cmd = build_ddrescue_cmd(block_pair, pass_name, settings) poweroff_source_after_idle = True state.update_progress_pane('Active') - std.clear_screen() + state.ui.clear_current_pane() + state.ui.clear_on_resize = True warning_message = '' - def _poweroff_source_drive(idle_minutes): + def _poweroff_source_drive(idle_minutes) -> None: """Power off source drive after a while.""" source_dev = state.source.path @@ -2197,8 +2074,8 @@ def run_ddrescue(state, block_pair, pass_name, settings, dry_run=True): return if i % 600 == 0 and i > 0: if i == 600: - std.print_standard(' ', flush=True) - std.print_warning( + cli.print_standard(' ', flush=True) + cli.print_warning( f'Powering off source in {int((idle_minutes*60-i)/60)} minutes...', ) std.sleep(5) @@ -2208,20 +2085,20 @@ def run_ddrescue(state, block_pair, pass_name, settings, dry_run=True): cmd = ['sudo', 'hdparm', '-Y', source_dev] proc = exe.run_program(cmd, check=False) if proc.returncode: - std.print_error(f'Failed to poweroff source {source_dev}') + cli.print_error(f'Failed to poweroff source {source_dev}') else: - std.print_warning(f'Powered off source {source_dev}') - std.print_standard( + cli.print_warning(f'Powered off source {source_dev}') + cli.print_standard( 'Press Enter to return to main menu...', end='', flush=True, ) - def _update_smart_pane(): + def _update_smart_pane() -> None: """Update SMART pane every 30 seconds.""" update_smart_details(state.source) now = datetime.datetime.now(tz=TIMEZONE).strftime('%Y-%m-%d %H:%M %Z') with open(f'{state.log_dir}/smart.out', 'w', encoding='utf-8') as _f: _f.write( - std.color_string( + ansi.color_string( ['SMART Attributes', f'Updated: {now}\n'], ['BLUE', 'YELLOW'], sep='\t\t', @@ -2254,12 +2131,8 @@ def run_ddrescue(state, block_pair, pass_name, settings, dry_run=True): if warning_message: # Error detected on destination, stop recovery exe.stop_process(proc) - std.print_error(warning_message) + cli.print_error(warning_message) break - - if _i % 60 == 0: - # Clear ddrescue pane - tmux.clear_pane() _i += 1 # Update progress @@ -2289,6 +2162,7 @@ def run_ddrescue(state, block_pair, pass_name, settings, dry_run=True): # NOTE: Using 'Active' here to avoid flickering between block pairs block_pair.update_progress(pass_name) state.update_progress_pane('Active') + state.ui.clear_on_resize = False # Check result if proc.poll(): @@ -2300,19 +2174,19 @@ def run_ddrescue(state, block_pair, pass_name, settings, dry_run=True): warning_message = 'Error(s) encountered, see message above' state.update_top_panes() if warning_message: - print(' ') - print(' ') - std.print_error('DDRESCUE PROCESS HALTED') - print(' ') - std.print_warning(warning_message) + cli.print_standard(' ') + cli.print_standard(' ') + cli.print_error('DDRESCUE PROCESS HALTED') + cli.print_standard(' ') + cli.print_warning(warning_message) # Needs attention? if str(proc.poll()) != '0': state.update_progress_pane('NEEDS ATTENTION') - std.pause('Press Enter to return to main menu...') + cli.pause('Press Enter to return to main menu...') # Stop source poweroff countdown - std.print_standard('Stopping device poweroff countdown...', flush=True) + cli.print_standard('Stopping device poweroff countdown...', flush=True) poweroff_source_after_idle = False poweroff_thread.join() @@ -2320,7 +2194,7 @@ def run_ddrescue(state, block_pair, pass_name, settings, dry_run=True): raise std.GenericAbort() -def run_recovery(state, main_menu, settings_menu, dry_run=True): +def run_recovery(state: State, main_menu, settings_menu, dry_run=True) -> None: """Run recovery passes.""" atexit.register(state.save_debug_reports) attempted_recovery = False @@ -2328,12 +2202,12 @@ def run_recovery(state, main_menu, settings_menu, dry_run=True): # Bail early if is_missing_source_or_destination(state): - std.print_standard('') - std.pause('Press Enter to return to main menu...') + cli.print_standard('') + cli.pause('Press Enter to return to main menu...') return if source_or_destination_changed(state): - std.print_standard('') - std.abort() + cli.print_standard('') + cli.abort() # Get settings for name, details in main_menu.toggles.items(): @@ -2344,14 +2218,14 @@ def run_recovery(state, main_menu, settings_menu, dry_run=True): state.retry_all_passes() # Start SMART/Journal - state.panes['SMART'] = tmux.split_window( - behind=True, lines=12, vertical=True, + state.ui.add_info_pane( + percent=50, + update_layout=False, watch_file=f'{state.log_dir}/smart.out', ) - if PLATFORM != 'Darwin': - state.panes['Journal'] = tmux.split_window( - lines=4, vertical=True, cmd='journalctl --dmesg --follow', - ) + if PLATFORM == 'Linux': + state.ui.add_worker_pane(lines=4, cmd='journal-datarec-monitor') + state.ui.set_current_pane_height(DDRESCUE_OUTPUT_HEIGHT) # Run pass(es) for pass_name in ('read-skip', 'read-full', 'trim', 'scrape'): @@ -2383,15 +2257,15 @@ def run_recovery(state, main_menu, settings_menu, dry_run=True): break # Stop SMART/Journal - for pane in ('SMART', 'Journal'): - if pane in state.panes: - tmux.kill_pane(state.panes.pop(pane)) + state.ui.remove_all_info_panes() + state.ui.remove_all_worker_panes() + state.ui.clear_current_pane_height() # Show warning if nothing was done if not attempted_recovery: - std.print_warning('No actions performed') - std.print_standard(' ') - std.pause('Press Enter to return to main menu...') + cli.print_warning('No actions performed') + cli.print_standard(' ') + cli.pause('Press Enter to return to main menu...') # Done state.save_debug_reports() @@ -2399,168 +2273,9 @@ def run_recovery(state, main_menu, settings_menu, dry_run=True): state.update_progress_pane('Idle') -def select_disk(prompt, skip_disk=None): - """Select disk from list, returns Disk().""" - std.print_info('Scanning disks...') - disks = hw_disk.get_disks() - menu = std.Menu( - title=std.color_string(f'ddrescue TUI: {prompt} Selection', 'GREEN'), - ) - menu.disabled_str = 'Already selected' - menu.separator = ' ' - menu.add_action('Quit') - for disk in disks: - disable_option = False - size = disk.size - - # Check if option should be disabled - if skip_disk: - if (disk.path.samefile(skip_disk.path) - or (skip_disk.parent and disk.path.samefile(skip_disk.parent))): - disable_option = True - - # Add to menu - menu.add_option( - name=( - f'{str(disk.path):<12} ' - f'{disk.bus:<5} ' - f'{std.bytes_to_string(size, decimals=1, use_binary=False):<8} ' - f'{disk.model} ' - f'{disk.serial}' - ), - details={'Disabled': disable_option, 'Object': disk}, - ) - - # Get selection - selection = menu.simple_select() - if 'Quit' in selection: - raise std.GenericAbort() - - # Update details to include child devices - selected_disk = selection[-1]['Object'] - selected_disk.update_details(skip_children=False) - - # Done - return selected_disk - - -def select_disk_parts(prompt, disk): - """Select disk parts from list, returns list of Disk().""" - title = std.color_string('ddrescue TUI: Partition Selection', 'GREEN') - title += f'\n\nDisk: {disk.path} {disk.description}' - menu = std.Menu(title) - menu.separator = ' ' - menu.add_action('All') - menu.add_action('None') - menu.add_action('Proceed', {'Separator': True}) - menu.add_action('Quit') - object_list = [] - - def _select_parts(menu): - """Loop over selection menu until at least one partition selected.""" - while True: - selection = menu.advanced_select( - f'Please select the parts to {prompt.lower()}: ', - ) - if 'All' in selection: - for option in menu.options.values(): - option['Selected'] = True - elif 'None' in selection: - for option in menu.options.values(): - option['Selected'] = False - elif 'Proceed' in selection: - if any(option['Selected'] for option in menu.options.values()): - # At least one partition/device selected/device selected - break - elif 'Quit' in selection: - raise std.GenericAbort() - - # Bail early if running under macOS - if PLATFORM == 'Darwin': - return [disk] - - # Bail early if child device selected - if disk.parent: - return [disk] - - # Add parts - whole_disk_str = f'{str(disk.path):<14} (Whole device)' - for part in disk.children: - size = part["size"] - name = ( - f'{str(part["path"]):<14} ' - f'({std.bytes_to_string(size, decimals=1, use_binary=False):>6})' - ) - menu.add_option(name, details={'Selected': True, 'Path': part['path']}) - - # Add whole disk if necessary - if not menu.options: - menu.add_option(whole_disk_str, {'Selected': True, 'Path': disk.path}) - menu.title += '\n\n' - menu.title += std.color_string(' No partitions detected.', 'YELLOW') - - # Get selection - _select_parts(menu) - - # Build list of Disk() object_list - for option in menu.options.values(): - if option['Selected']: - object_list.append(option['Path']) - - # Check if whole disk selected - if len(object_list) == len(disk.children): - # NOTE: This is not true if the disk has no partitions - msg = f'Preserve partition table and unused space in {prompt.lower()}?' - if std.ask(msg): - # Replace part list with whole disk obj - object_list = [disk.path] - - # Convert object_list to hw_disk.Disk() objects - print(' ') - std.print_info('Getting disk/partition details...') - object_list = [hw_disk.Disk(path) for path in object_list] - - # Done - return object_list - - -def select_path(prompt): - """Select path, returns pathlib.Path.""" - invalid = False - menu = std.Menu( - title=std.color_string(f'ddrescue TUI: {prompt} Path Selection', 'GREEN'), - ) - menu.separator = ' ' - menu.add_action('Quit') - menu.add_option('Current directory') - menu.add_option('Enter manually') - path = None - - # Make selection - selection = menu.simple_select() - if 'Current directory' in selection: - path = os.getcwd() - elif 'Enter manually' in selection: - path = std.input_text('Please enter path: ') - elif 'Quit' in selection: - raise std.GenericAbort() - - # Check - try: - path = pathlib.Path(path).resolve() - except TypeError: - invalid = True - if invalid or not path.is_dir(): - std.print_error(f'Invalid path: {path}') - raise std.GenericAbort() - - # Done - return path - - -def set_mode(docopt_args): +def set_mode(docopt_args) -> str: """Set mode from docopt_args or user selection, returns str.""" - mode = None + mode = '?' # Check docopt_args if docopt_args['clone']: @@ -2570,7 +2285,7 @@ def set_mode(docopt_args): # Ask user if necessary if not mode: - answer = std.choice(['C', 'I'], 'Are we cloning or imaging?') + answer = cli.choice('Are we cloning or imaging?', ['C', 'I']) if answer == 'C': mode = 'Clone' else: @@ -2580,7 +2295,7 @@ def set_mode(docopt_args): return mode -def unmount_loopback_device(path): +def unmount_loopback_device(path) -> None: """Unmount loopback device using OS specific methods.""" cmd = [] diff --git a/scripts/wk/clone/menus.py b/scripts/wk/clone/menus.py new file mode 100644 index 00000000..b0280395 --- /dev/null +++ b/scripts/wk/clone/menus.py @@ -0,0 +1,272 @@ +"""WizardKit: ddrescue TUI - Menus""" +# vim: sts=2 sw=2 ts=2 + +import logging +import pathlib + +from wk.cfg.ddrescue import DDRESCUE_SETTINGS +from wk.hw.disk import Disk, get_disks +from wk.std import GenericAbort, PLATFORM, bytes_to_string +from wk.ui import ansi, cli + + +# STATIC VARIABLES +LOG = logging.getLogger(__name__) +CLONE_SETTINGS = { + 'Source': None, + 'Destination': None, + 'Create Boot Partition': False, + 'First Run': True, + 'Needs Format': False, + 'Table Type': None, + 'Partition Mapping': [ + # (5, 1) ## Clone source partition #5 to destination partition #1 + ], + } +if PLATFORM == 'Darwin': + # TODO: Direct I/O needs more testing under macOS + DDRESCUE_SETTINGS['Default']['--idirect'] = {'Selected': False, 'Hidden': True} + DDRESCUE_SETTINGS['Default']['--odirect'] = {'Selected': False, 'Hidden': True} +MENU_ACTIONS = ( + 'Start', + 'Add tech note', + f'Change settings {ansi.color_string("(experts only)", "YELLOW")}', + f'Detect drives {ansi.color_string("(experts only)", "YELLOW")}', + f'Fresh start {ansi.color_string("(experts only)", "YELLOW")}', + 'Quit') +MENU_TOGGLES = { + 'Auto continue (if recovery % over threshold)': True, + 'Retry (mark non-rescued sectors "non-tried")': False, + } +SETTING_PRESETS = ( + 'Default', + 'Fast', + 'Safe', + ) + + +# Functions +def main() -> cli.Menu: + """Main menu, returns wk.ui.cli.Menu.""" + menu = cli.Menu(title=ansi.color_string('ddrescue TUI: Main Menu', 'GREEN')) + menu.separator = ' ' + + # Add actions, options, etc + for action in MENU_ACTIONS: + if not (PLATFORM == 'Darwin' and 'Detect drives' in action): + menu.add_action(action) + for toggle, selected in MENU_TOGGLES.items(): + menu.add_toggle(toggle, {'Selected': selected}) + + # Done + return menu + + +def settings(mode: str, silent: bool = True) -> cli.Menu: + """Settings menu, returns wk.ui.cli.Menu.""" + title_text = [ + ansi.color_string('ddrescue TUI: Expert Settings', 'GREEN'), + ' ', + ansi.color_string( + ['These settings can cause', 'MAJOR DAMAGE', 'to drives'], + ['YELLOW', 'RED', 'YELLOW'], + ), + 'Please read the manual before making changes', + ] + menu = cli.Menu(title='\n'.join(title_text)) + menu.separator = ' ' + preset = 'Default' + if not silent: + # Ask which preset to use + cli.print_standard( + f'Available ddrescue presets: {" / ".join(SETTING_PRESETS)}' + ) + preset = cli.choice('Please select a preset:', SETTING_PRESETS) + + # Fix selection + for _p in SETTING_PRESETS: + if _p.startswith(preset): + preset = _p + + # Add default settings + menu.add_action('Load Preset') + menu.add_action('Main Menu') + for name, details in DDRESCUE_SETTINGS['Default'].items(): + menu.add_option(name, details.copy()) + + # Update settings using preset + if preset != 'Default': + for name, details in DDRESCUE_SETTINGS[preset].items(): + menu.options[name].update(details.copy()) + + # Disable direct output when saving to an image + if mode == 'Image': + menu.options['--odirect']['Disabled'] = True + menu.options['--odirect']['Selected'] = False + + # Done + return menu + + +def disks() -> cli.Menu: + """Disk menu, returns wk.ui.cli.Menu().""" + cli.print_info('Scanning disks...') + available_disks = get_disks() + menu = cli.Menu('ddrescue TUI: Disk selection') + menu.disabled_str = 'Already selected' + menu.separator = ' ' + menu.add_action('Quit') + for disk in available_disks: + menu.add_option( + name=( + f'{str(disk.path):<12} ' + f'{disk.bus:<5} ' + f'{bytes_to_string(disk.size, decimals=1, use_binary=False):<8} ' + f'{disk.model} ' + f'{disk.serial}' + ), + details={'Object': disk}, + ) + + # Done + return menu + + +def select_disk(prompt_msg: str, menu: cli.Menu) -> Disk: + """Select disk from provided Menu, returns Disk().""" + menu.title = ansi.color_string( + f'ddrescue TUI: {prompt_msg} Selection', 'GREEN', + ) + + # Get selection + selection = menu.simple_select() + if 'Quit' in selection: + raise GenericAbort() + + # Disable selected disk's menu entry + menu.options[selection[0]]['Disabled'] = True + + # Update details to include child devices + selected_disk = selection[-1]['Object'] + selected_disk.update_details(skip_children=False) + + # Done + return selected_disk + + +def select_disk_parts(prompt_msg, disk) -> list[Disk]: + """Select disk parts from list, returns list of Disk().""" + title = ansi.color_string('ddrescue TUI: Partition Selection', 'GREEN') + title += f'\n\nDisk: {disk.path} {disk.description}' + menu = cli.Menu(title) + menu.separator = ' ' + menu.add_action('All') + menu.add_action('None') + menu.add_action('Proceed', {'Separator': True}) + menu.add_action('Quit') + object_list = [] + + def _select_parts(menu) -> None: + """Loop over selection menu until at least one partition selected.""" + while True: + selection = menu.advanced_select( + f'Please select the parts to {prompt_msg.lower()}: ', + ) + if 'All' in selection: + for option in menu.options.values(): + option['Selected'] = True + elif 'None' in selection: + for option in menu.options.values(): + option['Selected'] = False + elif 'Proceed' in selection: + if any(option['Selected'] for option in menu.options.values()): + # At least one partition/device selected/device selected + break + elif 'Quit' in selection: + raise GenericAbort() + + # Bail early if running under macOS + if PLATFORM == 'Darwin': + return [disk] + + # Bail early if child device selected + if disk.parent: + return [disk] + + # Add parts + whole_disk_str = f'{str(disk.path):<14} (Whole device)' + for part in disk.children: + size = part["size"] + name = ( + f'{str(part["path"]):<14} ' + f'({bytes_to_string(size, decimals=1, use_binary=False):>6})' + ) + menu.add_option(name, details={'Selected': True, 'pathlib.Path': part['path']}) + + # Add whole disk if necessary + if not menu.options: + menu.add_option(whole_disk_str, {'Selected': True, 'pathlib.Path': disk.path}) + menu.title += '\n\n' + menu.title += ansi.color_string(' No partitions detected.', 'YELLOW') + + # Get selection + _select_parts(menu) + + # Build list of Disk() object_list + for option in menu.options.values(): + if option['Selected']: + object_list.append(option['pathlib.Path']) + + # Check if whole disk selected + if len(object_list) == len(disk.children): + # NOTE: This is not true if the disk has no partitions + msg = f'Preserve partition table and unused space in {prompt_msg.lower()}?' + if cli.ask(msg): + # Replace part list with whole disk obj + object_list = [disk.path] + + # Convert object_list to Disk() objects + cli.print_standard(' ') + cli.print_info('Getting disk/partition details...') + object_list = [Disk(path) for path in object_list] + + # Done + return object_list + + +def select_path(prompt_msg) -> pathlib.Path: + """Select path, returns pathlib.Path.""" + invalid = False + menu = cli.Menu( + title=ansi.color_string(f'ddrescue TUI: {prompt_msg} Path Selection', 'GREEN'), + ) + menu.separator = ' ' + menu.add_action('Quit') + menu.add_option('Current directory') + menu.add_option('Enter manually') + path = pathlib.Path.cwd() + + # Make selection + selection = menu.simple_select() + if 'Current directory' in selection: + pass + elif 'Enter manually' in selection: + path = pathlib.Path(cli.input_text('Please enter path: ')) + elif 'Quit' in selection: + raise GenericAbort() + + # Check + try: + path = path.resolve() + except TypeError: + invalid = True + if invalid or not path.is_dir(): + cli.print_error(f'Invalid path: {path}') + raise GenericAbort() + + # Done + return path + + +if __name__ == '__main__': + print("This file is not meant to be called directly.") diff --git a/scripts/wk/debug.py b/scripts/wk/debug.py index 11da914d..6b4e3445 100644 --- a/scripts/wk/debug.py +++ b/scripts/wk/debug.py @@ -1,21 +1,87 @@ """WizardKit: Debug Functions""" # vim: sts=2 sw=2 ts=2 +import inspect +import logging +import lzma +import os +import pathlib +import pickle +import platform +import re +import socket +import sys +import time + +from typing import Any + +import requests + +from wk.cfg.net import CRASH_SERVER +from wk.log import get_root_logger_path # Classes class Debug(): """Object used when dumping debug data.""" - def method(self): + def method(self) -> None: """Dummy method used to identify functions vs data.""" # STATIC VARIABLES +LOG = logging.getLogger(__name__) DEBUG_CLASS = Debug() METHOD_TYPE = type(DEBUG_CLASS.method) # Functions -def generate_object_report(obj, indent=0): +def generate_debug_report() -> str: + """Generate debug report, returns str.""" + platform_function_list = ( + 'architecture', + 'machine', + 'platform', + 'python_version', + ) + report = [] + + # Logging data + try: + log_path = get_root_logger_path() + except RuntimeError: + # Assuming logging wasn't started + pass + else: + report.append('------ Start Log -------') + report.append('') + with open(log_path, 'r', encoding='utf-8') as log_file: + report.extend(log_file.read().splitlines()) + report.append('') + report.append('------- End Log --------') + + # System + report.append('--- Start debug info ---') + report.append('') + report.append('[System]') + report.append(f' {"FQDN":<24} {socket.getfqdn()}') + for func in platform_function_list: + func_name = func.replace('_', ' ').capitalize() + func_result = getattr(platform, func)() + report.append(f' {func_name:<24} {func_result}') + report.append(f' {"Python sys.argv":<24} {sys.argv}') + report.append('') + + # Environment + report.append('[Environment Variables]') + for key, value in sorted(os.environ.items()): + report.append(f' {key:<24} {value}') + report.append('') + + # Done + report.append('---- End debug info ----') + return '\n'.join(report) + + +def generate_object_report(obj: Any, indent: int = 0) -> list[str]: """Generate debug report for obj, returns list.""" report = [] attr_list = [] @@ -46,5 +112,78 @@ def generate_object_report(obj, indent=0): return report +def save_pickles( + obj_dict: dict[Any, Any], + out_path: pathlib.Path | str | None = None, + ) -> None: + """Save dict of objects using pickle.""" + LOG.info('Saving pickles') + + # Set path + if not out_path: + out_path = get_root_logger_path() + out_path = out_path.parent.joinpath('../debug').resolve() + + # Save pickles + try: + for name, obj in obj_dict.copy().items(): + if name.startswith('__') or inspect.ismodule(obj): + continue + with open(f'{out_path}/{name}.pickle', 'wb') as _f: + pickle.dump(obj, _f, protocol=pickle.HIGHEST_PROTOCOL) + except Exception: + LOG.error('Failed to save all the pickles', exc_info=True) + + +def upload_debug_report( + report: str, + compress: bool = True, + reason: str = 'DEBUG', + ) -> None: + """Upload debug report to CRASH_SERVER as specified in wk.cfg.main.""" + LOG.info('Uploading debug report to %s', CRASH_SERVER.get('Name', '?')) + headers = CRASH_SERVER.get('Headers', {'X-Requested-With': 'XMLHttpRequest'}) + if compress: + headers['Content-Type'] = 'application/octet-stream' + + # Check if the required server details are available + if not all(CRASH_SERVER.get(key, False) for key in ('Name', 'Url', 'User')): + msg = 'Server details missing, aborting upload.' + print(msg) + raise RuntimeError(msg) + + # Set filename (based on the logging config if possible) + filename = 'Unknown' + try: + log_path = get_root_logger_path() + except RuntimeError: + # Assuming logging wasn't started + pass + else: + # Strip everything but the prefix + filename = re.sub(r'^(.*)_(\d{4}-\d{2}-\d{2}.*)', r'\1', log_path.name) + filename = f'{filename}_{reason}_{time.strftime("%Y-%m-%d_%H%M%S%z")}.log' + LOG.debug('filename: %s', filename) + + # Compress report + if compress: + filename += '.xz' + xz_report = lzma.compress(report.encode('utf8')) + + # Upload report + url = f'{CRASH_SERVER["Url"]}/{filename}' + response = requests.put( + url, + auth=(CRASH_SERVER['User'], CRASH_SERVER.get('Pass', '')), + data=xz_report if compress else report, + headers=headers, + timeout=60, + ) + + # Check response + if not response.ok: + raise RuntimeError('Failed to upload report') + + if __name__ == '__main__': print("This file is not meant to be called directly.") diff --git a/scripts/wk/exe.py b/scripts/wk/exe.py index b0b95039..512d8232 100644 --- a/scripts/wk/exe.py +++ b/scripts/wk/exe.py @@ -1,15 +1,18 @@ """WizardKit: Execution functions""" -#vim: sts=2 sw=2 ts=2 +# vim: sts=2 sw=2 ts=2 import json import logging import os +import pathlib import re import subprocess import time -from threading import Thread +from io import BufferedReader, TextIOWrapper from queue import Queue, Empty +from threading import Thread +from typing import Any, Callable, Iterable import psutil @@ -25,11 +28,11 @@ class NonBlockingStreamReader(): ## https://gist.github.com/EyalAr/7915597 ## https://stackoverflow.com/a/4896288 - def __init__(self, stream): - self.stream = stream - self.queue = Queue() + def __init__(self, stream: BufferedReader | TextIOWrapper): + self.stream: BufferedReader | TextIOWrapper = stream + self.queue: Queue = Queue() - def populate_queue(stream, queue): + def populate_queue(stream: BufferedReader | TextIOWrapper, queue: Queue) -> None: """Collect lines from stream and put them in queue.""" while not stream.closed: try: @@ -45,18 +48,18 @@ class NonBlockingStreamReader(): args=(self.stream, self.queue), ) - def stop(self): + def stop(self) -> None: """Stop reading from input stream.""" self.stream.close() - def read(self, timeout=None): + def read(self, timeout: float | int | None = None) -> Any: """Read from queue if possible, returns item from queue.""" try: return self.queue.get(block=timeout is not None, timeout=timeout) except Empty: return None - def save_to_file(self, proc, out_path): + def save_to_file(self, proc: subprocess.Popen, out_path: pathlib.Path | str) -> None: """Continuously save output to file while proc is running.""" LOG.debug('Saving process %s output to %s', proc, out_path) while proc.poll() is None: @@ -75,7 +78,12 @@ class NonBlockingStreamReader(): # Functions def build_cmd_kwargs( - cmd, minimized=False, pipe=True, priority=False, shell=False, **kwargs): + cmd: list[str], + minimized: bool = False, + pipe: bool = True, + priority: bool = False, + shell: bool = False, + **kwargs) -> dict[str, Any]: """Build kwargs for use by subprocess functions, returns dict. Specifically subprocess.run() and subprocess.Popen(). @@ -107,8 +115,8 @@ def build_cmd_kwargs( # Start minimized if minimized: - startupinfo = subprocess.STARTUPINFO() - startupinfo.dwFlags = subprocess.STARTF_USESHOWWINDOW + startupinfo = subprocess.STARTUPINFO() # type: ignore + startupinfo.dwFlags = subprocess.STARTF_USESHOWWINDOW # type: ignore startupinfo.wShowWindow = 6 cmd_kwargs['startupinfo'] = startupinfo @@ -126,7 +134,12 @@ def build_cmd_kwargs( return cmd_kwargs -def get_json_from_command(cmd, check=True, encoding='utf-8', errors='ignore'): +def get_json_from_command( + cmd: list[str], + check: bool = True, + encoding: str = 'utf-8', + errors: str = 'ignore', + ) -> dict[Any, Any]: """Capture JSON content from cmd output, returns dict. If the data can't be decoded then either an exception is raised @@ -145,7 +158,11 @@ def get_json_from_command(cmd, check=True, encoding='utf-8', errors='ignore'): return json_data -def get_procs(name, exact=True, try_again=True): +def get_procs( + name: str, + exact: bool = True, + try_again: bool = True, + ) -> list[psutil.Process]: """Get process object(s) based on name, returns list of proc objects.""" LOG.debug('name: %s, exact: %s', name, exact) processes = [] @@ -165,7 +182,12 @@ def get_procs(name, exact=True, try_again=True): return processes -def kill_procs(name, exact=True, force=False, timeout=30): +def kill_procs( + name: str, + exact: bool = True, + force: bool = False, + timeout: float | int = 30, + ) -> None: """Kill all processes matching name (case-insensitively). NOTE: Under Posix systems this will send SIGINT to allow processes @@ -189,7 +211,13 @@ def kill_procs(name, exact=True, force=False, timeout=30): proc.kill() -def popen_program(cmd, minimized=False, pipe=False, shell=False, **kwargs): +def popen_program( + cmd: list[str], + minimized: bool = False, + pipe: bool = False, + shell: bool = False, + **kwargs, + ) -> subprocess.Popen: """Run program and return a subprocess.Popen object.""" LOG.debug( 'cmd: %s, minimized: %s, pipe: %s, shell: %s', @@ -213,7 +241,13 @@ def popen_program(cmd, minimized=False, pipe=False, shell=False, **kwargs): return proc -def run_program(cmd, check=True, pipe=True, shell=False, **kwargs): +def run_program( + cmd: list[str], + check: bool = True, + pipe: bool = True, + shell: bool = False, + **kwargs, + ) -> subprocess.CompletedProcess: """Run program and return a subprocess.CompletedProcess object.""" LOG.debug( 'cmd: %s, check: %s, pipe: %s, shell: %s', @@ -249,7 +283,11 @@ def set_proc_priority(name, priority, exact=True): proc.nice(psutil.HIGH_PRIORITY_CLASS) -def start_thread(function, args=None, daemon=True): +def start_thread( + function: Callable, + args: Iterable[Any] | None = None, + daemon: bool = True, + ) -> Thread: """Run function as thread in background, returns Thread object.""" LOG.debug( 'Starting background thread for function: %s, args: %s, daemon: %s', @@ -261,7 +299,7 @@ def start_thread(function, args=None, daemon=True): return thread -def stop_process(proc, graceful=True): +def stop_process(proc: subprocess.Popen, graceful: bool = True) -> None: """Stop process. NOTES: proc should be a subprocess.Popen obj. @@ -283,7 +321,11 @@ def stop_process(proc, graceful=True): proc.kill() -def wait_for_procs(name, exact=True, timeout=None): +def wait_for_procs( + name: str, + exact: bool = True, + timeout: float | int | None = None, + ) -> None: """Wait for all process matching name.""" LOG.debug('name: %s, exact: %s, timeout: %s', name, exact, timeout) target_procs = get_procs(name, exact=exact) diff --git a/scripts/wk/graph.py b/scripts/wk/graph.py index 762ba35a..c856d5d5 100644 --- a/scripts/wk/graph.py +++ b/scripts/wk/graph.py @@ -12,7 +12,7 @@ import requests import Gnuplot from wk.cfg.net import BENCHMARK_SERVER, IMGUR_CLIENT_ID -from wk.std import color_string +from wk.ui import ansi # Hack to hide X11 error when running in CLI mode @@ -87,7 +87,10 @@ def export_io_graph(disk, log_dir, read_rates): return out_path -def generate_horizontal_graph(rate_list, graph_width=40, oneline=False): +def generate_horizontal_graph( + rate_list: list[float], + graph_width: int = 40, + oneline: bool = False) -> list[str]: """Generate horizontal graph from rate_list, returns list.""" graph = ['', '', '', ''] scale = 8 if oneline else 32 @@ -106,27 +109,27 @@ def generate_horizontal_graph(rate_list, graph_width=40, oneline=False): rate_color = 'GREEN' # Build graph - full_block = color_string((GRAPH_HORIZONTAL[-1],), (rate_color,)) + full_block = ansi.color_string((GRAPH_HORIZONTAL[-1],), (rate_color,)) if step >= 24: - graph[0] += color_string((GRAPH_HORIZONTAL[step-24],), (rate_color,)) + graph[0] += ansi.color_string((GRAPH_HORIZONTAL[step-24],), (rate_color,)) graph[1] += full_block graph[2] += full_block graph[3] += full_block elif step >= 16: graph[0] += ' ' - graph[1] += color_string((GRAPH_HORIZONTAL[step-16],), (rate_color,)) + graph[1] += ansi.color_string((GRAPH_HORIZONTAL[step-16],), (rate_color,)) graph[2] += full_block graph[3] += full_block elif step >= 8: graph[0] += ' ' graph[1] += ' ' - graph[2] += color_string((GRAPH_HORIZONTAL[step-8],), (rate_color,)) + graph[2] += ansi.color_string((GRAPH_HORIZONTAL[step-8],), (rate_color,)) graph[3] += full_block else: graph[0] += ' ' graph[1] += ' ' graph[2] += ' ' - graph[3] += color_string((GRAPH_HORIZONTAL[step],), (rate_color,)) + graph[3] += ansi.color_string((GRAPH_HORIZONTAL[step],), (rate_color,)) # Done if oneline: @@ -134,7 +137,7 @@ def generate_horizontal_graph(rate_list, graph_width=40, oneline=False): return graph -def get_graph_step(rate, scale=16): +def get_graph_step(rate: float, scale: int = 16) -> int: """Get graph step based on rate and scale, returns int.""" rate_in_mb = rate / (1024**2) step = 0 @@ -149,14 +152,17 @@ def get_graph_step(rate, scale=16): return step -def merge_rates(rates, graph_width=40): +def merge_rates( + rates: list[float], + graph_width: int = 40, + ) -> list[int | float]: """Merge rates to have entries equal to the width, returns list.""" merged_rates = [] offset = 0 slice_width = int(len(rates) / graph_width) # Merge rates - for _i in range(graph_width): + for _ in range(graph_width): merged_rates.append(sum(rates[offset:offset+slice_width])/slice_width) offset += slice_width @@ -234,7 +240,7 @@ def upload_to_nextcloud(image_path, ticket_number, dev_name): return BENCHMARK_SERVER['Short Url'] -def vertical_graph_line(percent, rate, scale=32): +def vertical_graph_line(percent: float, rate: float, scale: int = 32) -> str: """Build colored graph string using thresholds, returns str.""" color_bar = None color_rate = None @@ -252,7 +258,7 @@ def vertical_graph_line(percent, rate, scale=32): color_rate = 'GREEN' # Build string - line = color_string( + line = ansi.color_string( strings=( f'{percent:5.1f}%', f'{GRAPH_VERTICAL[step]:<4}', diff --git a/scripts/wk/hw/benchmark.py b/scripts/wk/hw/benchmark.py index 31e5126b..9f7c0853 100644 --- a/scripts/wk/hw/benchmark.py +++ b/scripts/wk/hw/benchmark.py @@ -21,11 +21,9 @@ from wk.cfg.hw import ( THRESH_SSD_MIN, ) from wk.exe import run_program -from wk.std import ( - PLATFORM, - strip_colors, - color_string, - ) +from wk.std import PLATFORM +from wk.ui import ansi + if platform.system() != 'Windows': from wk import graph @@ -118,7 +116,7 @@ def check_io_results(state, test_obj, rate_list, graph_width) -> None: # Add horizontal graph to report for line in graph.generate_horizontal_graph(rate_list, graph_width): - if not strip_colors(line).strip(): + if not ansi.strip_colors(line).strip(): # Skip empty lines continue test_obj.report.append(line) @@ -183,7 +181,7 @@ def run_io_test(state, test_obj, log_path, test_mode=False) -> None: LOG.info('Using %s for better performance', dev_path) offset = 0 read_rates = [] - test_obj.report.append(color_string('I/O Benchmark', 'BLUE')) + test_obj.report.append(ansi.color_string('I/O Benchmark', 'BLUE')) # Get dd values or bail try: @@ -191,7 +189,7 @@ def run_io_test(state, test_obj, log_path, test_mode=False) -> None: except DeviceTooSmallError: test_obj.set_status('N/A') test_obj.report.append( - color_string('Disk too small to test', 'YELLOW'), + ansi.color_string('Disk too small to test', 'YELLOW'), ) return diff --git a/scripts/wk/hw/cpu.py b/scripts/wk/hw/cpu.py index 8458b918..e25f75d1 100644 --- a/scripts/wk/hw/cpu.py +++ b/scripts/wk/hw/cpu.py @@ -10,13 +10,8 @@ from typing import TextIO from wk import exe from wk.cfg.hw import CPU_FAILURE_TEMP from wk.os.mac import set_fans as macos_set_fans -from wk.std import ( - PLATFORM, - color_string, - print_error, - print_warning, - ) -from wk.tmux import respawn_pane as tmux_respawn_pane +from wk.std import PLATFORM +from wk.ui import ansi # STATIC VARIABLES @@ -70,7 +65,7 @@ def check_mprime_results(test_obj, working_dir) -> None: if re.search(r'(error|fail)', line, re.IGNORECASE): warning_lines[line] = None - # print.log (check if passed) + # prime.log (check if passed) for line in _read_file('prime.log'): line = line.strip() match = re.search( @@ -97,9 +92,9 @@ def check_mprime_results(test_obj, working_dir) -> None: for line in passing_lines: test_obj.report.append(f' {line}') for line in warning_lines: - test_obj.report.append(color_string(f' {line}', 'YELLOW')) + test_obj.report.append(ansi.color_string(f' {line}', 'YELLOW')) if not (passing_lines or warning_lines): - test_obj.report.append(color_string(' Unknown result', 'YELLOW')) + test_obj.report.append(ansi.color_string(' Unknown result', 'YELLOW')) def start_mprime(working_dir, log_path) -> subprocess.Popen: @@ -116,10 +111,10 @@ def start_mprime(working_dir, log_path) -> subprocess.Popen: stdin=proc_mprime.stdout, stdout=subprocess.PIPE, ) - proc_mprime.stdout.close() - save_nsbr = exe.NonBlockingStreamReader(proc_grep.stdout) + proc_mprime.stdout.close() # type: ignore[reportOptionalMemberAccess] + save_nbsr = exe.NonBlockingStreamReader(proc_grep.stdout) exe.start_thread( - save_nsbr.save_to_file, + save_nbsr.save_to_file, args=(proc_grep, log_path), ) @@ -127,7 +122,7 @@ def start_mprime(working_dir, log_path) -> subprocess.Popen: return proc_mprime -def start_sysbench(sensors, sensors_out, log_path, pane) -> SysbenchType: +def start_sysbench(sensors, sensors_out, log_path) -> SysbenchType: """Start sysbench, returns tuple with Popen object and file handle.""" set_apple_fan_speed('max') sysbench_cmd = [ @@ -146,9 +141,6 @@ def start_sysbench(sensors, sensors_out, log_path, pane) -> SysbenchType: thermal_action=('killall', 'sysbench', '-INT'), ) - # Update bottom pane - tmux_respawn_pane(pane, watch_file=log_path, watch_cmd='tail') - # Start sysbench filehandle_sysbench = open( log_path, 'a', encoding='utf-8', @@ -174,9 +166,9 @@ def set_apple_fan_speed(speed) -> None: except (RuntimeError, ValueError, subprocess.CalledProcessError) as err: LOG.error('Failed to set fans to %s', speed) LOG.error('Error: %s', err) - print_error(f'Failed to set fans to {speed}') - for line in str(err).splitlines(): - print_warning(f' {line.strip()}') + #ui.print_error(f'Failed to set fans to {speed}') + #for line in str(err).splitlines(): + # ui.print_warning(f' {line.strip()}') elif PLATFORM == 'Linux': cmd = ['apple-fans', speed] exe.run_program(cmd, check=False) diff --git a/scripts/wk/hw/diags.py b/scripts/wk/hw/diags.py index 191f3aa5..38be4584 100644 --- a/scripts/wk/hw/diags.py +++ b/scripts/wk/hw/diags.py @@ -11,7 +11,8 @@ import time from docopt import docopt -from wk import cfg, debug, exe, log, osticket, std, tmux +from wk import cfg, debug, exe, log, osticket, std +from wk.cfg.hw import STATUS_COLORS from wk.hw import benchmark as hw_benchmark from wk.hw import cpu as hw_cpu from wk.hw import disk as hw_disk @@ -27,6 +28,8 @@ from wk.hw.network import network_test from wk.hw.screensavers import screensaver from wk.hw.test import Test, TestGroup +from wk.ui import ansi, cli, tui + # STATIC VARIABLES DOCSTRING = f'''{cfg.main.KIT_NAME_FULL}: Hardware Diagnostics @@ -99,73 +102,48 @@ class State(): """Object for tracking hardware diagnostic data.""" def __init__(self, test_mode=False): self.cpu_max_temp = -1 - self.disks = [] - self.layout = cfg.hw.TMUX_LAYOUT.copy() - self.log_dir = None + self.disks: list[hw_disk.Disk] = [] + self.log_dir: pathlib.Path | None = None self.ost = osticket.osTicket() - self.panes = {} - self.system = None - self.test_groups = [] - self.top_text = std.color_string('Hardware Diagnostics', 'GREEN') + self.progress_file: pathlib.Path | None = None + self.system: hw_system.System | None = None + self.test_groups: list[TestGroup] = [] + self.title_text: str = ansi.color_string('Hardware Diagnostics', 'GREEN') if test_mode: - self.top_text += std.color_string(' (Test Mode)', 'YELLOW') - - # Init tmux and start a background process to maintain layout - self.init_tmux() - exe.start_thread(self.fix_tmux_layout_loop) + self.title_text += ansi.color_string(' (Test Mode)', 'YELLOW') + self.ui: tui.TUI = tui.TUI(f'{self.title_text}\nMain Menu') def abort_testing(self) -> None: - """Set unfinished tests as aborted and cleanup tmux panes.""" + """Set unfinished tests as aborted and cleanup panes.""" for group in self.test_groups: for test in group.test_objects: if test.status in ('Pending', 'Working'): test.set_status('Aborted') - # Cleanup tmux - self.panes.pop('Current', None) - for key, pane_ids in self.panes.copy().items(): - if key in ('Top', 'Started', 'Progress'): - continue - if isinstance(pane_ids, str): - tmux.kill_pane(self.panes.pop(key)) - else: - for _id in pane_ids: - tmux.kill_pane(_id) - self.panes.pop(key) + # Cleanup panes + self.ui.remove_all_info_panes() + self.ui.remove_all_worker_panes() def disk_safety_checks(self) -> None: """Check for mid-run SMART failures and failed test(s).""" for dev in self.disks: disk_smart_status_check(dev, mid_run=True) for test in dev.tests: - if test.failed and 'Attributes' not in test.name: + if test.failed: + # Skip acceptable failure states + if 'Attributes' in test.name: + continue + if 'Self-Test' in test.name and 'TimedOut' in test.status: + continue + # Disable remaining tests dev.disable_disk_tests() break - def fix_tmux_layout(self, forced=True) -> None: - """Fix tmux layout based on cfg.hw.TMUX_LAYOUT.""" - try: - tmux.fix_layout(self.panes, self.layout, forced=forced) - except RuntimeError: - # Assuming self.panes changed while running - pass - - def fix_tmux_layout_loop(self) -> None: - """Fix tmux layout on a loop. - - NOTE: This should be called as a thread. - """ - while True: - self.fix_tmux_layout(forced=False) - std.sleep(1) - def init_diags(self, menu) -> None: """Initialize diagnostic pass.""" # Reset objects self.disks.clear() - self.layout.clear() - self.layout.update(cfg.hw.TMUX_LAYOUT) self.test_groups.clear() # osTicket @@ -184,15 +162,13 @@ class State(): keep_history=False, timestamp=False, ) - std.clear_screen() - std.print_info('Initializing...') + cli.clear_screen() + cli.print_info('Initializing...') # Progress Pane - self.update_progress_pane() - tmux.respawn_pane( - pane_id=self.panes['Progress'], - watch_file=f'{self.log_dir}/progress.out', - ) + self.progress_file = pathlib.Path(f'{self.log_dir}/progress.out') + self.update_progress_file() + self.ui.set_progress_file(self.progress_file) # Add HW Objects self.system = hw_system.System() @@ -244,35 +220,6 @@ class State(): test_group.test_objects.append(test_obj) self.test_groups.append(test_group) - def init_tmux(self) -> None: - """Initialize tmux layout.""" - tmux.kill_all_panes() - - # Top - self.panes['Top'] = tmux.split_window( - behind=True, - lines=2, - vertical=True, - text=f'{self.top_text}\nMain Menu', - ) - - # Started - self.panes['Started'] = tmux.split_window( - lines=cfg.hw.TMUX_SIDE_WIDTH, - target_id=self.panes['Top'], - text=std.color_string( - ['Started', time.strftime("%Y-%m-%d %H:%M %Z")], - ['BLUE', None], - sep='\n', - ), - ) - - # Progress - self.panes['Progress'] = tmux.split_window( - lines=cfg.hw.TMUX_SIDE_WIDTH, - text=' ', - ) - def save_debug_reports(self) -> None: """Save debug reports to disk.""" LOG.info('Saving debug reports') @@ -281,7 +228,7 @@ class State(): debug_dir.mkdir() # State (self) - std.save_pickles({'state': self}, debug_dir) + debug.save_pickles({'state': self}, debug_dir) with open(f'{debug_dir}/state.report', 'a', encoding='utf-8') as _f: _f.write('\n'.join(debug.generate_object_report(self))) @@ -329,28 +276,16 @@ class State(): _f.write(f'\n{test.name}:\n') _f.write('\n'.join(debug.generate_object_report(test, indent=1))) - def update_clock(self) -> None: - """Update 'Started' pane following clock sync.""" - tmux.respawn_pane( - pane_id=self.panes['Started'], - text=std.color_string( - ['Started', time.strftime("%Y-%m-%d %H:%M %Z")], - ['BLUE', None], - sep='\n', - ), - ) - - def update_progress_pane(self) -> None: - """Update progress pane.""" + def update_progress_file(self) -> None: + """Update progress file.""" report = [] - width = cfg.hw.TMUX_SIDE_WIDTH for group in self.test_groups: - report.append(std.color_string(group.name, 'BLUE')) + report.append(ansi.color_string(group.name, 'BLUE')) for test in group.test_objects: - report.append(std.color_string( - [test.label, f'{test.status:>{width-len(test.label)}}'], - [None, cfg.hw.STATUS_COLORS.get(test.status, None)], + report.append(ansi.color_string( + [test.label, f'{test.status:>{self.ui.side_width-len(test.label)}}'], + [None, STATUS_COLORS.get(test.status, None)], sep='', )) @@ -358,19 +293,17 @@ class State(): report.append(' ') # Write to progress file - out_path = pathlib.Path(f'{self.log_dir}/progress.out') - with open(out_path, 'w', encoding='utf-8') as _f: - _f.write('\n'.join(report)) + self.progress_file.write_text('\n'.join(report), encoding='utf-8') - def update_top_pane(self, text) -> None: + def update_title_text(self, text) -> None: """Update top pane with text.""" - tmux.respawn_pane(self.panes['Top'], text=f'{self.top_text}\n{text}') + self.ui.set_title(self.title_text, text) # Functions -def build_menu(cli_mode=False, quick_mode=False) -> std.Menu: - """Build main menu, returns wk.std.Menu.""" - menu = std.Menu(title=None) +def build_menu(cli_mode=False, quick_mode=False) -> cli.Menu: + """Build main menu, returns wk.ui.cli.Menu.""" + menu = cli.Menu(title=None) # Add actions, options, etc for action in MENU_ACTIONS: @@ -440,7 +373,7 @@ def cpu_stress_tests(state, test_objects, test_mode=False) -> None: return # Prep - state.update_top_pane(test_mprime_obj.dev.cpu_description) + state.update_title_text(test_mprime_obj.dev.cpu_description) test_cooling_obj.set_status('Working') test_mprime_obj.set_status('Working') @@ -452,25 +385,24 @@ def cpu_stress_tests(state, test_objects, test_mode=False) -> None: ) # Create monitor and worker panes - state.update_progress_pane() - state.panes['Prime95'] = tmux.split_window( - lines=10, vertical=True, watch_file=prime_log, watch_cmd='tail') + state.update_progress_file() + state.ui.add_worker_pane(lines=10, watch_cmd='tail', watch_file=prime_log) if PLATFORM == 'Darwin': - state.panes['Temps'] = tmux.split_window( - behind=True, percent=80, vertical=True, cmd='./hw-sensors') + state.ui.add_info_pane( + percent=80, cmd='./hw-sensors', update_layout=False, + ) elif PLATFORM == 'Linux': - state.panes['Temps'] = tmux.split_window( - behind=True, percent=80, vertical=True, watch_file=sensors_out) - tmux.resize_pane(height=3) - state.panes['Current'] = '' - state.layout['Current'] = {'height': 3, 'Check': True} + state.ui.add_info_pane( + percent=80, watch_file=sensors_out, update_layout=False, + ) + state.ui.set_current_pane_height(3) # Get idle temps - std.print_standard('Saving idle temps...') + cli.print_standard('Saving idle temps...') sensors.save_average_temps(temp_label='Idle', seconds=5) # Stress CPU - std.print_info('Running stress test') + cli.print_info('Running stress test') hw_cpu.set_apple_fan_speed('max') proc_mprime = hw_cpu.start_mprime(state.log_dir, prime_log) @@ -488,17 +420,17 @@ def cpu_stress_tests(state, test_objects, test_mode=False) -> None: if sensors.cpu_reached_critical_temp() or aborted: test_cooling_obj.set_status('Aborted') test_mprime_obj.set_status('Aborted') - state.update_progress_pane() + state.update_progress_file() # Get cooldown temp - std.clear_screen() - std.print_standard('Letting CPU cooldown...') + state.ui.clear_current_pane() + cli.print_standard('Letting CPU cooldown...') std.sleep(5) - std.print_standard('Saving cooldown temps...') + cli.print_standard('Saving cooldown temps...') sensors.save_average_temps(temp_label='Cooldown', seconds=5) # Check Prime95 results - test_mprime_obj.report.append(std.color_string('Prime95', 'BLUE')) + test_mprime_obj.report.append(ansi.color_string('Prime95', 'BLUE')) hw_cpu.check_mprime_results( test_obj=test_mprime_obj, working_dir=state.log_dir, ) @@ -509,16 +441,19 @@ def cpu_stress_tests(state, test_objects, test_mode=False) -> None: ) if run_sysbench: LOG.info('CPU Test (Sysbench)') - std.print_standard('Letting CPU cooldown more...') - std.sleep(30) - std.clear_screen() - std.print_info('Running alternate stress test') + cli.print_standard('Letting CPU cooldown more...') + std.sleep(10) + state.ui.clear_current_pane() + cli.print_info('Running alternate stress test') print('') + sysbench_log = prime_log.with_name('sysbench.log') + sysbench_log.touch() + state.ui.remove_all_worker_panes() + state.ui.add_worker_pane(lines=10, watch_cmd='tail', watch_file=sysbench_log) proc_sysbench, filehandle_sysbench = hw_cpu.start_sysbench( sensors, sensors_out, - log_path=prime_log.with_name('sysbench.log'), - pane=state.panes['Prime95'], + log_path=sysbench_log, ) try: print_countdown(proc=proc_sysbench, seconds=test_minutes*60) @@ -535,10 +470,10 @@ def cpu_stress_tests(state, test_objects, test_mode=False) -> None: if sensors.cpu_reached_critical_temp() or aborted: test_cooling_obj.set_status('Aborted') test_mprime_obj.set_status('Aborted') - state.update_progress_pane() + state.update_progress_file() # Check Cooling results - test_cooling_obj.report.append(std.color_string('Temps', 'BLUE')) + test_cooling_obj.report.append(ansi.color_string('Temps', 'BLUE')) hw_cpu.check_cooling_results(test_cooling_obj, sensors, run_sysbench) # Post results to osTicket @@ -552,11 +487,11 @@ def cpu_stress_tests(state, test_objects, test_mode=False) -> None: ) # Cleanup - state.update_progress_pane() + state.update_progress_file() sensors.stop_background_monitor() - state.panes.pop('Current', None) - tmux.kill_pane(state.panes.pop('Prime95', None)) - tmux.kill_pane(state.panes.pop('Temps', None)) + state.ui.clear_current_pane_height() + state.ui.remove_all_info_panes() + state.ui.remove_all_worker_panes() # Done if aborted: @@ -574,7 +509,7 @@ def disk_attribute_check(state, test_objects, test_mode=False) -> None: continue # Done - state.update_progress_pane() + state.update_progress_file() def disk_io_benchmark( @@ -584,16 +519,10 @@ def disk_io_benchmark( aborted = False # Run benchmarks - state.update_top_pane( + state.update_title_text( f'Disk I/O Benchmark{"s" if len(test_objects) > 1 else ""}', ) - state.panes['I/O Benchmark'] = tmux.split_window( - percent=50, - vertical=True, - text=' ', - ) - - # Skip USB devices if requested + state.ui.set_current_pane_height(10) for test in test_objects: if ( skip_usb @@ -610,16 +539,18 @@ def disk_io_benchmark( continue # Start benchmark - std.clear_screen() - std.print_report(test.dev.generate_report()) + state.ui.clear_current_pane() + cli.print_report(test.dev.generate_report()) test.set_status('Working') test_log = f'{state.log_dir}/{test.dev.path.name}_benchmark.out' - tmux.respawn_pane( - state.panes['I/O Benchmark'], + state.ui.remove_all_worker_panes() + state.ui.add_worker_pane( + percent=50, + update_layout=False, watch_cmd='tail', watch_file=test_log, ) - state.update_progress_pane() + state.update_progress_file() try: hw_benchmark.run_io_test(state, test, test_log, test_mode=test_mode) except KeyboardInterrupt: @@ -628,20 +559,21 @@ def disk_io_benchmark( # Something went wrong LOG.error('%s', err) test.set_status('ERROR') - test.report.append(std.color_string(' Unknown Error', 'RED')) + test.report.append(ansi.color_string(' Unknown Error', 'RED')) # Mark test(s) aborted if necessary if aborted: test.set_status('Aborted') - test.report.append(std.color_string(' Aborted', 'YELLOW')) + test.report.append(ansi.color_string(' Aborted', 'YELLOW')) break # Update progress after each test - state.update_progress_pane() + state.update_progress_file() # Cleanup - state.update_progress_pane() - tmux.kill_pane(state.panes.pop('I/O Benchmark', None)) + state.update_progress_file() + state.ui.clear_current_pane_height() + state.ui.remove_all_worker_panes() # Done if aborted: @@ -653,13 +585,12 @@ def disk_self_test(state, test_objects, test_mode=False) -> None: LOG.info('Disk Self-Test(s)') aborted = False threads = [] - state.panes['SMART'] = [] # Run self-tests - state.update_top_pane( + state.update_title_text( f'Disk self-test{"s" if len(test_objects) > 1 else ""}', ) - std.print_info(f'Starting self-test{"s" if len(test_objects) > 1 else ""}') + cli.print_info(f'Starting self-test{"s" if len(test_objects) > 1 else ""}') show_failed_attributes(state) for test in reversed(test_objects): if test.disabled: @@ -673,12 +604,10 @@ def disk_self_test(state, test_objects, test_mode=False) -> None: # Show progress if threads[-1].is_alive(): - state.panes['SMART'].append( - tmux.split_window(lines=4, vertical=True, watch_file=test_log), - ) + state.ui.add_worker_pane(lines=4, watch_file=test_log) # Wait for all tests to complete - state.update_progress_pane() + state.update_progress_file() try: while True: if any(t.is_alive() for t in threads): @@ -693,10 +622,8 @@ def disk_self_test(state, test_objects, test_mode=False) -> None: hw_smart.build_self_test_report(test, aborted=True) # Cleanup - state.update_progress_pane() - for pane in state.panes['SMART']: - tmux.kill_pane(pane) - state.panes.pop('SMART', None) + state.update_progress_file() + state.ui.remove_all_worker_panes() # Done if aborted: @@ -750,13 +677,12 @@ def disk_surface_scan(state, test_objects, test_mode=False) -> None: LOG.info('Disk Surface Scan (badblocks)') aborted = False threads = [] - state.panes['badblocks'] = [] # Update panes - state.update_top_pane( + state.update_title_text( f'Disk Surface Scan{"s" if len(test_objects) > 1 else ""}', ) - std.print_info( + cli.print_info( f'Starting disk surface scan{"s" if len(test_objects) > 1 else ""}', ) show_failed_attributes(state) @@ -772,20 +698,13 @@ def disk_surface_scan(state, test_objects, test_mode=False) -> None: # Show progress if threads[-1].is_alive(): - state.panes['badblocks'].append( - tmux.split_window( - lines=5, - vertical=True, - watch_cmd='tail', - watch_file=test_log, - ), - ) + state.ui.add_worker_pane(lines=5, watch_cmd='tail', watch_file=test_log) # Wait for all tests to complete try: while True: if any(t.is_alive() for t in threads): - state.update_progress_pane() + state.update_progress_file() std.sleep(5) else: break @@ -796,13 +715,11 @@ def disk_surface_scan(state, test_objects, test_mode=False) -> None: for test in test_objects: if not (test.disabled or test.passed or test.failed): test.set_status('Aborted') - test.report.append(std.color_string(' Aborted', 'YELLOW')) + test.report.append(ansi.color_string(' Aborted', 'YELLOW')) # Cleanup - state.update_progress_pane() - for pane in state.panes['badblocks']: - tmux.kill_pane(pane) - state.panes.pop('badblocks', None) + state.update_progress_file() + state.ui.remove_all_worker_panes() # Done if aborted: @@ -816,7 +733,7 @@ def disk_volume_utilization(state, test_objects, test_mode=False) -> None: hw_volumes.check_volume_utilization(test) # Done - state.update_progress_pane() + state.update_progress_file() def main() -> None: @@ -830,7 +747,6 @@ def main() -> None: raise RuntimeError('tmux session not found') # Init - atexit.register(tmux.kill_all_panes) menu = build_menu(cli_mode=args['--cli'], quick_mode=args['--quick']) state = State(test_mode=args['--test-mode']) state.override_all_smart_errors = args['--ignore-smart-errors'] @@ -858,15 +774,15 @@ def main() -> None: # Run simple test if action: - state.update_top_pane(selection[0]) + state.update_title_text(selection[0]) try: action() except KeyboardInterrupt: - std.print_warning('Aborted.') - std.print_standard('') - std.pause('Press Enter to return to main menu...') + cli.print_warning('Aborted.') + cli.print_standard('') + cli.pause('Press Enter to return to main menu...') if 'Clock Sync' in selection: - state.update_clock() + state.ui.update_clock() # Secrets if 'Matrix' in selection: @@ -890,7 +806,7 @@ def main() -> None: run_diags(state, menu, quick_mode=False, test_mode=args['--test-mode']) # Reset top pane - state.update_top_pane('Main Menu') + state.update_title_text('Main Menu') def post_system_info(state, quick_mode=False, test_mode=False) -> None: @@ -946,8 +862,8 @@ def run_diags(state, menu, quick_mode=False, test_mode=False) -> None: # Just return if no tests were selected if not state.test_groups: - std.print_warning('No tests selected?') - std.pause() + cli.print_warning('No tests selected?') + cli.pause() return # osTicket @@ -975,13 +891,13 @@ def run_diags(state, menu, quick_mode=False, test_mode=False) -> None: args = [group.test_objects] if group.name == 'Disk I/O Benchmark': args.append(menu.toggles[IO_SIZE_SKIP_NAME]['Selected']) - std.clear_screen() + state.ui.clear_current_pane() try: function(state, *args, test_mode=test_mode) except (KeyboardInterrupt, std.GenericAbort): aborted = True state.abort_testing() - state.update_progress_pane() + state.update_progress_file() break else: # Run safety checks after disk tests @@ -1008,48 +924,48 @@ def run_diags(state, menu, quick_mode=False, test_mode=False) -> None: state.save_debug_reports() atexit.unregister(state.save_debug_reports) if quick_mode: - std.pause('Press Enter to exit...') + cli.pause('Press Enter to exit...') else: - std.pause('Press Enter to return to main menu...') + cli.pause('Press Enter to return to main menu...') def show_failed_attributes(state) -> None: """Show failed attributes for all disks.""" for dev in state.disks: - std.print_colored([dev.name, dev.description], ['CYAN', None]) - std.print_report( + cli.print_colored([dev.name, dev.description], ['CYAN', None]) + cli.print_report( hw_smart.generate_attribute_report(dev, only_failed=True), ) - std.print_standard('') + cli.print_standard('') def show_results(state) -> None: """Show test results by device.""" std.sleep(0.5) - std.clear_screen() - state.update_top_pane('Results') + state.ui.clear_current_pane() + state.update_title_text('Results') # CPU Tests cpu_tests_enabled = [ group.name for group in state.test_groups if 'CPU' in group.name ] if cpu_tests_enabled: - std.print_success('CPU:') - std.print_report(state.system.generate_cpu_ram_report()) - std.print_standard(' ') + cli.print_success('CPU:') + cli.print_report(state.system.generate_cpu_ram_report()) + cli.print_standard(' ') # Disk Tests disk_tests_enabled = [ group.name for group in state.test_groups if 'Disk' in group.name ] if disk_tests_enabled: - std.print_success(f'Disk{"s" if len(state.disks) > 1 else ""}:') + cli.print_success(f'Disk{"s" if len(state.disks) > 1 else ""}:') for disk in state.disks: - std.print_report(disk.generate_report()) - std.print_standard(' ') + cli.print_report(disk.generate_report()) + cli.print_standard(' ') if not state.disks: - std.print_warning('No devices') - std.print_standard(' ') + cli.print_warning('No devices') + cli.print_standard(' ') def sync_clock() -> None: diff --git a/scripts/wk/hw/disk.py b/scripts/wk/hw/disk.py index 3c7c727c..7b06bbac 100644 --- a/scripts/wk/hw/disk.py +++ b/scripts/wk/hw/disk.py @@ -9,7 +9,7 @@ import plistlib import re from dataclasses import dataclass, field -from typing import Any, Union +from typing import Any from wk.cfg.main import KIT_NAME_SHORT from wk.cfg.python import DATACLASS_DECORATOR_KWARGS @@ -19,7 +19,8 @@ from wk.hw.smart import ( generate_attribute_report, get_known_disk_attributes, ) -from wk.std import PLATFORM, color_string, strip_colors +from wk.std import PLATFORM +from wk.ui import ansi # STATIC VARIABLES @@ -45,19 +46,20 @@ class Disk: model: str = field(init=False) name: str = field(init=False) notes: list[str] = field(init=False, default_factory=list) - path: Union[pathlib.Path, str] + path: pathlib.Path = field(init=False) + path_str: pathlib.Path | str parent: str = field(init=False) phy_sec: int = field(init=False) raw_details: dict[str, Any] = field(init=False) - raw_smartctl: dict[str, Any] = field(init=False) + raw_smartctl: dict[str, Any] = field(init=False, default_factory=dict) serial: str = field(init=False) size: int = field(init=False) ssd: bool = field(init=False) tests: list[Test] = field(init=False, default_factory=list) use_sat: bool = field(init=False, default=False) - def __post_init__(self) -> None: - self.path = pathlib.Path(self.path).resolve() + def __post_init__(self): + self.path = pathlib.Path(self.path_str).resolve() self.update_details() self.set_description() self.known_attributes = get_known_disk_attributes(self.model) @@ -73,7 +75,7 @@ class Disk: def add_note(self, note, color=None) -> None: """Add note that will be included in the disk report.""" if color: - note = color_string(note, color) + note = ansi.color_string(note, color) if note not in self.notes: self.notes.append(note) self.notes.sort() @@ -82,7 +84,7 @@ class Disk: """Check if note is already present.""" present = False for note in self.notes: - if note_str == strip_colors(note): + if note_str == ansi.strip_colors(note): present = True return present @@ -98,18 +100,18 @@ class Disk: """Generate Disk report, returns list.""" report = [] if header: - report.append(color_string(f'Device ({self.path.name})', 'BLUE')) + report.append(ansi.color_string(f'Device ({self.path.name})', 'BLUE')) report.append(f' {self.description}') # Attributes if self.attributes: if header: - report.append(color_string('Attributes', 'BLUE')) + report.append(ansi.color_string('Attributes', 'BLUE')) report.extend(generate_attribute_report(self)) # Notes if self.notes: - report.append(color_string('Notes', 'BLUE')) + report.append(ansi.color_string('Notes', 'BLUE')) for note in self.notes: report.append(f' {note}') diff --git a/scripts/wk/hw/keyboard.py b/scripts/wk/hw/keyboard.py index 68e2d0a6..07d3a22d 100644 --- a/scripts/wk/hw/keyboard.py +++ b/scripts/wk/hw/keyboard.py @@ -4,7 +4,7 @@ import logging from wk.exe import run_program -from wk.std import PLATFORM, print_warning +from wk.std import PLATFORM # STATIC VARIABLES @@ -17,7 +17,8 @@ def keyboard_test() -> None: if PLATFORM == 'Linux': run_xev() else: - print_warning(f'Not supported under this OS: {PLATFORM}') + LOG.error('Not supported under this OS: %s', PLATFORM) + raise NotImplementedError(f'Not supported under this OS: {PLATFORM}') def run_xev() -> None: diff --git a/scripts/wk/hw/network.py b/scripts/wk/hw/network.py index 700ebfea..55d02eb2 100644 --- a/scripts/wk/hw/network.py +++ b/scripts/wk/hw/network.py @@ -9,11 +9,7 @@ from wk.net import ( show_valid_addresses, speedtest, ) -from wk.std import ( - TryAndPrint, - pause, - print_warning, - ) +from wk.ui import cli as ui # STATIC VARIABLES @@ -24,7 +20,7 @@ LOG = logging.getLogger(__name__) def network_test() -> None: """Run network tests.""" LOG.info('Network Test') - try_and_print = TryAndPrint() + try_and_print = ui.TryAndPrint() result = try_and_print.run( message='Network connection...', function=connected_to_private_network, @@ -34,8 +30,8 @@ def network_test() -> None: # Bail if not connected if result['Failed']: - print_warning('Please connect to a network and try again') - pause('Press Enter to return to main menu...') + ui.print_warning('Please connect to a network and try again') + ui.pause('Press Enter to return to main menu...') return # Show IP address(es) @@ -51,7 +47,7 @@ def network_test() -> None: try_and_print.run('Speedtest...', speedtest) # Done - pause('Press Enter to return to main menu...') + ui.pause('Press Enter to return to main menu...') if __name__ == '__main__': diff --git a/scripts/wk/hw/screensavers.py b/scripts/wk/hw/screensavers.py index 4417777b..71e67fcd 100644 --- a/scripts/wk/hw/screensavers.py +++ b/scripts/wk/hw/screensavers.py @@ -6,7 +6,7 @@ import logging from subprocess import PIPE from wk.exe import run_program -from wk.tmux import zoom_pane as tmux_zoom_pane +from wk.ui import tmux # STATIC VARIABLES @@ -31,9 +31,9 @@ def screensaver(name) -> None: ] # Switch pane to fullscreen and start screensaver - tmux_zoom_pane() + tmux.zoom_pane() run_program(cmd, check=False, pipe=False, stderr=PIPE) - tmux_zoom_pane() + tmux.zoom_pane() if __name__ == '__main__': diff --git a/scripts/wk/hw/sensors.py b/scripts/wk/hw/sensors.py index 49d693f0..de72d42b 100644 --- a/scripts/wk/hw/sensors.py +++ b/scripts/wk/hw/sensors.py @@ -7,12 +7,14 @@ import pathlib import re from subprocess import CalledProcessError +from threading import Thread from typing import Any from wk.cfg.hw import CPU_CRITICAL_TEMP, SMC_IDS, TEMP_COLORS from wk.exe import run_program, start_thread from wk.io import non_clobber_path -from wk.std import PLATFORM, color_string, sleep +from wk.std import PLATFORM, sleep +from wk.ui import ansi # STATIC VARIABLES @@ -36,9 +38,9 @@ class ThermalLimitReachedError(RuntimeError): class Sensors(): """Class for holding sensor specific data.""" def __init__(self): - self.background_thread = None - self.data = get_sensor_data() - self.out_path = None + self.background_thread: Thread | None = None + self.data: dict[Any, Any] = get_sensor_data() + self.out_path: pathlib.Path | str | None = None def clear_temps(self) -> None: """Clear saved temps but keep structure""" @@ -109,7 +111,7 @@ class Sensors(): # Handle empty reports if not report: report = [ - color_string('WARNING: No sensors found', 'YELLOW'), + ansi.color_string('WARNING: No sensors found', 'YELLOW'), '', 'Please monitor temps manually', ] @@ -425,7 +427,7 @@ def get_temp_str(temp, colored=True) -> str: temp = float(temp) except (TypeError, ValueError): # Invalid temp? - return color_string(temp, 'PURPLE') + return ansi.color_string(temp, 'PURPLE') # Determine color if colored: @@ -435,7 +437,7 @@ def get_temp_str(temp, colored=True) -> str: break # Done - return color_string(f'{"-" if temp < 0 else ""}{temp:2.0f}°C', temp_color) + return ansi.color_string(f'{"-" if temp < 0 else ""}{temp:2.0f}°C', temp_color) diff --git a/scripts/wk/hw/smart.py b/scripts/wk/hw/smart.py index d9f47f7a..13331efb 100644 --- a/scripts/wk/hw/smart.py +++ b/scripts/wk/hw/smart.py @@ -18,7 +18,8 @@ from wk.cfg.hw import ( SMART_SELF_TEST_START_TIMEOUT_IN_SECONDS, ) from wk.exe import get_json_from_command, run_program -from wk.std import bytes_to_string, color_string, sleep +from wk.std import bytes_to_string, sleep +from wk.ui import ansi # STATIC VARIABLES @@ -40,26 +41,25 @@ def build_self_test_report(test_obj, aborted=False) -> None: For instance if the test was aborted the report should include the last known progress instead of just "was aborted by host." """ - report = [color_string('Self-Test', 'BLUE')] - test_details = get_smart_self_test_details(test_obj.dev) - test_result = test_details.get('status', {}).get('string', 'Unknown') + report = [ansi.color_string('Self-Test', 'BLUE')] + test_result = get_smart_self_test_last_result(test_obj.dev) # Build report if test_obj.disabled or test_obj.status == 'Denied': - report.append(color_string(f' {test_obj.status}', 'RED')) + report.append(ansi.color_string(f' {test_obj.status}', 'RED')) elif test_obj.status == 'N/A' or not test_obj.dev.attributes: - report.append(color_string(f' {test_obj.status}', 'YELLOW')) - elif test_obj.status == 'TestInProgress': - report.append(color_string(' Failed to stop previous test', 'RED')) - test_obj.set_status('Failed') + report.append(ansi.color_string(f' {test_obj.status}', 'YELLOW')) else: # Other cases include self-test result string - report.append(f' {test_result.capitalize()}') - if aborted and not (test_obj.passed or test_obj.failed): - report.append(color_string(' Aborted', 'YELLOW')) - test_obj.set_status('Aborted') + if test_obj.status == 'TestInProgress': + report.append(ansi.color_string(' Failed to stop previous test', 'RED')) + test_obj.set_status('Failed') elif test_obj.status == 'TimedOut': - report.append(color_string(' TimedOut', 'YELLOW')) + report.append(ansi.color_string(' TimedOut', 'YELLOW')) + elif aborted and not (test_obj.passed or test_obj.failed): + report.append(ansi.color_string(' Aborted', 'YELLOW')) + test_obj.set_status('Aborted') + report.append(f' {test_result}') # Done test_obj.report.extend(report) @@ -136,7 +136,7 @@ def generate_attribute_report(dev, only_failed=False) -> list[str]: continue # Build colored string and append to report - line = color_string( + line = ansi.color_string( [label, get_attribute_value_string(dev, attr), note], [None, value_color, 'YELLOW'], ) @@ -200,7 +200,7 @@ def get_attribute_value_string(dev, attr) -> str: return value_str -def get_known_disk_attributes(model) -> None: +def get_known_disk_attributes(model) -> dict[str | int, dict[str, Any]]: """Get known disk attributes based on the device model.""" known_attributes = copy.deepcopy(KNOWN_DISK_ATTRIBUTES) @@ -218,7 +218,7 @@ def get_known_disk_attributes(model) -> None: return known_attributes -def get_smart_self_test_details(dev) -> dict[Any, Any]: +def get_smart_self_test_details(dev) -> dict[str, Any]: """Shorthand to get deeply nested self-test details, returns dict.""" details = {} try: @@ -231,6 +231,33 @@ def get_smart_self_test_details(dev) -> dict[Any, Any]: return details +def get_smart_self_test_last_result(dev) -> str: + """Get last SMART self-test result, returns str.""" + result = 'Unknown' + + # Parse SMART data + data = dev.raw_smartctl.get( + 'ata_smart_self_test_log', {}).get( + 'standard', {}).get( + 'table', []) + try: + data = data[0] + except IndexError: + # No results found + return result + + # Build result string + result = ( + f'Power-on hours: {data.get("lifetime_hours", "?")}' + f', Type: {data.get("type", {}).get("string", "?")}' + f', Passed: {data.get("status", {}).get("passed", "?")}' + f', Result: {data.get("status", {}).get("string", "?")}' + ) + + # Done + return result + + def monitor_smart_self_test(test_obj, header_str, log_path) -> bool: """Monitor SMART self-test status and update test_obj, returns bool.""" started = False @@ -262,6 +289,9 @@ def monitor_smart_self_test(test_obj, header_str, log_path) -> bool: if _i * 5 >= SMART_SELF_TEST_START_TIMEOUT_IN_SECONDS: # Test didn't start within limit, stop waiting abort_self_test(test_obj.dev) + result = get_smart_self_test_last_result(test_obj.dev) + if result == 'Unknown': + result = 'SMART self-test failed to start' test_obj.failed = True test_obj.set_status('TimedOut') break @@ -277,6 +307,11 @@ def monitor_smart_self_test(test_obj, header_str, log_path) -> bool: finished = True break + # Check if timed out + if started and not finished: + test_obj.failed = True + test_obj.set_status('TimedOut') + # Done return finished @@ -290,15 +325,15 @@ def run_self_test(test_obj, log_path) -> None: run_smart_self_test(test_obj, log_path) -def run_smart_self_test(test_obj, log_path) -> bool: - """Run SMART self-test and check if it passed, returns bool. +def run_smart_self_test(test_obj, log_path) -> None: + """Run SMART self-test and check if it passed, returns None. NOTE: An exception will be raised if the disk lacks SMART support. """ finished = False test_details = get_smart_self_test_details(test_obj.dev) size_str = bytes_to_string(test_obj.dev.size, use_binary=False) - header_str = color_string( + header_str = ansi.color_string( ['[', test_obj.dev.path.name, ' ', size_str, ']'], [None, 'BLUE', None, 'CYAN', None], sep='', @@ -348,11 +383,15 @@ def run_smart_self_test(test_obj, log_path) -> bool: # Check result if finished: + test_details = get_smart_self_test_details(test_obj.dev) test_obj.passed = test_details.get('status', {}).get('passed', False) test_obj.failed = test_obj.failed or not test_obj.passed # Set status - if test_obj.failed and test_obj.status != 'TimedOut': + if test_obj.status == 'TimedOut': + # Preserve TimedOut status + pass + elif test_obj.failed: test_obj.set_status('Failed') elif test_obj.passed: test_obj.set_status('Passed') diff --git a/scripts/wk/hw/surface_scan.py b/scripts/wk/hw/surface_scan.py index a7e10fb8..ef99cc20 100644 --- a/scripts/wk/hw/surface_scan.py +++ b/scripts/wk/hw/surface_scan.py @@ -14,12 +14,8 @@ from wk.cfg.hw import ( TEST_MODE_BADBLOCKS_LIMIT, ) from wk.exe import run_program -from wk.std import ( - PLATFORM, - bytes_to_string, - color_string, - strip_colors, - ) +from wk.std import PLATFORM, bytes_to_string +from wk.ui import ansi # STATIC VARIABLES @@ -35,7 +31,7 @@ def check_surface_scan_results(test_obj, log_path) -> None: # Read result with open(log_path, 'r', encoding='utf-8') as _f: for line in _f.readlines(): - line = strip_colors(line.strip()) + line = ansi.strip_colors(line.strip()) if not line: # Skip continue @@ -67,7 +63,7 @@ def check_surface_scan_results(test_obj, log_path) -> None: elif not test_obj.passed: report_color = 'YELLOW' for line in report: - test_obj.report.append(f' {color_string(line, report_color)}') + test_obj.report.append(f' {ansi.color_string(line, report_color)}') # Handle undefined result status if not (test_obj.passed or test_obj.failed): @@ -83,7 +79,7 @@ def run_scan(test_obj, log_path, test_mode=False) -> None: # Use "RAW" disks under macOS dev_path = dev_path.with_name(f'r{dev_path.name}') LOG.info('Using %s for better performance', dev_path) - test_obj.report.append(color_string('badblocks', 'BLUE')) + test_obj.report.append(ansi.color_string('badblocks', 'BLUE')) test_obj.set_status('Working') # Increase block size if necessary @@ -102,7 +98,7 @@ def run_scan(test_obj, log_path, test_mode=False) -> None: with open(log_path, 'a', encoding='utf-8') as _f: size_str = bytes_to_string(dev.size, use_binary=False) _f.write( - color_string( + ansi.color_string( ['[', dev.path.name, ' ', size_str, ']\n'], [None, 'BLUE', None, 'CYAN', None], sep='', diff --git a/scripts/wk/hw/system.py b/scripts/wk/hw/system.py index 300d41c4..80cb161e 100644 --- a/scripts/wk/hw/system.py +++ b/scripts/wk/hw/system.py @@ -13,12 +13,8 @@ from wk.cfg.hw import KNOWN_RAM_VENDOR_IDS from wk.cfg.python import DATACLASS_DECORATOR_KWARGS from wk.exe import get_json_from_command, run_program from wk.hw.test import Test -from wk.std import ( - PLATFORM, - bytes_to_string, - color_string, - string_to_bytes, - ) +from wk.std import PLATFORM, bytes_to_string, string_to_bytes +from wk.ui import ansi # STATIC VARIABLES @@ -42,11 +38,11 @@ class System: def generate_cpu_ram_report(self) -> list[str]: """Generate CPU & RAM report, returns list.""" report = [] - report.append(color_string('Device', 'BLUE')) + report.append(ansi.color_string('Device', 'BLUE')) report.append(f' {self.cpu_description}') # Include RAM details - report.append(color_string('RAM', 'BLUE')) + report.append(ansi.color_string('RAM', 'BLUE')) report.append(f' {self.ram_total} ({", ".join(self.ram_dimms)})') # Tests diff --git a/scripts/wk/hw/volumes.py b/scripts/wk/hw/volumes.py index 92795003..0f288788 100644 --- a/scripts/wk/hw/volumes.py +++ b/scripts/wk/hw/volumes.py @@ -10,7 +10,8 @@ from wk.cfg.hw import ( VOLUME_WARNING_THRESHOLD, VOLUME_SIZE_THRESHOLD, ) -from wk.std import PLATFORM, bytes_to_string, color_string +from wk.std import PLATFORM, bytes_to_string +from wk.ui.ansi import color_string # STATIC VARIABLES diff --git a/scripts/wk/io.py b/scripts/wk/io.py index 8e398da4..4599c898 100644 --- a/scripts/wk/io.py +++ b/scripts/wk/io.py @@ -13,7 +13,7 @@ LOG = logging.getLogger(__name__) # Functions -def case_insensitive_path(path): +def case_insensitive_path(path: pathlib.Path | str) -> pathlib.Path: """Find path case-insensitively, returns pathlib.Path obj.""" given_path = pathlib.Path(path).resolve() real_path = None @@ -37,7 +37,8 @@ def case_insensitive_path(path): return real_path -def case_insensitive_search(path, item): +def case_insensitive_search( + path: pathlib.Path | str, item: str) -> pathlib.Path: """Search path for item case insensitively, returns pathlib.Path obj.""" path = pathlib.Path(path).resolve() given_path = path.joinpath(item) @@ -61,7 +62,10 @@ def case_insensitive_search(path, item): return real_path -def copy_file(source, dest, overwrite=False): +def copy_file( + source: pathlib.Path | str, + dest: pathlib.Path | str, + overwrite: bool = False) -> None: """Copy file and optionally overwrite the destination.""" source = case_insensitive_path(source) dest = pathlib.Path(dest).resolve() @@ -72,7 +76,7 @@ def copy_file(source, dest, overwrite=False): shutil.copy2(source, dest) -def delete_empty_folders(path): +def delete_empty_folders(path: pathlib.Path | str) -> None: """Recursively delete all empty folders in path.""" LOG.debug('path: %s', path) @@ -89,7 +93,11 @@ def delete_empty_folders(path): pass -def delete_folder(path, force=False, ignore_errors=False): +def delete_folder( + path: pathlib.Path | str, + force: bool = False, + ignore_errors: bool = False, + ) -> None: """Delete folder if empty or if forced. NOTE: Exceptions are not caught by this function, @@ -106,7 +114,11 @@ def delete_folder(path, force=False, ignore_errors=False): os.rmdir(path) -def delete_item(path, force=False, ignore_errors=False): +def delete_item( + path: pathlib.Path | str, + force: bool = False, + ignore_errors: bool = False, + ) -> None: """Delete file or folder, optionally recursively. NOTE: Exceptions are not caught by this function, @@ -124,7 +136,11 @@ def delete_item(path, force=False, ignore_errors=False): os.remove(path) -def get_path_obj(path, expanduser=True, resolve=True): +def get_path_obj( + path: pathlib.Path | str, + expanduser: bool = True, + resolve: bool = True, + ) -> pathlib.Path: """Get based on path, returns pathlib.Path.""" path = pathlib.Path(path) if expanduser: @@ -134,7 +150,7 @@ def get_path_obj(path, expanduser=True, resolve=True): return path -def non_clobber_path(path): +def non_clobber_path(path: pathlib.Path | str) -> pathlib.Path: """Update path as needed to non-existing path, returns pathlib.Path.""" LOG.debug('path: %s', path) path = pathlib.Path(path) @@ -163,7 +179,10 @@ def non_clobber_path(path): return new_path -def recursive_copy(source, dest, overwrite=False): +def recursive_copy( + source: pathlib.Path | str, + dest: pathlib.Path | str, + overwrite: bool = False) -> None: """Copy source to dest recursively. NOTE: This uses rsync style source/dest syntax. @@ -213,7 +232,10 @@ def recursive_copy(source, dest, overwrite=False): raise FileExistsError(f'Refusing to delete file: {dest}') -def rename_item(path, new_path): +def rename_item( + path: pathlib.Path | str, + new_path: pathlib.Path | str, + ) -> pathlib.Path: """Rename item, returns pathlib.Path.""" path = pathlib.Path(path) return path.rename(new_path) diff --git a/scripts/wk/kit/build_win.py b/scripts/wk/kit/build_win.py index e4ac2943..201316f2 100644 --- a/scripts/wk/kit/build_win.py +++ b/scripts/wk/kit/build_win.py @@ -6,6 +6,7 @@ NOTE: This script is meant to be called from within a new kit in ConEmu. import logging import os +import pathlib import re from wk.cfg.launchers import LAUNCHERS @@ -21,16 +22,8 @@ from wk.kit.tools import ( get_tool_path, ) from wk.log import update_log_path -from wk.std import ( - GenericError, - TryAndPrint, - clear_screen, - pause, - print_info, - print_success, - set_title, - sleep, - ) +from wk.std import GenericError +from wk.ui import cli as ui # STATIC VARIABLES @@ -52,7 +45,7 @@ WIDTH = 50 # Functions -def compress_cbin_dirs(): +def compress_cbin_dirs() -> None: """Compress CBIN_DIR items using ARCHIVE_PASSWORD.""" current_dir = os.getcwd() for item in CBIN_DIR.iterdir(): @@ -70,25 +63,25 @@ def compress_cbin_dirs(): delete_item(item, force=True, ignore_errors=True) -def delete_from_temp(item_path): +def delete_from_temp(item_path) -> None: """Delete item from temp.""" delete_item(TMP_DIR.joinpath(item_path), force=True, ignore_errors=True) -def download_to_temp(filename, source_url, referer=None): +def download_to_temp(filename, source_url, referer=None) -> pathlib.Path: """Download file to temp dir, returns pathlib.Path.""" out_path = TMP_DIR.joinpath(filename) download_file(out_path, source_url, referer=referer) return out_path -def extract_to_bin(archive, folder): +def extract_to_bin(archive, folder) -> None: """Extract archive to folder under BIN_DIR.""" out_path = BIN_DIR.joinpath(folder) extract_archive(archive, out_path) -def generate_launcher(section, name, options): +def generate_launcher(section, name, options) -> None: """Generate launcher script.""" dest = ROOT_DIR.joinpath(f'{section+"/" if section else ""}{name}.cmd') out_text = [] @@ -115,27 +108,27 @@ def generate_launcher(section, name, options): # Download functions -def download_adobe_reader(): +def download_adobe_reader() -> None: """Download Adobe Reader.""" out_path = INSTALLERS_DIR.joinpath('Adobe Reader DC.exe') download_file(out_path, SOURCES['Adobe Reader DC']) -def download_aida64(): +def download_aida64() -> None: """Download AIDA64.""" archive = download_to_temp('AIDA64.zip', SOURCES['AIDA64']) extract_to_bin(archive, 'AIDA64') delete_from_temp('AIDA64.zip') -def download_autoruns(): +def download_autoruns() -> None: """Download Autoruns.""" for item in ('Autoruns32', 'Autoruns64'): out_path = BIN_DIR.joinpath(f'Sysinternals/{item}.exe') download_file(out_path, SOURCES[item]) -def download_bleachbit(): +def download_bleachbit() -> None: """Download BleachBit.""" out_path = BIN_DIR.joinpath('BleachBit') archive = download_to_temp('BleachBit.zip', SOURCES['BleachBit']) @@ -150,7 +143,7 @@ def download_bleachbit(): delete_from_temp('BleachBit.zip') -def download_bluescreenview(): +def download_bluescreenview() -> None: """Download BlueScreenView.""" archive_32 = download_to_temp( 'bluescreenview32.zip', SOURCES['BlueScreenView32'], @@ -188,14 +181,14 @@ def download_coretemp(): delete_from_temp('coretemp64.zip') -def download_erunt(): +def download_erunt() -> None: """Download ERUNT.""" archive = download_to_temp('erunt.zip', SOURCES['ERUNT']) extract_to_bin(archive, 'ERUNT') delete_from_temp('erunt.zip') -def download_everything(): +def download_everything() -> None: """Download Everything.""" archive_32 = download_to_temp('everything32.zip', SOURCES['Everything32']) archive_64 = download_to_temp('everything64.zip', SOURCES['Everything64']) @@ -210,7 +203,7 @@ def download_everything(): delete_from_temp('everything64.zip') -def download_fastcopy(): +def download_fastcopy() -> None: """Download FastCopy.""" installer = download_to_temp('FastCopyInstaller.exe', SOURCES['FastCopy']) out_path = BIN_DIR.joinpath('FastCopy') @@ -226,7 +219,7 @@ def download_fastcopy(): delete_item(BIN_DIR.joinpath('FastCopy/setup.exe')) -def download_furmark(): +def download_furmark() -> None: """Download FurMark.""" installer = download_to_temp( 'FurMark_Setup.exe', @@ -246,19 +239,19 @@ def download_furmark(): delete_from_temp('FurMarkInstall') -def download_hwinfo(): +def download_hwinfo() -> None: """Download HWiNFO.""" archive = download_to_temp('HWiNFO.zip', SOURCES['HWiNFO']) extract_to_bin(archive, 'HWiNFO') delete_from_temp('HWiNFO.zip') -def download_libreoffice(): +def download_libreoffice() -> None: """Download LibreOffice.""" for arch in 32, 64: out_path = INSTALLERS_DIR.joinpath(f'LibreOffice{arch}.msi') download_file(out_path, SOURCES[f'LibreOffice{arch}']) - sleep(1) + ui.sleep(1) def download_linux_reader(): @@ -269,13 +262,13 @@ def download_linux_reader(): delete_from_temp('LinuxReader.exe') -def download_macs_fan_control(): +def download_macs_fan_control() -> None: """Download Macs Fan Control.""" out_path = INSTALLERS_DIR.joinpath('Macs Fan Control.exe') download_file(out_path, SOURCES['Macs Fan Control']) -def download_neutron(): +def download_neutron() -> None: """Download Neutron.""" archive = download_to_temp('neutron.zip', SOURCES['Neutron']) out_path = BIN_DIR.joinpath('Neutron') @@ -283,7 +276,7 @@ def download_neutron(): delete_from_temp('neutron.zip') -def download_notepad_plus_plus(): +def download_notepad_plus_plus() -> None: """Download Notepad++.""" archive = download_to_temp('npp.7z', SOURCES['Notepad++']) extract_to_bin(archive, 'NotepadPlusPlus') @@ -295,7 +288,7 @@ def download_notepad_plus_plus(): delete_from_temp('npp.7z') -def download_openshell(): +def download_openshell() -> None: """Download OpenShell installer and Fluent-Metro skin.""" for name in ('OpenShell.exe', 'Fluent-Metro.zip'): out_path = BIN_DIR.joinpath(f'OpenShell/{name}') @@ -309,7 +302,7 @@ def download_prime95(): delete_from_temp('prime95.zip') -def download_putty(): +def download_putty() -> None: """Download PuTTY.""" archive = download_to_temp('putty.zip', SOURCES['PuTTY']) extract_to_bin(archive, 'PuTTY') @@ -322,7 +315,7 @@ def download_shutup10(): download_file(out_path, SOURCES['ShutUp10']) -def download_snappy_driver_installer_origin(): +def download_snappy_driver_installer_origin() -> None: """Download Snappy Driver Installer Origin.""" archive = download_to_temp('aria2.zip', SOURCES['Aria2']) aria2c = TMP_DIR.joinpath('aria2/aria2c.exe') @@ -355,7 +348,7 @@ def download_snappy_driver_installer_origin(): cmd.append('-new_console:n') cmd.append('-new_console:s33V') popen_program(cmd, cwd=aria2c.parent) - sleep(1) + ui.sleep(1) wait_for_procs('aria2c.exe') else: run_program(cmd) @@ -392,7 +385,7 @@ def download_snappy_driver_installer_origin(): delete_from_temp('fake.7z') -def download_uninstallview(): +def download_uninstallview() -> None: """Download UninstallView.""" archive_32 = download_to_temp('uninstallview32.zip', SOURCES['UninstallView32']) archive_64 = download_to_temp('uninstallview64.zip', SOURCES['UninstallView64']) @@ -430,14 +423,14 @@ def download_winscp(): delete_from_temp('winscp.zip') -def download_wiztree(): +def download_wiztree() -> None: """Download WizTree.""" archive = download_to_temp('wiztree.zip', SOURCES['WizTree']) extract_to_bin(archive, 'WizTree') delete_from_temp('wiztree.zip') -def download_xmplay(): +def download_xmplay() -> None: """Download XMPlay.""" archives = [ download_to_temp('xmplay.zip', SOURCES['XMPlay']), @@ -465,7 +458,7 @@ def download_xmplay(): delete_from_temp('xmp-rar.zip') delete_from_temp('Innocuous.zip') -def download_xmplay_music(): +def download_xmplay_music() -> None: """Download XMPlay Music.""" music_tmp = TMP_DIR.joinpath('music') music_tmp.mkdir(exist_ok=True) @@ -518,17 +511,17 @@ def download_xmplay_music(): # "Main" Function -def build_kit(): +def build_kit() -> None: """Build Kit.""" update_log_path(dest_name='Build Tool', timestamp=True) title = f'{KIT_NAME_FULL}: Build Tool' - clear_screen() - set_title(title) - print_info(title) + ui.clear_screen() + ui.set_title(title) + ui.print_info(title) print('') # Set up TryAndPrint - try_print = TryAndPrint() + try_print = ui.TryAndPrint() try_print.width = WIDTH try_print.verbose = True for error in ('CalledProcessError', 'FileNotFoundError'): @@ -565,15 +558,15 @@ def build_kit(): # Pause print('', flush=True) - pause('Please review and press Enter to continue...') + ui.pause('Please review and press Enter to continue...') # Compress .cbin try_print.run('Compress cbin...', compress_cbin_dirs) # Generate launcher scripts - print_success('Generating launchers') + ui.print_success('Generating launchers') for section, launchers in sorted(LAUNCHERS.items()): - print_info(f' {section if section else "(Root)"}') + ui.print_info(f' {section if section else "(Root)"}') for name, options in sorted(launchers.items()): try_print.run( f' {name}...', generate_launcher, @@ -583,7 +576,7 @@ def build_kit(): # Done print('') print('Done.') - pause('Press Enter to exit...') + ui.pause('Press Enter to exit...') if __name__ == '__main__': diff --git a/scripts/wk/kit/tools.py b/scripts/wk/kit/tools.py index 641beb2b..8d7251f0 100644 --- a/scripts/wk/kit/tools.py +++ b/scripts/wk/kit/tools.py @@ -1,11 +1,13 @@ """WizardKit: Tool Functions""" # vim: sts=2 sw=2 ts=2 -from datetime import datetime, timedelta import logging import pathlib import platform +from datetime import datetime, timedelta +from subprocess import CompletedProcess, Popen + import requests from wk.cfg.main import ARCHIVE_PASSWORD @@ -30,7 +32,9 @@ CACHED_DIRS = {} # Functions -def download_file(out_path, source_url, as_new=False, overwrite=False, referer=None): +def download_file( + out_path, source_url, + as_new=False, overwrite=False, referer=None) -> pathlib.Path: """Download a file using requests, returns pathlib.Path.""" out_path = pathlib.Path(out_path).resolve() name = out_path.name @@ -38,6 +42,7 @@ def download_file(out_path, source_url, as_new=False, overwrite=False, referer=N download_msg = f'Downloading {name}...' if as_new: out_path = out_path.with_suffix(f'{out_path.suffix}.new') + overwrite = True print(download_msg, end='', flush=True) # Avoid clobbering @@ -94,7 +99,7 @@ def download_file(out_path, source_url, as_new=False, overwrite=False, referer=N return out_path -def download_tool(folder, name, suffix=None): +def download_tool(folder, name, suffix=None) -> None: """Download tool.""" name_arch = f'{name}{ARCH}' out_path = get_tool_path(folder, name, check=False, suffix=suffix) @@ -129,7 +134,7 @@ def download_tool(folder, name, suffix=None): raise -def extract_archive(archive, out_path, *args, mode='x', silent=True): +def extract_archive(archive, out_path, *args, mode='x', silent=True) -> None: """Extract an archive to out_path.""" out_path = pathlib.Path(out_path).resolve() out_path.parent.mkdir(parents=True, exist_ok=True) @@ -141,7 +146,7 @@ def extract_archive(archive, out_path, *args, mode='x', silent=True): run_program(cmd) -def extract_tool(folder): +def extract_tool(folder) -> None: """Extract tool.""" extract_archive( find_kit_dir('.cbin').joinpath(folder).with_suffix('.7z'), @@ -150,7 +155,7 @@ def extract_tool(folder): ) -def find_kit_dir(name=None): +def find_kit_dir(name=None) -> pathlib.Path: """Find folder in kit, returns pathlib.Path. Search is performed in the script's path and then recursively upwards. @@ -177,7 +182,7 @@ def find_kit_dir(name=None): return cur_path -def get_tool_path(folder, name, check=True, suffix=None): +def get_tool_path(folder, name, check=True, suffix=None) -> pathlib.Path: """Get tool path, returns pathlib.Path""" bin_dir = find_kit_dir('.bin') if not suffix: @@ -202,7 +207,7 @@ def run_tool( folder, name, *run_args, cbin=False, cwd=False, download=False, popen=False, **run_kwargs, - ): + ) -> CompletedProcess | Popen: """Run tool from the kit or the Internet, returns proc obj. proc will be either subprocess.CompletedProcess or subprocess.Popen.""" diff --git a/scripts/wk/kit/ufd.py b/scripts/wk/kit/ufd.py index 86705ade..1fa3241d 100644 --- a/scripts/wk/kit/ufd.py +++ b/scripts/wk/kit/ufd.py @@ -4,13 +4,13 @@ import logging import math import os +import pathlib import shutil from subprocess import CalledProcessError -from collections import OrderedDict from docopt import docopt -from wk import io, log, std +from wk import io, log from wk.cfg.main import KIT_NAME_FULL, KIT_NAME_SHORT from wk.cfg.ufd import ( BOOT_ENTRIES, @@ -23,6 +23,8 @@ from wk.cfg.ufd import ( from wk.exe import get_json_from_command, run_program from wk.os import linux +from wk.ui import cli as ui + # STATIC VARIABLES DOCSTRING = '''WizardKit: Build UFD @@ -59,7 +61,7 @@ UFD_LABEL = f'{KIT_NAME_SHORT}_UFD' # Functions -def apply_image(part_path, image_path, hide_macos_boot=True): +def apply_image(part_path, image_path, hide_macos_boot=True) -> None: """Apply raw image to dev_path using dd.""" cmd = [ 'sudo', @@ -89,16 +91,16 @@ def apply_image(part_path, image_path, hide_macos_boot=True): linux.unmount(source_or_mountpoint='/mnt/TMP') -def build_ufd(): +def build_ufd() -> None: """Build UFD using selected sources.""" args = docopt(DOCSTRING) if args['--debug']: log.enable_debug_mode() if args['--update'] and args['EXTRA_IMAGES']: - std.print_warning('Extra images are ignored when updating') + ui.print_warning('Extra images are ignored when updating') args['EXTRA_IMAGES'] = [] log.update_log_path(dest_name='build-ufd', timestamp=True) - try_print = std.TryAndPrint() + try_print = ui.TryAndPrint() try_print.add_error('FileNotFoundError') try_print.catch_all = False try_print.indent = 2 @@ -106,9 +108,9 @@ def build_ufd(): try_print.width = 64 # Show header - std.print_success(KIT_NAME_FULL) - std.print_warning('UFD Build Tool') - std.print_warning(' ') + ui.print_success(KIT_NAME_FULL) + ui.print_warning('UFD Build Tool') + ui.print_warning(' ') # Verify selections ufd_dev = verify_ufd(args['--ufd-device']) @@ -120,7 +122,7 @@ def build_ufd(): # Prep UFD if not args['--update']: - std.print_info('Prep UFD') + ui.print_info('Prep UFD') try_print.run( message='Zeroing first 64MiB...', function=zero_device, @@ -172,8 +174,8 @@ def build_ufd(): ) # Copy sources - std.print_standard(' ') - std.print_info('Copy Sources') + ui.print_standard(' ') + ui.print_info('Copy Sources') try_print.run( 'Copying Memtest86...', io.recursive_copy, '/usr/share/memtest86-efi/', '/mnt/UFD/EFI/Memtest86/', overwrite=True, @@ -189,8 +191,8 @@ def build_ufd(): # Apply extra images if not args['--update']: - std.print_standard(' ') - std.print_info('Apply Extra Images') + ui.print_standard(' ') + ui.print_info('Apply Extra Images') for part_num, image_path in enumerate(extra_images): try_print.run( message=f'Applying {image_path.name}...', @@ -205,8 +207,8 @@ def build_ufd(): _f.write('\n'.join([image.name for image in extra_images])) # Update boot entries - std.print_standard(' ') - std.print_info('Boot Setup') + ui.print_standard(' ') + ui.print_info('Boot Setup') try_print.run( message='Updating boot entries...', function=update_boot_entries, @@ -237,8 +239,8 @@ def build_ufd(): ) # Hide items - std.print_standard(' ') - std.print_info('Final Touches') + ui.print_standard(' ') + ui.print_info('Final Touches') try_print.run( message='Hiding items...', function=hide_items, @@ -247,33 +249,33 @@ def build_ufd(): ) # Done - std.print_standard('\nDone.') + ui.print_standard('\nDone.') if not args['--force']: - std.pause('Press Enter to exit...') + ui.pause('Press Enter to exit...') -def confirm_selections(update=False): +def confirm_selections(update=False) -> None: """Ask tech to confirm selections, twice if necessary.""" - if not std.ask('Is the above information correct?'): - std.abort() + if not ui.ask('Is the above information correct?'): + ui.abort() # Safety check if not update: - std.print_standard(' ') - std.print_warning('SAFETY CHECK') - std.print_standard( + ui.print_standard(' ') + ui.print_warning('SAFETY CHECK') + ui.print_standard( 'All data will be DELETED from the disk and partition(s) listed above.') - std.print_colored( + ui.print_colored( ['This is irreversible and will lead to', 'DATA LOSS'], [None, 'RED'], ) - if not std.ask('Asking again to confirm, is this correct?'): - std.abort() + if not ui.ask('Asking again to confirm, is this correct?'): + ui.abort() - std.print_standard(' ') + ui.print_standard(' ') -def copy_source(source, items, overwrite=False): +def copy_source(source, items, overwrite=False) -> None: """Copy source items to /mnt/UFD.""" is_image = source.is_file() items_not_found = False @@ -300,7 +302,7 @@ def copy_source(source, items, overwrite=False): raise FileNotFoundError('One or more items not found') -def create_table(dev_path, use_mbr=False, images=None): +def create_table(dev_path, use_mbr=False, images=None) -> None: """Create GPT or DOS partition table.""" cmd = [ 'sudo', @@ -338,7 +340,7 @@ def create_table(dev_path, use_mbr=False, images=None): run_program(cmd) -def find_first_partition(dev_path): +def find_first_partition(dev_path) -> str: """Find path to first partition of dev, returns str.""" cmd = [ 'lsblk', @@ -357,7 +359,7 @@ def find_first_partition(dev_path): return part_path -def format_partition(dev_path, label): +def format_partition(dev_path, label) -> None: """Format first partition on device FAT32.""" cmd = [ 'sudo', @@ -369,7 +371,7 @@ def format_partition(dev_path, label): run_program(cmd) -def get_block_device_size(dev_path): +def get_block_device_size(dev_path) -> int: """Get block device size via lsblk, returns int.""" cmd = [ 'lsblk', @@ -388,7 +390,7 @@ def get_block_device_size(dev_path): return int(proc.stdout.strip()) -def get_uuid(path): +def get_uuid(path) -> str: """Get filesystem UUID via findmnt, returns str.""" cmd = [ 'findmnt', @@ -404,7 +406,7 @@ def get_uuid(path): return proc.stdout.strip() -def hide_items(ufd_dev_first_partition, items): +def hide_items(ufd_dev_first_partition, items) -> None: """Set FAT32 hidden flag for items.""" with open('/root/.mtoolsrc', 'w', encoding='utf-8') as _f: _f.write(f'drive U: file="{ufd_dev_first_partition}"\n') @@ -416,7 +418,7 @@ def hide_items(ufd_dev_first_partition, items): run_program(cmd, shell=True, check=False) -def install_syslinux_to_dev(ufd_dev, use_mbr): +def install_syslinux_to_dev(ufd_dev, use_mbr) -> None: """Install Syslinux to UFD (dev).""" cmd = [ 'sudo', @@ -429,7 +431,7 @@ def install_syslinux_to_dev(ufd_dev, use_mbr): run_program(cmd) -def install_syslinux_to_partition(partition): +def install_syslinux_to_partition(partition) -> None: """Install Syslinux to UFD (partition).""" cmd = [ 'sudo', @@ -442,7 +444,7 @@ def install_syslinux_to_partition(partition): run_program(cmd) -def is_valid_path(path_obj, path_type): +def is_valid_path(path_obj, path_type) -> bool: """Verify path_obj is valid by type, returns bool.""" valid_path = False if path_type == 'DIR': @@ -459,7 +461,7 @@ def is_valid_path(path_obj, path_type): return valid_path -def set_boot_flag(dev_path, use_mbr=False): +def set_boot_flag(dev_path, use_mbr=False) -> None: """Set modern or legacy boot flag.""" cmd = [ 'sudo', @@ -471,7 +473,7 @@ def set_boot_flag(dev_path, use_mbr=False): run_program(cmd) -def remove_arch(): +def remove_arch() -> None: """Remove arch dir from UFD. This ensures a clean installation to the UFD and resets the boot files @@ -479,16 +481,16 @@ def remove_arch(): shutil.rmtree(io.case_insensitive_path('/mnt/UFD/arch')) -def show_selections(args, sources, ufd_dev, ufd_sources, extra_images): +def show_selections(args, sources, ufd_dev, ufd_sources, extra_images) -> None: """Show selections including non-specified options.""" # Sources - std.print_info('Sources') + ui.print_info('Sources') for label in ufd_sources.keys(): if label in sources: - std.print_standard(f' {label+":":<18} {sources[label]}') + ui.print_standard(f' {label+":":<18} {sources[label]}') else: - std.print_colored( + ui.print_colored( [f' {label+":":<18}', 'Not Specified'], [None, 'YELLOW'], ) @@ -500,15 +502,15 @@ def show_selections(args, sources, ufd_dev, ufd_sources, extra_images): print(f' {" ":<18} {image}') # Destination - std.print_standard(' ') - std.print_info('Destination') + ui.print_standard(' ') + ui.print_info('Destination') cmd = [ 'lsblk', '--nodeps', '--noheadings', '--paths', '--output', 'NAME,FSTYPE,TRAN,SIZE,VENDOR,MODEL,SERIAL', ufd_dev, ] proc = run_program(cmd, check=False) - std.print_standard(proc.stdout.strip()) + ui.print_standard(proc.stdout.strip()) cmd = [ 'lsblk', '--noheadings', '--paths', '--output', 'NAME,SIZE,FSTYPE,LABEL,MOUNTPOINT', @@ -516,17 +518,17 @@ def show_selections(args, sources, ufd_dev, ufd_sources, extra_images): ] proc = run_program(cmd, check=False) for line in proc.stdout.splitlines()[1:]: - std.print_standard(line) + ui.print_standard(line) # Notes if args['--update']: - std.print_warning('Updating kit in-place') + ui.print_warning('Updating kit in-place') elif args['--use-mbr']: - std.print_warning('Formatting using legacy MBR') - std.print_standard(' ') + ui.print_warning('Formatting using legacy MBR') + ui.print_standard(' ') -def update_boot_entries(ufd_dev, images=None): +def update_boot_entries(ufd_dev, images=None) -> None: """Update boot files for UFD usage""" configs = [] uuids = [get_uuid('/mnt/UFD')] @@ -616,9 +618,9 @@ def update_boot_entries(ufd_dev, images=None): break -def verify_sources(args, ufd_sources): +def verify_sources(args, ufd_sources) -> dict[str, pathlib.Path]: """Check all sources and abort if necessary, returns dict.""" - sources = OrderedDict() + sources = {} for label, data in ufd_sources.items(): s_path = args[data['Arg']] @@ -626,34 +628,35 @@ def verify_sources(args, ufd_sources): try: s_path_obj = io.case_insensitive_path(s_path) except FileNotFoundError: - std.print_error(f'ERROR: {label} not found: {s_path}') - std.abort() - if not is_valid_path(s_path_obj, data['Type']): - std.print_error(f'ERROR: Invalid {label} source: {s_path}') - std.abort() - sources[label] = s_path_obj + ui.print_error(f'ERROR: {label} not found: {s_path}') + ui.abort() + else: + if not is_valid_path(s_path_obj, data['Type']): + ui.print_error(f'ERROR: Invalid {label} source: {s_path}') + ui.abort() + sources[label] = s_path_obj return sources -def verify_ufd(dev_path): +def verify_ufd(dev_path) -> pathlib.Path: """Check that dev_path is a valid UFD, returns pathlib.Path obj.""" ufd_dev = None try: ufd_dev = io.case_insensitive_path(dev_path) except FileNotFoundError: - std.print_error(f'ERROR: UFD device not found: {dev_path}') - std.abort() + ui.print_error(f'ERROR: UFD device not found: {dev_path}') + ui.abort() if not is_valid_path(ufd_dev, 'UFD'): - std.print_error(f'ERROR: Invalid UFD device: {ufd_dev}') - std.abort() + ui.print_error(f'ERROR: Invalid UFD device: {ufd_dev}') + ui.abort() - return ufd_dev + return ufd_dev # type: ignore[reportGeneralTypeIssues] -def zero_device(dev_path): +def zero_device(dev_path) -> None: """Zero-out first 64MB of device.""" cmd = [ 'sudo', diff --git a/scripts/wk/log.py b/scripts/wk/log.py index 3c771940..768b9f13 100644 --- a/scripts/wk/log.py +++ b/scripts/wk/log.py @@ -26,7 +26,7 @@ DEFAULT_LOG_NAME = cfg.main.KIT_NAME_FULL # Functions -def enable_debug_mode(): +def enable_debug_mode() -> None: """Configures logging for better debugging.""" root_logger = logging.getLogger() for handler in root_logger.handlers: @@ -39,8 +39,11 @@ def enable_debug_mode(): def format_log_path( - log_dir=None, log_name=None, timestamp=False, - kit=False, tool=False, append=False): + log_dir: None | pathlib.Path | str = None, + log_name: None | str = None, + timestamp: bool = False, + kit: bool = False, tool: bool = False, append: bool = False, + ) -> pathlib.Path: """Format path based on args passed, returns pathlib.Path obj.""" log_path = pathlib.Path( f'{log_dir if log_dir else DEFAULT_LOG_DIR}/' @@ -61,22 +64,24 @@ def format_log_path( return log_path -def get_root_logger_path(): - """Get path to log file from root logger, returns pathlib.Path obj.""" - log_path = None +def get_root_logger_path() -> pathlib.Path: + """Get the log filepath from the root logger, returns pathlib.Path obj. + + NOTE: This will use the first handler baseFilename it finds (if any). + """ root_logger = logging.getLogger() - # Check all handlers and use the first fileHandler found + # Check handlers for handler in root_logger.handlers: - if isinstance(handler, logging.FileHandler): - log_path = pathlib.Path(handler.baseFilename).resolve() - break + if hasattr(handler, 'baseFilename'): + log_file = handler.baseFilename # type: ignore[reportGeneralTypeIssues] + return pathlib.Path(log_file).resolve() - # Done - return log_path + # No log file found + raise RuntimeError('Log path not found.') -def remove_empty_log(log_path=None): +def remove_empty_log(log_path: None | pathlib.Path = None) -> None: """Remove log if empty. NOTE: Under Windows an empty log is 2 bytes long. @@ -99,7 +104,7 @@ def remove_empty_log(log_path=None): log_path.unlink() -def start(config=None): +def start(config: dict[str, str] | None = None) -> None: """Configure and start logging using safe defaults.""" log_path = format_log_path(timestamp=os.name != 'nt') root_logger = logging.getLogger() @@ -122,7 +127,10 @@ def start(config=None): def update_log_path( - dest_dir=None, dest_name=None, keep_history=True, timestamp=True, append=False): + dest_dir: None | pathlib.Path | str = None, + dest_name: None | str = None, + keep_history: bool = True, timestamp: bool = True, append: bool = False, + ) -> None: """Moves current log file to new path and updates the root logger.""" root_logger = logging.getLogger() new_path = format_log_path(dest_dir, dest_name, timestamp=timestamp, append=append) diff --git a/scripts/wk/net.py b/scripts/wk/net.py index 3e9382a9..14b5fc0e 100644 --- a/scripts/wk/net.py +++ b/scripts/wk/net.py @@ -5,12 +5,16 @@ import os import pathlib import re +from subprocess import CompletedProcess +from typing import Any + import psutil from wk.exe import get_json_from_command, run_program -from wk.std import PLATFORM, GenericError, show_data +from wk.std import PLATFORM, GenericError from wk.cfg.net import BACKUP_SERVERS +from wk.ui import cli as ui # REGEX @@ -22,7 +26,7 @@ REGEX_VALID_IP = re.compile( # Functions -def connected_to_private_network(raise_on_error=False): +def connected_to_private_network(raise_on_error: bool = False) -> bool: """Check if connected to a private network, returns bool. This checks for a valid private IP assigned to this system. @@ -48,12 +52,10 @@ def connected_to_private_network(raise_on_error=False): raise GenericError('Not connected to a network') # Done - if raise_on_error: - connected = None return connected -def mount_backup_shares(read_write=False): +def mount_backup_shares(read_write: bool = False) -> list[str]: """Mount backup shares using OS specific methods.""" report = [] for name, details in BACKUP_SERVERS.items(): @@ -96,7 +98,10 @@ def mount_backup_shares(read_write=False): return report -def mount_network_share(details, mount_point=None, read_write=False): +def mount_network_share( + details: dict[str, Any], + mount_point: None | pathlib.Path | str = None, + read_write: bool = False) -> CompletedProcess: """Mount network share using OS specific methods.""" cmd = None address = details['Address'] @@ -147,7 +152,7 @@ def mount_network_share(details, mount_point=None, read_write=False): return run_program(cmd, check=False) -def ping(addr='google.com'): +def ping(addr: str = 'google.com') -> None: """Attempt to ping addr.""" cmd = ( 'ping', @@ -158,7 +163,7 @@ def ping(addr='google.com'): run_program(cmd) -def share_is_mounted(details): +def share_is_mounted(details: dict[str, Any]) -> bool: """Check if dev/share/etc is mounted, returns bool.""" mounted = False @@ -192,18 +197,20 @@ def share_is_mounted(details): return mounted -def show_valid_addresses(): +def show_valid_addresses() -> None: """Show all valid private IP addresses assigned to the system.""" + # TODO: Refactor to remove ui dependancy devs = psutil.net_if_addrs() for dev, families in sorted(devs.items()): for family in families: if REGEX_VALID_IP.search(family.address): # Valid IP found - show_data(message=dev, data=family.address) + ui.show_data(message=dev, data=family.address) -def speedtest(): +def speedtest() -> list[str]: """Run a network speedtest using speedtest-cli.""" + # TODO: Refactor to use speedtest-cli's JSON output cmd = ['speedtest-cli', '--simple'] proc = run_program(cmd, check=False) output = [line.strip() for line in proc.stdout.splitlines() if line.strip()] @@ -212,7 +219,7 @@ def speedtest(): return [f'{a:<10}{b:6.2f} {c}' for a, b, c in output] -def unmount_backup_shares(): +def unmount_backup_shares() -> list[str]: """Unmount backup shares.""" report = [] for name, details in BACKUP_SERVERS.items(): @@ -241,7 +248,10 @@ def unmount_backup_shares(): return report -def unmount_network_share(details=None, mount_point=None): +def unmount_network_share( + details: dict[str, Any] | None = None, + mount_point: None | pathlib.Path | str = None, + ) -> CompletedProcess: """Unmount network share""" cmd = [] diff --git a/scripts/wk/os/linux.py b/scripts/wk/os/linux.py index 481e687c..ecc2eca4 100644 --- a/scripts/wk/os/linux.py +++ b/scripts/wk/os/linux.py @@ -10,7 +10,8 @@ import subprocess from wk.cfg.hw import VOLUME_FAILURE_THRESHOLD, VOLUME_WARNING_THRESHOLD from wk.exe import get_json_from_command, popen_program, run_program from wk.log import format_log_path -from wk.std import bytes_to_string, color_string +from wk.std import bytes_to_string +from wk.ui import ansi # STATIC VARIABLES @@ -19,12 +20,12 @@ UUID_CORESTORAGE = '53746f72-6167-11aa-aa11-00306543ecac' # Functions -def build_volume_report(device_path=None) -> list: +def build_volume_report(device_path=None) -> list[str]: """Build volume report using lsblk, returns list. If device_path is provided the report is limited to that device. """ - def _get_volumes(dev, indent=0) -> list: + def _get_volumes(dev, indent=0) -> list[dict]: """Convert lsblk JSON tree to a flat list of items, returns list.""" dev['name'] = f'{" "*indent}{dev["name"]}' volumes = [dev] @@ -82,20 +83,20 @@ def build_volume_report(device_path=None) -> list: vol['mountpoint'] = f'Mounted on {vol["mountpoint"]}' # Name and size - line = color_string( + line = ansi.color_string( [f'{vol["name"]:<20}', f'{vol["size"]:>9}'], [None, 'CYAN'], ) # Mountpoint and type - line = color_string( + line = ansi.color_string( [line, f'{vol["mountpoint"]:<{m_width}}', f'{vol["fstype"]:<11}'], [None, None, 'BLUE'], ) # Used and free if any([vol['fsused'], vol['fsavail']]): - line = color_string( + line = ansi.color_string( [line, f'({vol["fsused"]:>9} used, {vol["fsavail"]:>9} free)'], [None, size_color], ) @@ -107,7 +108,7 @@ def build_volume_report(device_path=None) -> list: return report -def get_user_home(user): +def get_user_home(user) -> pathlib.Path: """Get path to user's home dir, returns pathlib.Path obj.""" home = None @@ -128,7 +129,7 @@ def get_user_home(user): return pathlib.Path(home) -def get_user_name(): +def get_user_name() -> str: """Get real user name, returns str.""" user = None @@ -145,7 +146,7 @@ def get_user_name(): return user -def make_temp_file(suffix=None): +def make_temp_file(suffix=None) -> pathlib.Path: """Make temporary file, returns pathlib.Path() obj.""" cmd = ['mktemp'] if suffix: @@ -154,7 +155,7 @@ def make_temp_file(suffix=None): return pathlib.Path(proc.stdout.strip()) -def mount(source, mount_point=None, read_write=False): +def mount(source, mount_point=None, read_write=False) -> None: """Mount source (on mount_point if provided). NOTE: If not running_as_root() then udevil will be used. @@ -177,13 +178,13 @@ def mount(source, mount_point=None, read_write=False): raise RuntimeError(f'Failed to mount: {source} on {mount_point}') -def mount_volumes(device_path=None, read_write=False, scan_corestorage=False): +def mount_volumes(device_path=None, read_write=False, scan_corestorage=False) -> None: """Mount all detected volumes. NOTE: If device_path is specified then only volumes under that path will be mounted. """ - def _get_volumes(dev) -> list: + def _get_volumes(dev) -> list[dict]: """Convert lsblk JSON tree to a flat list of items, returns list.""" volumes = [dev] for child in dev.get('children', []): @@ -232,12 +233,12 @@ def mount_volumes(device_path=None, read_write=False, scan_corestorage=False): pass -def running_as_root(): +def running_as_root() -> bool: """Check if running with effective UID of 0, returns bool.""" return os.geteuid() == 0 -def scan_corestorage_container(container, timeout=300): +def scan_corestorage_container(container, timeout=300) -> list[dict]: """Scan CoreStorage container for inner volumes, returns list.""" container_path = pathlib.Path(container) detected_volumes = {} @@ -284,7 +285,7 @@ def scan_corestorage_container(container, timeout=300): return inner_volumes -def unmount(source_or_mountpoint): +def unmount(source_or_mountpoint) -> None: """Unmount source_or_mountpoint. NOTE: If not running_as_root() then udevil will be used. diff --git a/scripts/wk/os/mac.py b/scripts/wk/os/mac.py index 1a461557..096e21a3 100644 --- a/scripts/wk/os/mac.py +++ b/scripts/wk/os/mac.py @@ -16,7 +16,7 @@ REGEX_FANS = re.compile(r'^.*\(bytes (?P.*)\)$') # Functions -def decode_smc_bytes(text): +def decode_smc_bytes(text) -> int: """Decode SMC bytes, returns int.""" result = None @@ -218,7 +218,7 @@ def mount_disk(device_path=None): return report -def set_fans(mode): +def set_fans(mode) -> None: """Set fans to auto or max.""" if mode == 'auto': set_fans_auto() @@ -228,14 +228,14 @@ def set_fans(mode): raise RuntimeError(f'Invalid fan mode: {mode}') -def set_fans_auto(): +def set_fans_auto() -> None: """Set fans to auto.""" LOG.info('Setting fans to auto') cmd = ['sudo', 'smc', '-k', 'FS! ', '-w', '0000'] run_program(cmd) -def set_fans_max(): +def set_fans_max() -> None: """Set fans to their max speeds.""" LOG.info('Setting fans to max') num_fans = 0 diff --git a/scripts/wk/os/win.py b/scripts/wk/os/win.py index ac899c6b..b6968713 100644 --- a/scripts/wk/os/win.py +++ b/scripts/wk/os/win.py @@ -9,6 +9,8 @@ import platform import re from contextlib import suppress +from typing import Any + import psutil try: @@ -31,11 +33,10 @@ from wk.std import ( GenericError, GenericWarning, bytes_to_string, - color_string, - input_text, - pause, sleep, ) +from wk.ui import cli as ui +from wk.ui import ansi # STATIC VARIABLES @@ -94,7 +95,7 @@ else: # Activation Functions -def activate_with_bios(): +def activate_with_bios() -> None: """Attempt to activate Windows with a key stored in the BIOS.""" # Code borrowed from https://github.com/aeruder/get_win8key ##################################################### @@ -134,7 +135,7 @@ def activate_with_bios(): raise GenericError('Activation Failed') -def get_activation_string(): +def get_activation_string() -> str: """Get activation status, returns str.""" cmd = ['cscript', '//nologo', SLMGR, '/xpr'] proc = run_program(cmd, check=False) @@ -144,7 +145,7 @@ def get_activation_string(): return act_str -def is_activated(): +def is_activated() -> bool: """Check if Windows is activated via slmgr.vbs and return bool.""" act_str = get_activation_string() @@ -153,22 +154,22 @@ def is_activated(): # Date / Time functions -def get_timezone(): +def get_timezone() -> str: """Get current timezone using tzutil, returns str.""" cmd = ['tzutil', '/g'] proc = run_program(cmd, check=False) return proc.stdout -def set_timezone(zone): +def set_timezone(zone) -> None: """Set current timezone using tzutil.""" cmd = ['tzutil', '/s', zone] run_program(cmd, check=False) # Info Functions -def check_4k_alignment(show_alert=False): - """Check if all partitions are 4K aligned, returns book.""" +def check_4k_alignment(show_alert=False) -> list[str]: + """Check if all partitions are 4K aligned, returns list.""" cmd = ['WMIC', 'partition', 'get', 'Caption,Size,StartingOffset'] report = [] show_alert = False @@ -183,9 +184,9 @@ def check_4k_alignment(show_alert=False): if not match: LOG.error('Failed to parse partition info for: %s', line) continue - if int(match.group('offset')) % 4096 == 0: + if int(match.group('offset')) % 4096 != 0: report.append( - color_string( + ansi.color_string( f'{match.group("description")}' f' ({bytes_to_string(match.group("size"), decimals=1)})' , @@ -201,7 +202,7 @@ def check_4k_alignment(show_alert=False): if report: report.insert( 0, - color_string('One or more partitions not 4K aligned', 'YELLOW'), + ansi.color_string('One or more partitions not 4K aligned', 'YELLOW'), ) return report @@ -214,8 +215,8 @@ def defender_is_disabled(): return bool(disabled) -def export_bitlocker_info(): - """Get Bitlocker info and save to the current directory.""" +def export_bitlocker_info() -> None: + """Get Bitlocker info and save to either the base directory of the kit or osTicket.""" commands = [ ['manage-bde', '-status', SYSTEMDRIVE], ['manage-bde', '-protectors', '-get', SYSTEMDRIVE], @@ -260,7 +261,7 @@ def export_bitlocker_info(): if ost.disabled or ost.errors: result = 'Unknown' print( - color_string( + ansi.color_string( ['\nPost info to osTicket... ', result], [None, 'GREEN' if result == 'OK' else 'YELLOW'], ) @@ -270,16 +271,16 @@ def export_bitlocker_info(): if ost.ticket_name: file_name = f'{ost.ticket_id}_{ost.ticket_name.replace(" ", "-")}' if not file_name: - file_name = input_text(prompt='Enter filename', allow_empty_response=False) + file_name = ui.input_text(prompt_msg='Enter filename: ', allow_empty=False) file_path = pathlib.Path(f'../../../Bitlocker_{file_name}.txt').resolve() with open(file_path, 'a', encoding='utf-8') as _f: _f.write('\n'.join(output_raw)) # Done - pause('\nPress Enter to exit...') + ui.pause('\nPress Enter to exit...') -def get_installed_antivirus(): +def get_installed_antivirus() -> list[str]: """Get list of installed antivirus programs, returns list.""" cmd = [ 'WMIC', r'/namespace:\\root\SecurityCenter2', @@ -308,20 +309,20 @@ def get_installed_antivirus(): state = proc.stdout.split('=')[1] state = hex(int(state)) if str(state)[3:5] not in ['10', '11']: - report.append(color_string(f'[Disabled] {product}', 'YELLOW')) + report.append(ansi.color_string(f'[Disabled] {product}', 'YELLOW')) else: report.append(product) # Final check if not report: - report.append(color_string('No products detected', 'RED')) + report.append(ansi.color_string('No products detected', 'RED')) # Done return report -def get_installed_ram(as_list=False, raise_exceptions=False): - """Get installed RAM.""" +def get_installed_ram(as_list=False, raise_exceptions=False) -> list | str: + """Get installed RAM, returns list or str.""" mem = psutil.virtual_memory() mem_str = bytes_to_string(mem.total, decimals=1) @@ -336,8 +337,8 @@ def get_installed_ram(as_list=False, raise_exceptions=False): return [mem_str] if as_list else mem_str -def get_os_activation(as_list=False, check=True): - """Get OS activation status, returns str. +def get_os_activation(as_list=False, check=True) -> list | str: + """Get OS activation status, returns list or str. NOTE: If check=True then raise an exception if OS isn't activated. """ @@ -353,7 +354,7 @@ def get_os_activation(as_list=False, check=True): return [act_str] if as_list else act_str -def get_os_name(as_list=False, check=True): +def get_os_name(as_list=False, check=True) -> str: """Build OS display name, returns str. NOTE: If check=True then an exception is raised if the OS version is @@ -379,7 +380,7 @@ def get_os_name(as_list=False, check=True): return [display_name] if as_list else display_name -def get_raw_disks(): +def get_raw_disks() -> list[str]: """Get all disks without a partiton table, returns list.""" script_path = find_kit_dir('Scripts').joinpath('get_raw_disks.ps1') cmd = ['PowerShell', '-ExecutionPolicy', 'Bypass', '-File', script_path] @@ -404,7 +405,7 @@ def get_raw_disks(): return raw_disks -def get_volume_usage(use_colors=False): +def get_volume_usage(use_colors=False) -> list[str]: """Get space usage info for all fixed volumes, returns list.""" report = [] for disk in psutil.disk_partitions(): @@ -421,14 +422,14 @@ def get_volume_usage(use_colors=False): f' ({bytes_to_string(free, 2):>10} / {bytes_to_string(total, 2):>10})' ) if use_colors: - display_str = color_string(display_str, color) + display_str = ansi.color_string(display_str, color) report.append(f'{disk.device} {display_str}') # Done return report -def show_alert_box(message, title=None): +def show_alert_box(message, title=None) -> None: """Show Windows alert box with message.""" title = title if title else f'{KIT_NAME_FULL} Warning' message_box = ctypes.windll.user32.MessageBoxW @@ -436,7 +437,7 @@ def show_alert_box(message, title=None): # Registry Functions -def reg_delete_key(hive, key, recurse=False): +def reg_delete_key(hive, key, recurse=False) -> None: """Delete a key from the registry. NOTE: If recurse is False then it will only work on empty keys. @@ -458,7 +459,7 @@ def reg_delete_key(hive, key, recurse=False): except FileNotFoundError: # Ignore pass - except PermissionError: + except PermissionError as _e: LOG.error(r'Failed to delete registry key: %s\%s', hive_name, key) if recurse: # Re-raise exception @@ -466,10 +467,10 @@ def reg_delete_key(hive, key, recurse=False): # recurse is not True so assuming we tried to remove a non-empty key msg = fr'Refusing to remove non-empty key: {hive_name}\{key}' - raise FileExistsError(msg) + raise FileExistsError(msg) from _e -def reg_delete_value(hive, key, value): +def reg_delete_value(hive, key, value) -> None: """Delete a value from the registry.""" access = winreg.KEY_ALL_ACCESS hive = reg_get_hive(hive) @@ -493,8 +494,9 @@ def reg_delete_value(hive, key, value): raise -def reg_get_hive(hive): +def reg_get_hive(hive) -> Any: """Get winreg HKEY constant from string, returns HKEY constant.""" + # TODO: Fix type hint if isinstance(hive, int): # Assuming we're already a winreg HKEY constant pass @@ -505,8 +507,9 @@ def reg_get_hive(hive): return hive -def reg_get_data_type(data_type): +def reg_get_data_type(data_type) -> Any: """Get registry data type from string, returns winreg constant.""" + # TODO: Fix type hint if isinstance(data_type, int): # Assuming we're already a winreg value type constant pass @@ -517,7 +520,7 @@ def reg_get_data_type(data_type): return data_type -def reg_key_exists(hive, key): +def reg_key_exists(hive, key) -> bool: """Test if the specified hive/key exists, returns bool.""" exists = False hive = reg_get_hive(hive) @@ -535,7 +538,7 @@ def reg_key_exists(hive, key): return exists -def reg_read_value(hive, key, value, force_32=False, force_64=False): +def reg_read_value(hive, key, value, force_32=False, force_64=False) -> Any: """Query value from hive/hey, returns multiple types. NOTE: Set value='' to read the default value. @@ -559,7 +562,7 @@ def reg_read_value(hive, key, value, force_32=False, force_64=False): return data -def reg_write_settings(settings): +def reg_write_settings(settings) -> None: """Set registry values in bulk from a custom data structure. Data structure should be as follows: @@ -599,7 +602,7 @@ def reg_write_settings(settings): reg_set_value(hive, key, *value) -def reg_set_value(hive, key, name, data, data_type, option=None): +def reg_set_value(hive, key, name, data, data_type, option=None) -> None: """Set value for hive/key.""" access = winreg.KEY_WRITE data_type = reg_get_data_type(data_type) @@ -631,25 +634,25 @@ def reg_set_value(hive, key, name, data, data_type, option=None): # Safe Mode Functions -def disable_safemode(): +def disable_safemode() -> None: """Edit BCD to remove safeboot value.""" cmd = ['bcdedit', '/deletevalue', '{default}', 'safeboot'] run_program(cmd) -def disable_safemode_msi(): +def disable_safemode_msi() -> None: """Disable MSI access under safemode.""" cmd = ['reg', 'delete', REG_MSISERVER, '/f'] run_program(cmd) -def enable_safemode(): +def enable_safemode() -> None: """Edit BCD to set safeboot as default.""" cmd = ['bcdedit', '/set', '{default}', 'safeboot', 'network'] run_program(cmd) -def enable_safemode_msi(): +def enable_safemode_msi() -> None: """Enable MSI access under safemode.""" cmd = ['reg', 'add', REG_MSISERVER, '/f'] run_program(cmd) @@ -662,7 +665,7 @@ def enable_safemode_msi(): # Secure Boot Functions -def is_booted_uefi(): +def is_booted_uefi() -> bool: """Check if booted UEFI or legacy, returns bool.""" kernel = ctypes.windll.kernel32 firmware_type = ctypes.c_uint() @@ -678,7 +681,7 @@ def is_booted_uefi(): return firmware_type.value == 2 -def is_secure_boot_enabled(raise_exceptions=False, show_alert=False): +def is_secure_boot_enabled(raise_exceptions=False, show_alert=False) -> bool: """Check if Secure Boot is enabled, returns bool. If raise_exceptions is True then an exception is raised with details. @@ -728,7 +731,7 @@ def is_secure_boot_enabled(raise_exceptions=False, show_alert=False): # Service Functions -def disable_service(service_name): +def disable_service(service_name) -> None: """Set service startup to disabled.""" cmd = ['sc', 'config', service_name, 'start=', 'disabled'] run_program(cmd, check=False) @@ -738,7 +741,7 @@ def disable_service(service_name): raise GenericError(f'Failed to disable service {service_name}') -def enable_service(service_name, start_type='auto'): +def enable_service(service_name, start_type='auto') -> None: """Enable service by setting start type.""" cmd = ['sc', 'config', service_name, 'start=', start_type] psutil_type = 'automatic' @@ -753,7 +756,7 @@ def enable_service(service_name, start_type='auto'): raise GenericError(f'Failed to enable service {service_name}') -def get_service_status(service_name): +def get_service_status(service_name) -> str: """Get service status using psutil, returns str.""" status = 'unknown' try: @@ -765,7 +768,7 @@ def get_service_status(service_name): return status -def get_service_start_type(service_name): +def get_service_start_type(service_name) -> str: """Get service startup type using psutil, returns str.""" start_type = 'unknown' try: @@ -777,7 +780,7 @@ def get_service_start_type(service_name): return start_type -def start_service(service_name): +def start_service(service_name) -> None: """Stop service.""" cmd = ['net', 'start', service_name] run_program(cmd, check=False) @@ -787,7 +790,7 @@ def start_service(service_name): raise GenericError(f'Failed to start service {service_name}') -def stop_service(service_name): +def stop_service(service_name) -> None: """Stop service.""" cmd = ['net', 'stop', service_name] run_program(cmd, check=False) diff --git a/scripts/wk/repairs/win.py b/scripts/wk/repairs/win.py index 5e3d64cf..b91008a5 100644 --- a/scripts/wk/repairs/win.py +++ b/scripts/wk/repairs/win.py @@ -4,11 +4,13 @@ import atexit import logging import os +import pathlib import re import sys import time from subprocess import CalledProcessError, DEVNULL +from typing import Any from xml.dom.minidom import parse as xml_parse from wk.cfg.main import KIT_NAME_FULL, KIT_NAME_SHORT, WINDOWS_TIME_ZONE @@ -67,21 +69,10 @@ from wk.os.win import ( from wk.std import ( GenericError, GenericWarning, - Menu, - TryAndPrint, - abort, - ask, - clear_screen, - color_string, - pause, - print_info, - print_standard, - print_warning, - set_title, - show_data, sleep, - strip_colors, ) +from wk.ui import cli as ui +from wk.ui import ansi # STATIC VARIABLES @@ -96,7 +87,7 @@ GPUPDATE_SUCCESS_STRINGS = ( 'User Policy update has completed successfully.', ) IN_CONEMU = 'ConEmuPID' in os.environ -MENU_PRESETS = Menu() +MENU_PRESETS = ui.Menu() PROGRAMDATA = os.environ.get('{ALLUSERSPROFILE}', r'C:\ProgramData') PROGRAMFILES_32 = os.environ.get( 'PROGRAMFILES(X86)', os.environ.get( @@ -123,7 +114,7 @@ WHITELIST = '\n'.join(( fr'{PROGRAMFILES_32}\TeamViewer\tv_x64.exe', sys.executable, )) -TRY_PRINT = TryAndPrint() +TRY_PRINT = ui.TryAndPrint() TRY_PRINT.width = WIDTH TRY_PRINT.verbose = True for error in ('CalledProcessError', 'FileNotFoundError'): @@ -131,10 +122,10 @@ for error in ('CalledProcessError', 'FileNotFoundError'): # Auto Repairs -def build_menus(base_menus, title, presets): +def build_menus(base_menus, title, presets) -> dict[str, ui.Menu]: """Build menus, returns dict.""" menus = {} - menus['Main'] = Menu(title=f'{title}\n{color_string("Main Menu", "GREEN")}') + menus['Main'] = ui.Menu(title=f'{title}\n{ansi.color_string("Main Menu", "GREEN")}') # Main Menu for entry in base_menus['Actions']: @@ -143,7 +134,7 @@ def build_menus(base_menus, title, presets): menus['Main'].add_option(group, {'Selected': True}) # Options - menus['Options'] = Menu(title=f'{title}\n{color_string("Options", "GREEN")}') + menus['Options'] = ui.Menu(title=f'{title}\n{ansi.color_string("Options", "GREEN")}') for entry in base_menus['Options']: menus['Options'].add_option(entry.name, entry.details) menus['Options'].add_action('All') @@ -153,7 +144,7 @@ def build_menus(base_menus, title, presets): # Run groups for group, entries in base_menus['Groups'].items(): - menus[group] = Menu(title=f'{title}\n{color_string(group, "GREEN")}') + menus[group] = ui.Menu(title=f'{title}\n{ansi.color_string(group, "GREEN")}') menus[group].disabled_str = 'Locked' for entry in entries: menus[group].add_option(entry.name, entry.details) @@ -185,7 +176,7 @@ def build_menus(base_menus, title, presets): ) # Update presets Menu - MENU_PRESETS.title = f'{title}\n{color_string("Load Preset", "GREEN")}' + MENU_PRESETS.title = f'{title}\n{ansi.color_string("Load Preset", "GREEN")}' MENU_PRESETS.add_option('Default') for name in presets: MENU_PRESETS.add_option(name) @@ -198,7 +189,7 @@ def build_menus(base_menus, title, presets): return menus -def end_session(menus): +def end_session(menus: dict[str, ui.Menu]) -> None: """End Auto Repairs session.""" # Remove logon task cmd = [ @@ -245,7 +236,7 @@ def end_session(menus): ) -def get_entry_settings(group, name): +def get_entry_settings(group, name) -> dict[str, Any]: """Get menu entry settings from the registry, returns dict.""" key_path = fr'{AUTO_REPAIR_KEY}\{group}\{name}' settings = {} @@ -264,7 +255,7 @@ def get_entry_settings(group, name): return settings -def init(menus, presets): +def init(menus, presets) -> None: """Initialize Auto Repairs.""" session_started = is_session_started() @@ -287,14 +278,14 @@ def init(menus, presets): # Resume session load_settings(menus) - print_info('Resuming session, press CTRL+c to cancel') + ui.print_info('Resuming session, press CTRL+c to cancel') 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('') -def init_run(options): +def init_run(options) -> None: """Initialize Auto Repairs Run.""" update_scheduled_task() if options['Kill Explorer']['Selected']: @@ -322,7 +313,7 @@ def init_run(options): TRY_PRINT.run('Running RKill...', run_rkill, msg_good='DONE') -def init_session(options): +def init_session(options) -> None: """Initialize Auto Repairs session.""" reg_set_value('HKCU', AUTO_REPAIR_KEY, 'SessionStarted', 1, 'DWORD') reg_set_value('HKCU', AUTO_REPAIR_KEY, 'LogName', get_root_logger_path().stem, 'SZ') @@ -333,7 +324,7 @@ def init_session(options): 'The timezone is currently set to ' f'{zone}, switch it to {WINDOWS_TIME_ZONE}?' ) - if zone != WINDOWS_TIME_ZONE and ask(msg): + if zone != WINDOWS_TIME_ZONE and ui.ask(msg): set_timezone(WINDOWS_TIME_ZONE) # One-time tasks @@ -346,7 +337,7 @@ def init_session(options): print('') -def is_autologon_enabled(): +def is_autologon_enabled() -> bool: """Check if Autologon is enabled, returns bool.""" auto_admin_logon = False try: @@ -364,7 +355,7 @@ def is_autologon_enabled(): return auto_admin_logon -def is_session_started(): +def is_session_started() -> bool: """Check if session was started, returns bool.""" session_started = False try: @@ -376,7 +367,7 @@ def is_session_started(): return session_started -def load_preset(menus, presets, enable_menu_exit=True): +def load_preset(menus, presets, enable_menu_exit=True) -> None: """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}) @@ -402,26 +393,26 @@ def load_preset(menus, presets, enable_menu_exit=True): MENU_PRESETS.actions['Main Menu'].update({'Disabled':False, 'Hidden':False}) -def load_settings(menus): +def load_settings(menus) -> None: """Load session settings from the registry.""" for group, menu in menus.items(): if group == 'Main': continue for name in menu.options: - menu.options[name].update(get_entry_settings(group, strip_colors(name))) + menu.options[name].update(get_entry_settings(group, ansi.strip_colors(name))) -def run_auto_repairs(base_menus, presets): +def run_auto_repairs(base_menus, presets) -> None: """Run Auto Repairs.""" set_log_path() title = f'{KIT_NAME_FULL}: Auto Repairs' - clear_screen() - set_title(title) - print_info(title) + ui.clear_screen() + ui.set_title(title) + ui.print_info(title) print('') # Generate menus - print_standard('Initializing...') + ui.print_standard('Initializing...') menus = build_menus(base_menus, title, presets) # Init @@ -439,21 +430,21 @@ def run_auto_repairs(base_menus, presets): try: show_main_menu(base_menus, menus, presets, title) except SystemExit: - if ask('End session?'): + if ui.ask('End session?'): end_session(menus) raise # Start or resume repairs - clear_screen() - print_standard(title) + ui.clear_screen() + ui.print_standard(title) print('') save_selection_settings(menus) - print_info('Initializing...') + ui.print_info('Initializing...') init_run(menus['Options'].options) save_selection_settings(menus) if not session_started: init_session(menus['Options'].options) - print_info('Running repairs') + ui.print_info('Running repairs') # Run repairs for group, menu in menus.items(): @@ -462,19 +453,19 @@ def run_auto_repairs(base_menus, presets): try: run_group(group, menu) except KeyboardInterrupt: - abort() + ui.abort() # Done end_session(menus) - print_info('Done') - pause('Press Enter to exit...') + ui.print_info('Done') + ui.pause('Press Enter to exit...') -def run_group(group, menu): +def run_group(group, menu) -> None: """Run entries in group if appropriate.""" - print_info(f' {group}') + ui.print_info(f' {group}') for name, details in menu.options.items(): - name_str = strip_colors(name) + name_str = ansi.strip_colors(name) skipped = details.get('Skipped', False) done = details.get('Done', False) disabled = details.get('Disabled', False) @@ -488,7 +479,7 @@ def run_group(group, menu): # Previously skipped if skipped: - show_data(f'{name_str}...', 'Skipped', 'YELLOW', width=WIDTH) + ui.show_data(f'{name_str}...', 'Skipped', 'YELLOW', width=WIDTH) continue # Previously ran @@ -498,7 +489,7 @@ def run_group(group, menu): color = 'YELLOW' elif details.get('Failed', False): color = 'RED' - show_data( + ui.show_data( f'{name_str}...', details.get('Message', 'Unknown'), color, width=WIDTH, ) @@ -506,7 +497,7 @@ def run_group(group, menu): # Not selected if not selected: - show_data(f'{name_str}...', 'Skipped', 'YELLOW', width=WIDTH) + ui.show_data(f'{name_str}...', 'Skipped', 'YELLOW', width=WIDTH) save_settings(group, name, skipped=True) continue @@ -514,7 +505,7 @@ def run_group(group, menu): details['Function'](group, name) -def save_selection_settings(menus): +def save_selection_settings(menus) -> None: """Save selections in the registry.""" for group, menu in menus.items(): if group == 'Main': @@ -527,9 +518,9 @@ def save_selection_settings(menus): ) -def save_settings(group, name, result=None, **kwargs): +def save_settings(group, name, result=None, **kwargs) -> None: """Save entry settings in the registry.""" - key_path = fr'{AUTO_REPAIR_KEY}\{group}\{strip_colors(name)}' + key_path = fr'{AUTO_REPAIR_KEY}\{group}\{ansi.strip_colors(name)}' # Get values from TryAndPrint result if result: @@ -543,7 +534,7 @@ def save_settings(group, name, result=None, **kwargs): # Write values to registry for value_name, data in kwargs.items(): - value_name = strip_colors(value_name) + value_name = ansi.strip_colors(value_name) if isinstance(data, bool): data = 1 if data else 0 if isinstance(data, int): @@ -555,7 +546,7 @@ def save_settings(group, name, result=None, **kwargs): reg_set_value('HKCU', key_path, value_name, data, data_type) -def set_log_path(): +def set_log_path() -> None: """Set log name using defaults or the saved registry value.""" try: log_path = reg_read_value('HKCU', AUTO_REPAIR_KEY, 'LogName') @@ -571,7 +562,7 @@ def set_log_path(): ) -def show_main_menu(base_menus, menus, presets, title): +def show_main_menu(base_menus, menus, presets, title) -> None: """Show main menu and handle actions.""" while True: update_main_menu(menus) @@ -586,7 +577,7 @@ def show_main_menu(base_menus, menus, presets, title): raise SystemExit -def show_sub_menu(menu): +def show_sub_menu(menu) -> None: """Show sub-menu and handle sub-menu actions.""" while True: selection = menu.advanced_select() @@ -617,7 +608,7 @@ def show_sub_menu(menu): menu.options[name][key] = value -def update_main_menu(menus): +def update_main_menu(menus) -> None: """Update main menu based on current selections.""" index = 1 skip = 'Reboot' @@ -650,7 +641,7 @@ def update_scheduled_task(): # Auto Repairs: Wrapper Functions -def auto_adwcleaner(group, name): +def auto_adwcleaner(group, name) -> None: """Run AdwCleaner scan. save_settings() is called first since AdwCleaner may kill this script. @@ -662,25 +653,25 @@ def auto_adwcleaner(group, name): save_settings(group, name, result=result) -def auto_backup_browser_profiles(group, name): +def auto_backup_browser_profiles(group, name) -> None: """Backup browser profiles.""" backup_all_browser_profiles(use_try_print=True) save_settings(group, name, done=True, failed=False, message='DONE') -def auto_backup_power_plans(group, name): +def auto_backup_power_plans(group, name) -> None: """Backup power plans.""" result = TRY_PRINT.run('Backup Power Plans...', export_power_plans) save_settings(group, name, result=result) -def auto_backup_registry(group, name): +def auto_backup_registry(group, name) -> None: """Backup registry.""" result = TRY_PRINT.run('Backup Registry...', backup_registry) save_settings(group, name, result=result) -def auto_bleachbit(group, name): +def auto_bleachbit(group, name) -> None: """Run BleachBit to clean files.""" result = TRY_PRINT.run( 'BleachBit...', run_bleachbit, BLEACH_BIT_CLEANERS, msg_good='DONE', @@ -688,7 +679,7 @@ def auto_bleachbit(group, name): save_settings(group, name, result=result) -def auto_chkdsk(group, name): +def auto_chkdsk(group, name) -> None: """Run CHKDSK repairs.""" needs_reboot = False result = TRY_PRINT.run(f'CHKDSK ({SYSTEMDRIVE})...', run_chkdsk_online) @@ -734,7 +725,7 @@ def auto_enable_defender(group, name): save_settings(group, name, result=result) -def auto_disable_pending_renames(group, name): +def auto_disable_pending_renames(group, name) -> None: """Disable pending renames.""" result = TRY_PRINT.run( 'Disabling pending renames...', disable_pending_renames, @@ -742,7 +733,7 @@ def auto_disable_pending_renames(group, name): save_settings(group, name, result=result) -def auto_dism(group, name): +def auto_dism(group, name) -> None: """Run DISM repairs.""" needs_reboot = False result = TRY_PRINT.run('DISM (RestoreHealth)...', run_dism) @@ -782,7 +773,7 @@ def auto_emsisoft_cmd_uninstall(group, name): save_settings(group, name, result=result) -def auto_enable_regback(group, name): +def auto_enable_regback(group, name) -> None: """Enable RegBack.""" result = TRY_PRINT.run( 'Enable RegBack...', reg_set_value, 'HKLM', @@ -800,13 +791,13 @@ def auto_fix_file_associations(group, name): save_settings(group, name, result=result) -def auto_hitmanpro(group, name): +def auto_hitmanpro(group, name) -> None: """Run HitmanPro scan.""" result = TRY_PRINT.run('HitmanPro...', run_hitmanpro, msg_good='DONE') save_settings(group, name, result=result) -def auto_kvrt(group, name): +def auto_kvrt(group, name) -> None: """Run KVRT scan.""" result = TRY_PRINT.run('KVRT...', run_kvrt, msg_good='DONE') save_settings(group, name, result=result) @@ -834,7 +825,7 @@ def auto_mbam_uninstall(group, name): save_settings(group, name, result=result) -def auto_microsoft_defender(group, name): +def auto_microsoft_defender(group, name) -> None: """Run Microsoft Defender scan.""" result = TRY_PRINT.run( 'Microsoft Defender...', run_microsoft_defender, msg_good='DONE', @@ -842,14 +833,14 @@ def auto_microsoft_defender(group, name): save_settings(group, name, result=result) -def auto_reboot(group, name): +def auto_reboot(group, name) -> None: """Reboot the system.""" save_settings(group, name, done=True, failed=False, message='DONE') print('') reboot(30) -def auto_remove_power_plan(group, name): +def auto_remove_power_plan(group, name) -> None: """Remove custom power plan and set to Balanced.""" result = TRY_PRINT.run( 'Remove Custom Power Plan...', remove_custom_power_plan, @@ -858,7 +849,7 @@ def auto_remove_power_plan(group, name): save_settings(group, name, result=result) -def auto_repair_registry(group, name): +def auto_repair_registry(group, name) -> None: """Delete registry keys with embedded null characters.""" result = TRY_PRINT.run( 'Running Registry repairs...', delete_registry_null_keys, @@ -866,19 +857,19 @@ def auto_repair_registry(group, name): save_settings(group, name, result=result) -def auto_reset_power_plans(group, name): +def auto_reset_power_plans(group, name) -> None: """Reset power plans.""" result = TRY_PRINT.run('Reset Power Plans...', reset_power_plans) save_settings(group, name, result=result) -def auto_reset_proxy(group, name): +def auto_reset_proxy(group, name) -> None: """Reset proxy settings.""" result = TRY_PRINT.run('Clearing proxy settings...', reset_proxy) save_settings(group, name, result=result) -def auto_reset_windows_policies(group, name): +def auto_reset_windows_policies(group, name) -> None: """Reset Windows policies to defaults.""" result = TRY_PRINT.run( 'Resetting Windows policies...', reset_windows_policies, @@ -886,13 +877,13 @@ def auto_reset_windows_policies(group, name): save_settings(group, name, result=result) -def auto_restore_uac_defaults(group, name): +def auto_restore_uac_defaults(group, name) -> None: """Restore UAC default settings.""" result = TRY_PRINT.run('Restoring UAC defaults...', restore_uac_defaults) save_settings(group, name, result=result) -def auto_set_custom_power_plan(group, name): +def auto_set_custom_power_plan(group, name) -> None: """Set custom power plan.""" result = TRY_PRINT.run( 'Set Custom Power Plan...', create_custom_power_plan, @@ -902,13 +893,13 @@ def auto_set_custom_power_plan(group, name): save_settings(group, name, result=result) -def auto_sfc(group, name): +def auto_sfc(group, name) -> None: """Run SFC repairs.""" result = TRY_PRINT.run('SFC Scan...', run_sfc_scan) save_settings(group, name, result=result) -def auto_system_restore_create(group, name): +def auto_system_restore_create(group, name) -> None: """Create System Restore point.""" result = TRY_PRINT.run( 'Create System Restore...', create_system_restore_point, @@ -916,7 +907,7 @@ def auto_system_restore_create(group, name): save_settings(group, name, result=result) -def auto_system_restore_enable(group, name): +def auto_system_restore_enable(group, name) -> None: """Enable System Restore.""" cmd = [ 'powershell', '-Command', 'Enable-ComputerRestore', @@ -926,13 +917,13 @@ def auto_system_restore_enable(group, name): save_settings(group, name, result=result) -def auto_system_restore_set_size(group, name): +def auto_system_restore_set_size(group, name) -> None: """Set System Restore size.""" result = TRY_PRINT.run('Set System Restore Size...', set_system_restore_size) save_settings(group, name, result=result) -def auto_uninstallview(group, name): +def auto_uninstallview(group, name) -> None: """Run UninstallView.""" result = TRY_PRINT.run( 'UninstallView...', run_uninstallview, msg_good='DONE', @@ -940,7 +931,7 @@ def auto_uninstallview(group, name): save_settings(group, name, result=result) -def auto_windows_updates_disable(group, name): +def auto_windows_updates_disable(group, name) -> None: """Disable Windows Updates.""" result = TRY_PRINT.run('Disable Windows Updates...', disable_windows_updates) if result['Failed']: @@ -949,13 +940,13 @@ def auto_windows_updates_disable(group, name): save_settings(group, name, result=result) -def auto_windows_updates_enable(group, name): +def auto_windows_updates_enable(group, name) -> None: """Enable Windows Updates.""" result = TRY_PRINT.run('Enable Windows Updates...', enable_windows_updates) save_settings(group, name, result=result) -def auto_windows_updates_reset(group, name): +def auto_windows_updates_reset(group, name) -> None: """Reset Windows Updates.""" result = TRY_PRINT.run('Reset Windows Updates...', reset_windows_updates) if result['Failed']: @@ -965,12 +956,12 @@ def auto_windows_updates_reset(group, name): # Misc Functions -def set_backup_path(name, date=False): +def set_backup_path(name, date=False) -> pathlib.Path: """Set backup path, returns pathlib.Path.""" return set_local_storage_path('Backups', name, date) -def set_local_storage_path(folder, name, date=False): +def set_local_storage_path(folder, name, date=False) -> pathlib.Path: """Get path for local storage, returns pathlib.Path.""" local_path = get_path_obj(f'{SYSTEMDRIVE}/{KIT_NAME_SHORT}/{folder}/{name}') if date: @@ -978,22 +969,22 @@ def set_local_storage_path(folder, name, date=False): return local_path -def set_quarantine_path(name, date=False): +def set_quarantine_path(name, date=False) -> pathlib.Path: """Set quarantine path, returns pathlib.Path.""" return set_local_storage_path('Quarantine', name, date) # Tool Functions -def backup_all_browser_profiles(use_try_print=False): +def backup_all_browser_profiles(use_try_print=False) -> None: """Backup browser profiles for all users.""" users = get_path_obj(f'{SYSTEMDRIVE}/Users') for userprofile in users.iterdir(): if use_try_print: - print_info(f'{" "*6}{userprofile.name}') + ui.print_info(f'{" "*6}{userprofile.name}') backup_browser_profiles(userprofile, use_try_print) -def backup_browser_chromium(backup_path, browser, search_path, use_try_print): +def backup_browser_chromium(backup_path, browser, search_path, use_try_print) -> None: """Backup Chromium-based browser profile.""" for item in search_path.iterdir(): match = re.match(r'^(Default|Profile).*', item.name, re.IGNORECASE) @@ -1003,7 +994,7 @@ def backup_browser_chromium(backup_path, browser, search_path, use_try_print): if output_path.exists(): # Assuming backup was already done if use_try_print: - show_data( + ui.show_data( f'{" "*8}{browser} ({item.name})...', 'Backup already exists.', color='YELLOW', width=WIDTH, ) @@ -1023,7 +1014,7 @@ def backup_browser_chromium(backup_path, browser, search_path, use_try_print): run_program(cmd, check=False) -def backup_browser_firefox(backup_path, search_path, use_try_print): +def backup_browser_firefox(backup_path, search_path, use_try_print) -> None: """Backup Firefox browser profile.""" output_path = backup_path.joinpath('Firefox.7z') @@ -1031,7 +1022,7 @@ def backup_browser_firefox(backup_path, search_path, use_try_print): if output_path.exists(): # Assuming backup was already done if use_try_print: - show_data( + ui.show_data( f'{" "*8}Firefox (All)...', 'Backup already exists.', color='YELLOW', width=WIDTH, ) @@ -1048,7 +1039,7 @@ def backup_browser_firefox(backup_path, search_path, use_try_print): run_program(cmd, check=False) -def backup_browser_profiles(userprofile, use_try_print=False): +def backup_browser_profiles(userprofile, use_try_print=False) -> None: """Backup browser profiles for userprofile.""" backup_path = set_backup_path('Browsers', date=True) backup_path = backup_path.joinpath(userprofile.name) @@ -1077,7 +1068,7 @@ def backup_browser_profiles(userprofile, use_try_print=False): pass -def backup_registry(): +def backup_registry() -> None: """Backup Registry.""" backup_path = set_backup_path('Registry', date=True) backup_path.parent.mkdir(parents=True, exist_ok=True) @@ -1101,7 +1092,7 @@ def delete_emsisoft_cmd_service(): run_program(['sc', 'delete', 'epp'], check=False) -def delete_registry_null_keys(): +def delete_registry_null_keys() -> None: """Delete registry keys with embedded null characters.""" run_tool('RegDelNull', 'RegDelNull', '-s', '-y', download=True) @@ -1145,7 +1136,7 @@ def install_mbam(): proc.wait() -def log_kvrt_results(log_path, report_path): +def log_kvrt_results(log_path, report_path) -> None: """Parse KVRT report and log results in plain text.""" log_text = '' report_file = None @@ -1186,7 +1177,7 @@ def log_kvrt_results(log_path, report_path): log_path.write_text(log_text, encoding='utf-8') -def run_adwcleaner(): +def run_adwcleaner() -> None: """Run AdwCleaner.""" settings_path = get_tool_path('AdwCleaner', 'AdwCleaner', check=False) settings_path = settings_path.with_name('settings') @@ -1196,7 +1187,7 @@ def run_adwcleaner(): run_tool('AdwCleaner', 'AdwCleaner', download=True) -def run_bleachbit(cleaners, preview=True): +def run_bleachbit(cleaners, preview=True) -> None: """Run BleachBit to either clean or preview files.""" cmd_args = ( '--preview' if preview else '--clean', @@ -1207,8 +1198,12 @@ def run_bleachbit(cleaners, preview=True): proc = run_tool('BleachBit', 'bleachbit_console', *cmd_args) # Save logs - log_path.write_text(proc.stdout, encoding='utf-8') - log_path.with_suffix('.err').write_text(proc.stderr, encoding='utf-8') + log_path.write_text( + proc.stdout, encoding='utf-8', # type: ignore[reportGeneralTypeIssues] + ) + log_path.with_suffix('.err').write_text( + proc.stderr, encoding='utf-8', # type: ignore[reportGeneralTypeIssues] + ) def run_emsisoft_cmd_scan(): @@ -1246,7 +1241,7 @@ def run_emsisoft_cmd_scan(): proc.wait() -def run_hitmanpro(): +def run_hitmanpro() -> None: """Run HitmanPro scan.""" log_path = format_log_path(log_name='HitmanPro', timestamp=True, tool=True) log_path = log_path.with_suffix('.xml') @@ -1259,7 +1254,7 @@ def run_hitmanpro(): proc.wait() -def run_kvrt(): +def run_kvrt() -> None: """Run KVRT scan.""" log_path = format_log_path(log_name='KVRT', timestamp=True, tool=True) log_path.parent.mkdir(parents=True, exist_ok=True) @@ -1325,11 +1320,11 @@ def run_mbam(): run_program(exe_path, check=False) -def run_microsoft_defender(full=True): +def run_microsoft_defender(full=True) -> None: """Run Microsoft Defender scan.""" reg_key = r'Software\Microsoft\Windows Defender' - def _get_defender_path(): + def _get_defender_path() -> str: install_path = reg_read_value('HKLM', reg_key, 'InstallLocation') return fr'{install_path}\MpCmdRun.exe' @@ -1369,7 +1364,7 @@ def run_microsoft_defender(full=True): raise GenericError('Failed to run scan or clean items.') -def run_rkill(): +def run_rkill() -> None: """Run RKill scan.""" log_path = format_log_path(log_name='RKill', timestamp=True, tool=True) log_path.parent.mkdir(parents=True, exist_ok=True) @@ -1383,7 +1378,7 @@ def run_rkill(): run_tool('RKill', 'RKill', *cmd_args, download=True) -def run_tdsskiller(): +def run_tdsskiller() -> None: """Run TDSSKiller scan.""" log_path = format_log_path(log_name='TDSSKiller', timestamp=True, tool=True) log_path.parent.mkdir(parents=True, exist_ok=True) @@ -1434,7 +1429,7 @@ def update_emsisoft_cmd(): # OS Built-in Functions -def create_custom_power_plan(enable_sleep=True, keep_display_on=False): +def create_custom_power_plan(enable_sleep=True, keep_display_on=False) -> None: """Create new power plan and set as active.""" custom_guid = POWER_PLANS['Custom'] sleep_timeouts = POWER_PLAN_SLEEP_TIMEOUTS['High Performance'] @@ -1483,7 +1478,7 @@ def create_custom_power_plan(enable_sleep=True, keep_display_on=False): run_program(cmd) -def create_system_restore_point(): +def create_system_restore_point() -> None: """Create System Restore point.""" cmd = [ 'powershell', '-Command', 'Checkpoint-Computer', @@ -1498,7 +1493,7 @@ def create_system_restore_point(): raise GenericWarning('Skipped, a restore point was created too recently') -def disable_pending_renames(): +def disable_pending_renames() -> None: """Disable pending renames.""" reg_set_value( 'HKLM', r'SYSTEM\CurrentControlSet\Control\Session Manager', @@ -1506,18 +1501,18 @@ def disable_pending_renames(): ) -def disable_windows_updates(): +def disable_windows_updates() -> None: """Disable and stop Windows Updates.""" disable_service('wuauserv') stop_service('wuauserv') -def enable_windows_updates(): +def enable_windows_updates() -> None: """Enable Windows Updates.""" enable_service('wuauserv', 'demand') -def export_power_plans(): +def export_power_plans() -> None: """Export existing power plans.""" backup_path = set_backup_path('Power Plans', date=True) @@ -1548,7 +1543,7 @@ def export_power_plans(): run_program(cmd) -def kill_explorer(): +def kill_explorer() -> None: """Kill all Explorer processes.""" cmd = ['taskkill', '/im', 'explorer.exe', '/f'] run_program(cmd, check=False) @@ -1611,17 +1606,17 @@ def open_defender_settings(disable=False, enable=False): kill_explorer() -def reboot(timeout=10): +def reboot(timeout=10) -> None: """Reboot the system.""" atexit.unregister(start_explorer) - print_warning(f'Rebooting the system in {timeout} seconds...') + ui.print_warning(f'Rebooting the system in {timeout} seconds...') sleep(timeout) cmd = ['shutdown', '-r', '-t', '0'] run_program(cmd, check=False) raise SystemExit -def remove_custom_power_plan(high_performance=False): +def remove_custom_power_plan(high_performance=False) -> None: """Remove custom power plan and set to a built-in plan. If high_performance is True then set to High Performance and set @@ -1648,13 +1643,13 @@ def remove_custom_power_plan(high_performance=False): run_program(cmd) -def reset_power_plans(): +def reset_power_plans() -> None: """Reset power plans to their default settings.""" cmd = ['powercfg', '-RestoreDefaultSchemes'] run_program(cmd) -def reset_proxy(): +def reset_proxy() -> None: """Reset WinHTTP proxy settings.""" cmd = ['netsh', 'winhttp', 'reset', 'proxy'] proc = run_program(cmd, check=False) @@ -1664,7 +1659,7 @@ def reset_proxy(): raise GenericError('Failed to reset proxy settings.') -def reset_windows_policies(): +def reset_windows_policies() -> None: """Reset Windows policies to defaults.""" cmd = ['gpupdate', '/force'] proc = run_program(cmd, check=False) @@ -1674,7 +1669,7 @@ def reset_windows_policies(): raise GenericError('Failed to reset one or more policies.') -def reset_windows_updates(): +def reset_windows_updates() -> None: """Reset Windows Updates.""" system_root = os.environ.get('SYSTEMROOT', 'C:/Windows') src_path = f'{system_root}/SoftwareDistribution' @@ -1687,7 +1682,7 @@ def reset_windows_updates(): pass -def restore_uac_defaults(): +def restore_uac_defaults() -> None: """Restore UAC default settings.""" settings = REG_UAC_DEFAULTS_WIN10 if OS_VERSION in (7, 8, 8.1): @@ -1696,7 +1691,7 @@ def restore_uac_defaults(): reg_write_settings(settings) -def run_chkdsk_offline(): +def run_chkdsk_offline() -> None: """Set filesystem 'dirty bit' to force a CHKDSK during startup.""" cmd = ['fsutil', 'dirty', 'set', SYSTEMDRIVE] proc = run_program(cmd, check=False) @@ -1706,7 +1701,7 @@ def run_chkdsk_offline(): raise GenericError('Failed to set dirty bit.') -def run_chkdsk_online(): +def run_chkdsk_online() -> None: """Run CHKDSK. NOTE: If run on Windows 8+ online repairs are attempted. @@ -1746,7 +1741,7 @@ def run_chkdsk_online(): raise GenericError('Issue(s) detected') -def run_dism(repair=True): +def run_dism(repair=True) -> None: """Run DISM to either scan or repair component store health.""" conemu_args = ['-new_console:nb', '-new_console:s33V'] if IN_CONEMU else [] @@ -1785,7 +1780,7 @@ def run_dism(repair=True): raise GenericError('Issue(s) detected') -def run_sfc_scan(): +def run_sfc_scan() -> None: """Run SFC and save results.""" cmd = ['sfc', '/scannow'] log_path = format_log_path(log_name='SFC', timestamp=True, tool=True) @@ -1812,12 +1807,12 @@ def run_sfc_scan(): raise OSError -def run_uninstallview(): +def run_uninstallview() -> None: """Run UninstallView.""" run_tool('UninstallView', 'UninstallView') -def set_system_restore_size(size=8): +def set_system_restore_size(size=8) -> None: """Set System Restore size.""" cmd = [ 'vssadmin', 'Resize', 'ShadowStorage', @@ -1826,7 +1821,7 @@ def set_system_restore_size(size=8): run_program(cmd, pipe=False, stderr=DEVNULL, stdout=DEVNULL) -def start_explorer(): +def start_explorer() -> None: """Start Explorer.""" popen_program(['explorer.exe']) diff --git a/scripts/wk/setup/win.py b/scripts/wk/setup/win.py index f473f256..e8e526cd 100644 --- a/scripts/wk/setup/win.py +++ b/scripts/wk/setup/win.py @@ -10,6 +10,8 @@ import re import sys import webbrowser +from typing import Any + from wk.cfg.main import KIT_NAME_FULL from wk.cfg.setup import ( BROWSER_PATHS, @@ -65,22 +67,10 @@ from wk.repairs.win import ( from wk.std import ( GenericError, GenericWarning, - Menu, - TryAndPrint, - abort, - ask, - clear_screen, - color_string, - pause, - print_error, - print_info, - print_standard, - print_warning, - set_title, - show_data, sleep, - strip_colors, ) +from wk.ui import cli as ui +from wk.ui import ansi # STATIC VARIABLES @@ -96,7 +86,7 @@ KNOWN_ENCODINGS = ( 'utf-32-le', ) IN_CONEMU = 'ConEmuPID' in os.environ -MENU_PRESETS = Menu() +MENU_PRESETS = ui.Menu() PROGRAMFILES_32 = os.environ.get( 'PROGRAMFILES(X86)', os.environ.get( 'PROGRAMFILES', r'C:\Program Files (x86)', @@ -108,7 +98,7 @@ PROGRAMFILES_64 = os.environ.get( ), ) SYSTEMDRIVE = os.environ.get('SYSTEMDRIVE', 'C:') -TRY_PRINT = TryAndPrint() +TRY_PRINT = ui.TryAndPrint() TRY_PRINT.width = WIDTH TRY_PRINT.verbose = True for error in ('CalledProcessError', 'FileNotFoundError'): @@ -116,10 +106,10 @@ for error in ('CalledProcessError', 'FileNotFoundError'): # Auto Setup -def build_menus(base_menus, title, presets): +def build_menus(base_menus, title, presets) -> dict[str, ui.Menu]: """Build menus, returns dict.""" menus = {} - menus['Main'] = Menu(title=f'{title}\n{color_string("Main Menu", "GREEN")}') + menus['Main'] = ui.Menu(title=f'{title}\n{ansi.color_string("Main Menu", "GREEN")}') # Main Menu for entry in base_menus['Actions']: @@ -129,7 +119,7 @@ def build_menus(base_menus, title, presets): # Run groups for group, entries in base_menus['Groups'].items(): - menus[group] = Menu(title=f'{title}\n{color_string(group, "GREEN")}') + menus[group] = ui.Menu(title=f'{title}\n{ansi.color_string(group, "GREEN")}') for entry in entries: menus[group].add_option(entry.name, entry.details) menus[group].add_action('All') @@ -158,7 +148,7 @@ def build_menus(base_menus, title, presets): ) # Update presets Menu - MENU_PRESETS.title = f'{title}\n{color_string("Load Preset", "GREEN")}' + MENU_PRESETS.title = f'{title}\n{ansi.color_string("Load Preset", "GREEN")}' MENU_PRESETS.add_option('Default') for name in presets: MENU_PRESETS.add_option(name) @@ -184,38 +174,38 @@ def check_if_av_scan_is_needed(): # Check date and prompt tech if necessary last_run_date = datetime.strptime(last_run, '%Y-%m-%d') if datetime.now() - last_run_date < timedelta(days=FAB_TIMEFRAME): - print_warning('Fab was recently run, an AV scan may be needed.') - if not ask('Continue with setup?'): - abort() + ui.print_warning('Fab was recently run, an AV scan may be needed.') + if not ui.ask('Continue with setup?'): + ui.abort() -def check_os_and_set_menu_title(title): +def check_os_and_set_menu_title(title) -> str: """Check OS version and update title for menus, returns str.""" color = None os_name = get_os_name(check=False) - print_standard(f'Operating System: {os_name}') + ui.print_standard(f'Operating System: {os_name}') # Check support status and set color try: get_os_name() except GenericWarning: # Outdated version - print_warning('OS version is outdated, updating is recommended.') - if not ask('Continue anyway?'): - abort() + ui.print_warning('OS version is outdated, updating is recommended.') + if not ui.ask('Continue anyway?'): + ui.abort() color = 'YELLOW' except GenericError: # Unsupported version - print_error('OS version is unsupported, updating is recommended.') - if not ask('Continue anyway? (NOT RECOMMENDED)'): - abort() + ui.print_error('OS version is unsupported, updating is recommended.') + if not ui.ask('Continue anyway? (NOT RECOMMENDED)'): + ui.abort() color = 'RED' # Done - return f'{title} ({color_string(os_name, color)})' + return f'{title} ({ansi.color_string(os_name, color)})' -def load_preset(menus, presets, title, enable_menu_exit=True): +def load_preset(menus, presets, title, enable_menu_exit=True) -> None: """Load menu settings from preset and ask selection question(s).""" msp = False @@ -241,37 +231,37 @@ def load_preset(menus, presets, title, enable_menu_exit=True): menu.options[name]['Selected'] = value # Ask selection question(s) - clear_screen() - print_standard(f'{title}') + ui.clear_screen() + ui.print_standard(f'{title}') print('') if selection[0] == 'Default': # OpenShell - if OS_VERSION != 11 and ask('Install OpenShell?'): + if OS_VERSION != 11 and ui.ask('Install OpenShell?'): menus['Install Software'].options['Open-Shell']['Selected'] = True menus['Configure System'].options['Open-Shell']['Selected'] = True # LibreOffice - if ask('Install LibreOffice?'): + if ui.ask('Install LibreOffice?'): menus['Install Software'].options['LibreOffice']['Selected'] = True # Hiberboot & Hibernation print('') msg = 'Disable Fast Startup and enable Hibernation? (Recommended for SSDs)' - if ask(msg): + if ui.ask(msg): for option in ('Disable Fast Startup', 'Enable Hibernation'): menus['Configure System'].options[option]['Selected'] = True # Apply ITS settings? - msp = ask('Is this an ITS system?') + msp = ui.ask('Is this an ITS system?') if msp: option = 'Apply ITS Settings' menus['Configure System'].options[option]['Selected'] = True # ESET NOD32 AV print('') - if msp or ask('Install ESET NOD32 AV?'): + if msp or ui.ask('Install ESET NOD32 AV?'): option = 'ESET NOD32 AV' - if msp or ask(' For VIP?'): + if msp or ui.ask(' For VIP?'): option = f'{option} (MSP)' menus['Install Software'].options[option]['Selected'] = True @@ -285,15 +275,15 @@ def load_preset(menus, presets, title, enable_menu_exit=True): menus[group_name].options[entry_name]['Selected'] = False -def run_auto_setup(base_menus, presets): +def run_auto_setup(base_menus, presets) -> None: """Run Auto Setup.""" update_log_path(dest_name='Auto Setup', timestamp=True) title = f'{KIT_NAME_FULL}: Auto Setup' - clear_screen() - set_title(title) - print_info(title) + ui.clear_screen() + ui.set_title(title) + ui.print_info(title) print('') - print_standard('Initializing...') + ui.print_standard('Initializing...') # Check OS and update title for menus title = check_os_and_set_menu_title(title) @@ -311,10 +301,10 @@ def run_auto_setup(base_menus, presets): show_main_menu(base_menus, menus, presets, title) # Start setup - clear_screen() - print_standard(title) + ui.clear_screen() + ui.print_standard(title) print('') - print_info('Running setup') + ui.print_info('Running setup') # Run setup for group, menu in menus.items(): @@ -323,29 +313,29 @@ def run_auto_setup(base_menus, presets): try: run_group(group, menu) except KeyboardInterrupt: - abort() + ui.abort() # Done - print_info('Done') - pause('Press Enter to exit...') + ui.print_info('Done') + ui.pause('Press Enter to exit...') -def run_group(group, menu): +def run_group(group, menu) -> None: """Run entries in group if appropriate.""" - print_info(f' {group}') + ui.print_info(f' {group}') for name, details in menu.options.items(): - name_str = strip_colors(name) + name_str = ansi.strip_colors(name) # Not selected if not details.get('Selected', False): - show_data(f'{name_str}...', 'Skipped', 'YELLOW', width=WIDTH) + ui.show_data(f'{name_str}...', 'Skipped', 'YELLOW', width=WIDTH) continue # Selected details['Function']() -def show_main_menu(base_menus, menus, presets, title): +def show_main_menu(base_menus, menus, presets, title) -> None: """Show main menu and handle actions.""" while True: update_main_menu(menus) @@ -360,7 +350,7 @@ def show_main_menu(base_menus, menus, presets, title): raise SystemExit -def show_sub_menu(menu): +def show_sub_menu(menu) -> None: """Show sub-menu and handle sub-menu actions.""" while True: selection = menu.advanced_select() @@ -376,7 +366,7 @@ def show_sub_menu(menu): menu.options[name]['Selected'] = value -def update_main_menu(menus): +def update_main_menu(menus) -> None: """Update main menu based on current selections.""" index = 1 skip = 'Reboot' @@ -395,37 +385,37 @@ def update_main_menu(menus): # Auto Repairs: Wrapper Functions -def auto_backup_registry(): +def auto_backup_registry() -> None: """Backup registry.""" TRY_PRINT.run('Backup Registry...', backup_registry) -def auto_backup_browser_profiles(): +def auto_backup_browser_profiles() -> None: """Backup browser profiles.""" backup_all_browser_profiles(use_try_print=True) -def auto_backup_power_plans(): +def auto_backup_power_plans() -> None: """Backup power plans.""" TRY_PRINT.run('Backup Power Plans...', export_power_plans) -def auto_reset_power_plans(): +def auto_reset_power_plans() -> None: """Reset power plans.""" TRY_PRINT.run('Reset Power Plans...', reset_power_plans) -def auto_set_custom_power_plan(): +def auto_set_custom_power_plan() -> None: """Set custom power plan.""" TRY_PRINT.run('Set Custom Power Plan...', create_custom_power_plan) -def auto_enable_bsod_minidumps(): +def auto_enable_bsod_minidumps() -> None: """Enable saving minidumps during BSoDs.""" TRY_PRINT.run('Enable BSoD mini dumps...', enable_bsod_minidumps) -def auto_enable_regback(): +def auto_enable_regback() -> None: """Enable RegBack.""" TRY_PRINT.run( 'Enable RegBack...', reg_set_value, 'HKLM', @@ -439,7 +429,7 @@ def auto_apply_its_settings(): TRY_PRINT.run('Apply ITS settings...', apply_its_settings) -def auto_system_restore_enable(): +def auto_system_restore_enable() -> None: """Enable System Restore.""" cmd = [ 'powershell', '-Command', 'Enable-ComputerRestore', @@ -448,28 +438,28 @@ def auto_system_restore_enable(): TRY_PRINT.run('Enable System Restore...', run_program, cmd=cmd) -def auto_system_restore_set_size(): +def auto_system_restore_set_size() -> None: """Set System Restore size.""" TRY_PRINT.run('Set System Restore Size...', set_system_restore_size) -def auto_system_restore_create(): +def auto_system_restore_create() -> None: """Create System Restore point.""" TRY_PRINT.run('Create System Restore...', create_system_restore_point) -def auto_windows_updates_enable(): +def auto_windows_updates_enable() -> None: """Enable Windows Updates.""" TRY_PRINT.run('Enable Windows Updates...', enable_windows_updates) # Auto Setup: Wrapper Functions -def auto_activate_windows(): +def auto_activate_windows() -> None: """Attempt to activate Windows using BIOS key.""" TRY_PRINT.run('Windows Activation...', activate_with_bios) -def auto_config_browsers(): +def auto_config_browsers() -> None: """Configure Browsers.""" prompt = ' Press Enter to continue...' TRY_PRINT.run('Chrome Notifications...', disable_chrome_notifications) @@ -480,18 +470,18 @@ def auto_config_browsers(): 'Set default browser...', set_default_browser, msg_good='STARTED', ) print(prompt, end='', flush=True) - pause('') + ui.pause('') # Move cursor to beginning of the previous line and clear prompt print(f'\033[F\r{" "*len(prompt)}\r', end='', flush=True) -def auto_config_explorer(): +def auto_config_explorer() -> None: """Configure Windows Explorer and restart the process.""" TRY_PRINT.run('Windows Explorer...', config_explorer) -def auto_config_open_shell(): +def auto_config_open_shell() -> None: """Configure Open Shell.""" TRY_PRINT.run('Open Shell...', config_open_shell) @@ -512,7 +502,7 @@ def auto_enable_hibernation(): ) -def auto_export_aida64_report(): +def auto_export_aida64_report() -> None: """Export AIDA64 reports.""" TRY_PRINT.run('AIDA64 Report...', export_aida64_report) @@ -527,12 +517,12 @@ def auto_install_eset_nod32_av_msp(): TRY_PRINT.run('ESET NOD32 (MSP)...', install_eset_nod32_av, msp=True) -def auto_install_firefox(): +def auto_install_firefox() -> None: """Install Firefox.""" TRY_PRINT.run('Firefox...', install_firefox) -def auto_install_libreoffice(): +def auto_install_libreoffice() -> None: """Install LibreOffice. NOTE: It is assumed that auto_install_vcredists() will be run @@ -544,32 +534,32 @@ def auto_install_libreoffice(): ) -def auto_install_open_shell(): +def auto_install_open_shell() -> None: """Install Open Shell.""" TRY_PRINT.run('Open Shell...', install_open_shell) -def auto_install_software_bundle(): +def auto_install_software_bundle() -> None: """Install standard software bundle.""" TRY_PRINT.run('Software Bundle...', install_software_bundle) -def auto_install_vcredists(): +def auto_install_vcredists() -> None: """Install latest supported Visual C++ runtimes.""" TRY_PRINT.run('Visual C++ Runtimes...', install_vcredists) -def auto_open_device_manager(): +def auto_open_device_manager() -> None: """Open Device Manager.""" TRY_PRINT.run('Device Manager...', open_device_manager) -def auto_open_hwinfo_sensors(): +def auto_open_hwinfo_sensors() -> None: """Open HWiNFO Sensors.""" TRY_PRINT.run('HWiNFO Sensors...', open_hwinfo_sensors) -def auto_open_snappy_driver_installer_origin(): +def auto_open_snappy_driver_installer_origin() -> None: """Open Snappy Driver Installer Origin.""" TRY_PRINT.run('Snappy Driver Installer...', open_snappy_driver_installer_origin) @@ -579,57 +569,57 @@ def auto_open_webcam_tests(): TRY_PRINT.run('Webcam Tests...', open_webcam_tests) -def auto_open_windows_activation(): +def auto_open_windows_activation() -> None: """Open Windows Activation.""" if not is_activated(): TRY_PRINT.run('Windows Activation...', open_windows_activation) -def auto_open_windows_updates(): +def auto_open_windows_updates() -> None: """Open Windows Updates.""" TRY_PRINT.run('Windows Updates...', open_windows_updates) -def auto_open_xmplay(): +def auto_open_xmplay() -> None: """Open XMPlay.""" TRY_PRINT.run('XMPlay...', open_xmplay) -def auto_show_4k_alignment_check(): +def auto_show_4k_alignment_check() -> None: """Display 4K alignment check.""" TRY_PRINT.run('4K alignment Check...', check_4k_alignment, show_alert=True) -def auto_show_installed_antivirus(): +def auto_show_installed_antivirus() -> None: """Display installed antivirus.""" TRY_PRINT.run('Virus Protection...', get_installed_antivirus) -def auto_show_installed_ram(): +def auto_show_installed_ram() -> None: """Display installed RAM.""" TRY_PRINT.run('Installed RAM...', get_installed_ram, as_list=True, raise_exceptions=True, ) -def auto_show_os_activation(): +def auto_show_os_activation() -> None: """Display OS activation status.""" TRY_PRINT.run('Activation...', get_os_activation, as_list=True) -def auto_show_os_name(): +def auto_show_os_name() -> None: """Display OS Name.""" TRY_PRINT.run('Operating System...', get_os_name, as_list=True) -def auto_show_secure_boot_status(): +def auto_show_secure_boot_status() -> None: """Display Secure Boot status.""" TRY_PRINT.run( 'Secure Boot...', check_secure_boot_status, msg_good='Enabled', ) -def auto_show_storage_status(): +def auto_show_storage_status() -> None: """Display storage status.""" TRY_PRINT.run('Storage Status...', get_storage_status) @@ -639,20 +629,20 @@ def auto_shutup_10(): TRY_PRINT.run('Disable Telemetry...', run_shutup_10) -def auto_windows_temp_fix(): +def auto_windows_temp_fix() -> None: """Restore default ACLs for Windows\\Temp.""" TRY_PRINT.run(r'Windows\Temp fix...', fix_windows_temp) # Configure Functions -def config_explorer(): +def config_explorer() -> None: """Configure Windows Explorer and restart the process.""" reg_write_settings(REG_WINDOWS_EXPLORER) kill_procs('explorer.exe', force=True) popen_program(['explorer.exe']) -def config_open_shell(): +def config_open_shell() -> None: """Configure Open Shell.""" has_low_power_idle = False @@ -672,7 +662,7 @@ def config_open_shell(): reg_write_settings(REG_OPEN_SHELL_LOW_POWER_IDLE) -def disable_chrome_notifications(): +def disable_chrome_notifications() -> None: """Disable notifications in Google Chrome.""" defaults_key = 'default_content_setting_values' profiles = [] @@ -714,13 +704,13 @@ def disable_chrome_notifications(): pref_file.write_text(json.dumps(pref_data, separators=(',', ':'))) -def enable_bsod_minidumps(): +def enable_bsod_minidumps() -> None: """Enable saving minidumps during BSoDs.""" cmd = ['wmic', 'RECOVEROS', 'set', 'DebugInfoType', '=', '3'] run_program(cmd) -def enable_ublock_origin(): +def enable_ublock_origin() -> None: """Enable uBlock Origin in supported browsers.""" base_paths = [ PROGRAMFILES_64, PROGRAMFILES_32, os.environ.get('LOCALAPPDATA'), @@ -747,10 +737,10 @@ def enable_ublock_origin(): # Open detected browsers for cmd in cmds: - popen_program(cmd) + popen_program(cmd, pipe=True) -def fix_windows_temp(): +def fix_windows_temp() -> None: """Restore default permissions for Windows\\Temp.""" permissions = ( 'Users:(CI)(X,WD,AD)', @@ -776,7 +766,7 @@ def install_eset_nod32_av(msp=False): ) -def install_firefox(): +def install_firefox() -> None: """Install Firefox. As far as I can tell if you use the EXE installers then it will use @@ -881,7 +871,7 @@ def install_libreoffice( run_program(cmd) -def install_open_shell(): +def install_open_shell() -> None: """Install Open Shell (just the Start Menu).""" skin_zip = get_tool_path('OpenShell', 'Fluent-Metro', suffix='zip') @@ -911,7 +901,7 @@ def install_open_shell(): run_program(cmd) -def install_software_bundle(): +def install_software_bundle() -> None: """Install standard software bundle.""" download_tool('Ninite', 'Software Bundle') installer = get_tool_path('Ninite', 'Software Bundle') @@ -919,8 +909,8 @@ def install_software_bundle(): warning = 'NOTE: Press CTRL+c to manually resume if it gets stuck...' # Start installations and wait for them to finish - print_standard(msg) - print_warning(warning, end='', flush=True) + ui.print_standard(msg) + ui.print_warning(warning, end='', flush=True) proc = popen_program([installer]) try: proc.wait() @@ -936,7 +926,7 @@ def install_software_bundle(): end='', flush=True) -def install_vcredists(): +def install_vcredists() -> None: """Install latest supported Visual C++ runtimes.""" for year in (2012, 2013, 2022): cmd_args = ['/install', '/passive', '/norestart'] @@ -953,7 +943,7 @@ def install_vcredists(): run_program([installer, *cmd_args]) -def uninstall_firefox(): +def uninstall_firefox() -> None: """Uninstall all copies of Firefox.""" json_file = format_log_path(log_name='Installed Programs', timestamp=True) json_file = json_file.with_name(f'{json_file.stem}.json') @@ -979,13 +969,14 @@ def apply_its_settings(): create_custom_power_plan(enable_sleep=False) -def check_secure_boot_status(): +def check_secure_boot_status() -> None: """Check Secure Boot status.""" is_secure_boot_enabled(raise_exceptions=True, show_alert=True) -def get_firefox_default_profile(profiles_ini): +def get_firefox_default_profile(profiles_ini) -> Any: """Get Firefox default profile, returns(pathlib.Path, encoding) or None.""" + # TODO: Refactor to remove dependancy on Any default_profile = None encoding = None parser = None @@ -1022,24 +1013,24 @@ def get_firefox_default_profile(profiles_ini): return (default_profile, encoding) -def get_storage_status(): +def get_storage_status() -> list[str]: """Get storage status for fixed disks, returns list.""" report = get_volume_usage(use_colors=True) for disk in get_raw_disks(): - report.append(color_string(f'Uninitialized Disk: {disk}', 'RED')) + report.append(ansi.color_string(f'Uninitialized Disk: {disk}', 'RED')) # Done return report -def set_default_browser(): +def set_default_browser() -> None: """Open Windows Settings to the default apps section.""" cmd = ['start', '', 'ms-settings:defaultapps'] popen_program(cmd, shell=True) # Tool Functions -def export_aida64_report(): +def export_aida64_report() -> None: """Export AIDA64 report.""" report_path = format_log_path( log_name='AIDA64 System Report', @@ -1060,12 +1051,12 @@ def export_aida64_report(): raise GenericError('Error(s) encountered exporting report.') -def open_device_manager(): +def open_device_manager() -> None: """Open Device Manager.""" popen_program(['mmc', 'devmgmt.msc']) -def open_hwinfo_sensors(): +def open_hwinfo_sensors() -> None: """Open HWiNFO sensors.""" hwinfo_path = get_tool_path('HWiNFO', 'HWiNFO') base_config = hwinfo_path.with_name('general.ini') @@ -1081,7 +1072,7 @@ def open_hwinfo_sensors(): run_tool('HWiNFO', 'HWiNFO', popen=True) -def open_snappy_driver_installer_origin(): +def open_snappy_driver_installer_origin() -> None: """Open Snappy Driver Installer Origin.""" run_tool('SDIO', 'SDIO', cwd=True, pipe=True, popen=True) @@ -1091,17 +1082,17 @@ def open_webcam_tests(): webbrowser.open('http://webcamtests.com/') -def open_windows_activation(): +def open_windows_activation() -> None: """Open Windows Activation.""" popen_program(['slui']) -def open_windows_updates(): +def open_windows_updates() -> None: """Open Windows Updates.""" popen_program(['control', '/name', 'Microsoft.WindowsUpdate']) -def open_xmplay(): +def open_xmplay() -> None: """Open XMPlay.""" sleep(2) run_tool('XMPlay', 'XMPlay', 'music.7z', cwd=True, popen=True) diff --git a/scripts/wk/std.py b/scripts/wk/std.py index 65fd69f0..92c50b79 100644 --- a/scripts/wk/std.py +++ b/scripts/wk/std.py @@ -1,55 +1,13 @@ """WizardKit: Standard Functions""" # vim: sts=2 sw=2 ts=2 -import inspect -import itertools import logging -import lzma -import os -import pathlib -import pickle import platform import re -import socket -import subprocess -import sys import time -import traceback - -from collections import OrderedDict - -try: - from functools import cache -except ImportError: - # Assuming Python is < 3.9 - from functools import lru_cache as cache - -import requests - -from wk.cfg.main import ( - ENABLED_UPLOAD_DATA, - INDENT, - SUPPORT_MESSAGE, - WIDTH, - ) -from wk.cfg.net import CRASH_SERVER -from wk.log import get_root_logger_path # STATIC VARIABLES -COLORS = { - 'CLEAR': '\033[0m', - 'RED': '\033[31m', - 'RED_BLINK': '\033[31;5m', - 'ORANGE': '\033[31;1m', - 'ORANGE_RED': '\033[1;31;41m', - 'GREEN': '\033[32m', - 'YELLOW': '\033[33m', - 'YELLOW_BLINK': '\033[33;5m', - 'BLUE': '\033[34m', - 'PURPLE': '\033[35m', - 'CYAN': '\033[36m', - } LOG = logging.getLogger(__name__) PLATFORM = platform.system() REGEX_SIZE_STRING = re.compile( @@ -72,575 +30,11 @@ class GenericWarning(Exception): """ -# Classes -class Menu(): - """Object for tracking menu specific data and methods. - - Menu items are added to an OrderedDict so the order is preserved. - - ASSUMPTIONS: - 1. All entry names are unique. - 2. All action entry names start with different letters. - """ - def __init__(self, title='[Untitled Menu]'): - self.actions = OrderedDict() - self.options = OrderedDict() - self.sets = OrderedDict() - self.toggles = OrderedDict() - self.disabled_str = 'Disabled' - self.separator = '─' - self.title = title - - def _generate_menu_text(self): - """Generate menu text, returns str.""" - separator_string = self._get_separator_string() - menu_lines = [self.title, separator_string] if self.title else [] - - # Sets & toggles - for section in (self.sets, self.toggles): - for details in section.values(): - if details.get('Hidden', False): - continue - if details.get('Separator', False): - menu_lines.append(separator_string) - menu_lines.append(details['Display Name']) - if self.sets or self.toggles: - menu_lines.append(separator_string) - - # Options - for details in self.options.values(): - if details.get('Hidden', False): - continue - if details.get('Separator', False): - menu_lines.append(separator_string) - menu_lines.append(details['Display Name']) - if self.options: - menu_lines.append(separator_string) - - # Actions - for details in self.actions.values(): - if details.get('Hidden', False): - continue - if details.get('Separator', False): - menu_lines.append(separator_string) - menu_lines.append(details['Display Name']) - - # Show menu - menu_lines.append('') - menu_lines = [str(line) for line in menu_lines] - return '\n'.join(menu_lines) - - def _get_display_name( - self, name, details, - index=None, no_checkboxes=True, setting_item=False): - """Format display name based on details and args, returns str.""" - disabled = details.get('Disabled', False) - if setting_item and not details['Selected']: - # Display item in YELLOW - disabled = True - checkmark = '*' - if 'DISPLAY' in os.environ or PLATFORM == 'Darwin': - checkmark = '✓' - display_name = f'{index if index else name[:1].upper()}: ' - if not (index and index >= 10): - display_name = f' {display_name}' - if setting_item and 'Value' in details: - name = f'{name} = {details["Value"]}' - - # Add enabled status if necessary - if not no_checkboxes: - display_name += f'[{checkmark if details["Selected"] else " "}] ' - - # Add name - if disabled: - display_name += color_string(f'{name} ({self.disabled_str})', 'YELLOW') - else: - display_name += name - - # Done - return display_name - - def _get_separator_string(self): - """Format separator length based on name lengths, returns str.""" - separator_length = 0 - - # Check title line(s) - if self.title: - for line in self.title.split('\n'): - separator_length = max(separator_length, len(strip_colors(line))) - - # Loop over all item names - for section in (self.actions, self.options, self.sets, self.toggles): - for details in section.values(): - if details.get('Hidden', False): - # Skip hidden lines - continue - line = strip_colors(details['Display Name']) - separator_length = max(separator_length, len(line)) - separator_length += 1 - - # Done - return self.separator * separator_length - - def _get_valid_answers(self): - """Get valid answers based on menu items, returns list.""" - valid_answers = [] - - # Numbered items - index = 0 - for section in (self.sets, self.toggles, self.options): - for details in section.values(): - if details.get('Hidden', False): - # Don't increment index or add to valid_answers - continue - index += 1 - if not details.get('Disabled', False): - valid_answers.append(str(index)) - - # Action items - for name, details in self.actions.items(): - if not details.get('Disabled', False): - valid_answers.append(name[:1].upper()) - - # Done - return valid_answers - - def _resolve_selection(self, selection): - """Get menu item based on user selection, returns tuple.""" - offset = 1 - resolved_selection = None - if selection.isnumeric(): - # Enumerate over numbered entries - entries = [ - *self.sets.items(), - *self.toggles.items(), - *self.options.items(), - ] - for _i, details in enumerate(entries): - if details[1].get('Hidden', False): - offset -= 1 - elif str(_i+offset) == selection: - resolved_selection = (details) - break - else: - # Just check actions - for action, details in self.actions.items(): - if action.lower().startswith(selection.lower()): - resolved_selection = (action, details) - break - - # 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: - # Update targets not the set itself - new_status = not self.sets[entry]['Selected'] if toggle else status - targets = self.sets[entry]['Targets'] - self._update_set_selection_status(targets, new_status) - for section in (self.toggles, self.options, self.actions): - if entry in section: - if toggle: - section[entry]['Selected'] = not section[entry]['Selected'] - else: - section[entry]['Selected'] = status - - def _update_set_selection_status(self, targets, status): - """Select or deselect options based on targets and status.""" - for option, details in self.options.items(): - # If (new) status is True and this option is a target then select - # Otherwise deselect - details['Selected'] = status and option in targets - - def _user_select(self, prompt): - """Show menu and select an entry, returns str.""" - menu_text = self._generate_menu_text() - valid_answers = self._get_valid_answers() - - # Menu loop - while True: - clear_screen() - print(menu_text) - sleep(0.01) - answer = input_text(prompt).strip() - if answer.upper() in valid_answers: - break - - # Done - return answer - - def add_action(self, name, details=None): - """Add action to menu.""" - details = details if details else {} - details['Selected'] = details.get('Selected', False) - self.actions[name] = details - - def add_option(self, name, details=None): - """Add option to menu.""" - details = details if details else {} - details['Selected'] = details.get('Selected', False) - self.options[name] = details - - def add_set(self, name, details=None): - """Add set to menu.""" - details = details if details else {} - details['Selected'] = details.get('Selected', False) - - # Safety check - if 'Targets' not in details: - raise KeyError('Menu set has no targets') - - # Add set - self.sets[name] = details - - def add_toggle(self, name, details=None): - """Add toggle to menu.""" - details = details if details else {} - details['Selected'] = details.get('Selected', False) - self.toggles[name] = details - - def advanced_select(self, prompt='Please make a selection: '): - """Display menu and make multiple selections, returns tuple. - - NOTE: Menu is displayed until an action entry is selected. - """ - while True: - self._update(single_selection=False) - user_selection = self._user_select(prompt) - selected_entry = self._resolve_selection(user_selection) - if user_selection.isnumeric(): - # Update selection(s) - self._update_entry_selection_status(selected_entry[0]) - else: - # Action selected - break - - # Done - return selected_entry - - def settings_select(self, prompt='Please make a selection: '): - """Display menu and make multiple selections, returns tuple. - - NOTE: Menu is displayed until an action entry is selected. - """ - choice_kwargs = { - 'choices': ['T', 'C'], - 'prompt': 'Toggle or change value?', - } - - while 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(): - if 'Value' in selected_entry[-1] and choice(**choice_kwargs) == 'C': - # Change - selected_entry[-1]['Value'] = input_text('Enter new value: ') - else: - # Toggle - self._update_entry_selection_status(selected_entry[0]) - else: - # Action selected - break - - # Done - return selected_entry - - def simple_select(self, prompt='Please make a selection: ', update=True): - """Display menu and make a single selection, returns tuple.""" - if update: - self._update() - user_selection = self._user_select(prompt) - return self._resolve_selection(user_selection) - - def update(self): - """Update menu with default settings.""" - self._update() - - -class TryAndPrint(): - """Object used to standardize running functions and returning the result. - - The errors and warning attributes are used to allow fine-tuned results - based on exception names. - """ - def __init__(self, msg_bad='FAILED', msg_good='SUCCESS'): - self.catch_all = True - self.indent = INDENT - self.list_errors = ['GenericError'] - self.list_warnings = ['GenericWarning'] - self.msg_bad = msg_bad - self.msg_good = msg_good - self.verbose = False - self.width = WIDTH - - def _format_exception_message(self, _exception): - """Format using the exception's args or name, returns str.""" - LOG.debug( - 'Formatting exception: %s, %s', - _exception.__class__.__name__, - _exception, - ) - message = '' - - # Format message string from _exception - try: - if isinstance(_exception, subprocess.CalledProcessError): - message = _exception.stderr - if not isinstance(message, str): - message = message.decode('utf-8') - message = message.strip() - elif isinstance(_exception, ZeroDivisionError): - # Skip and just use exception name below - pass - else: - message = str(_exception) - except Exception: - # Just use the exception name instead - pass - - # Prepend exception name - if _exception.__class__.__name__ not in ('GenericError', 'GenericWarning'): - try: - message = f'{_exception.__class__.__name__}: {message}' - except Exception: - message = f'UNKNOWN ERROR: {message}' - - # Fix multi-line messages - if '\n' in message: - try: - lines = [ - f'{" "*(self.indent+self.width)}{line.strip()}' - for line in message.splitlines() if line.strip() - ] - lines[0] = lines[0].strip() - message = '\n'.join(lines) - except Exception: - pass - - # Done - return message - - def _format_function_output(self, output, msg_good): - """Format function output for use in try_and_print(), returns str.""" - LOG.debug('Formatting output: %s', output) - - if not output: - raise GenericWarning('No output') - - # Ensure we're working with a list - if isinstance(output, subprocess.CompletedProcess): - stdout = output.stdout - if not isinstance(stdout, str): - stdout = stdout.decode('utf8') - output = stdout.strip().splitlines() - if not output: - # Going to treat these as successes (for now) - LOG.warning('Program output was empty, assuming good result.') - return color_string(msg_good, 'GREEN') - else: - try: - output = list(output) - except TypeError: - output = [output] - - # Safety check - if not output: - # Going to ignore empty function output for now - LOG.error('Output is empty') - return 'UNKNOWN' - - # Build result_msg - result_msg = f'{output.pop(0)}' - if output: - output = [f'{" "*(self.indent+self.width)}{line}' for line in output] - result_msg += '\n' + '\n'.join(output) - - # Done - return result_msg - - def _log_result(self, message, result_msg): - """Log result text without color formatting.""" - log_text = f'{" "*self.indent}{message:<{self.width}}{result_msg}' - for line in log_text.splitlines(): - line = strip_colors(line) - LOG.info(line) - - def add_error(self, exception_name): - """Add exception name to error list.""" - if exception_name not in self.list_errors: - self.list_errors.append(exception_name) - - def add_warning(self, exception_name): - """Add exception name to warning list.""" - if exception_name not in self.list_warnings: - self.list_warnings.append(exception_name) - - def run( - self, message, function, *args, - catch_all=None, msg_good=None, verbose=None, **kwargs): - """Run a function and print the results, returns results as dict. - - If catch_all is True then (nearly) all exceptions will be caught. - Otherwise if an exception occurs that wasn't specified it will be - re-raised. - - If the function returns data it will be used instead of msg_good, - msg_bad, or exception text. - The output should be a list or a subprocess.CompletedProcess object. - - If msg_good is passed it will override self.msg_good for this call. - - If verbose is True then exception names or messages will be used for - the result message. Otherwise it will simply be set to result_bad. - - If catch_all and/or verbose are passed it will override - self.catch_all and/or self.verbose for this call. - - args and kwargs are passed to the function. - """ - LOG.debug('function: %s.%s', function.__module__, function.__name__) - LOG.debug('args: %s', args) - LOG.debug('kwargs: %s', kwargs) - LOG.debug( - 'catch_all: %s, msg_good: %s, verbose: %s', - catch_all, - msg_good, - verbose, - ) - f_exception = None - catch_all = catch_all if catch_all is not None else self.catch_all - msg_good = msg_good if msg_good is not None else self.msg_good - output = None - result_msg = 'UNKNOWN' - verbose = verbose if verbose is not None else self.verbose - - # Build exception tuples - e_exceptions = tuple(get_exception(e) for e in self.list_errors) - w_exceptions = tuple(get_exception(e) for e in self.list_warnings) - - # Run function and catch exceptions - print(f'{" "*self.indent}{message:<{self.width}}', end='', flush=True) - LOG.debug('Running function: %s.%s', function.__module__, function.__name__) - try: - output = function(*args, **kwargs) - except w_exceptions as _exception: - # Warnings - result_msg = self._format_exception_message(_exception) - print_warning(result_msg, log=False) - f_exception = _exception - except e_exceptions as _exception: - # Exceptions - result_msg = self._format_exception_message(_exception) - print_error(result_msg, log=False) - f_exception = _exception - except Exception as _exception: - # Unexpected exceptions - if verbose: - result_msg = self._format_exception_message(_exception) - else: - result_msg = self.msg_bad - print_error(result_msg, log=False) - f_exception = _exception - if not catch_all: - # Re-raise error as necessary - raise - else: - # Success - if output: - result_msg = self._format_function_output(output, msg_good) - print(result_msg) - else: - result_msg = msg_good - print_success(result_msg, log=False) - - # Done - self._log_result(message, result_msg) - return { - 'Exception': f_exception, - 'Failed': bool(f_exception), - 'Message': result_msg, - 'Output': output, - } - - # Functions -def abort(prompt='Aborted.', show_prompt=True, return_code=1): - """Abort script.""" - print_warning(prompt) - if show_prompt: - sleep(0.5) - pause(prompt='Press Enter to exit... ') - sys.exit(return_code) - - -def ask(prompt='Kotaero!'): - """Prompt the user with a Y/N question, returns bool.""" - answer = None - prompt = f'{prompt} [Y/N]: ' - - # Loop until acceptable answer is given - while answer is None: - tmp = input_text(prompt) - if re.search(r'^y(es|up|)$', tmp, re.IGNORECASE): - answer = True - elif re.search(r'^n(o|ope|)$', tmp, re.IGNORECASE): - answer = False - - # Done - LOG.info('%s%s', prompt, 'Yes' if answer else 'No') - return answer - - -def beep(repeat=1): - """Play system bell with optional repeat.""" - while repeat >= 1: - # Print bell char without a newline - print('\a', end='', flush=True) - sleep(0.5) - repeat -= 1 - - -def bytes_to_string(size, decimals=0, use_binary=True): +def bytes_to_string( + size: float | int, + decimals: int = 0, + use_binary: bool = True) -> str: """Convert size into a human-readable format, returns str. [Doctest] @@ -682,330 +76,13 @@ def bytes_to_string(size, decimals=0, use_binary=True): return size_str -def choice(choices, prompt='答えろ！'): - """Choose an option from a provided list, returns str. - - Choices provided will be converted to uppercase and returned as such. - Similar to the commands choice (Windows) and select (Linux). - """ - LOG.debug('choices: %s, prompt: %s', choices, prompt) - answer = None - choices = [str(c).upper()[:1] for c in choices] - prompt = f'{prompt} [{"/".join(choices)}]' - regex = f'^({"|".join(choices)})$' - - # Loop until acceptable answer is given - while answer is None: - tmp = input_text(prompt=prompt) - if re.search(regex, tmp, re.IGNORECASE): - answer = tmp.upper() - - # Done - LOG.info('%s %s', prompt, answer) - return answer - - -def clear_screen(): - """Simple wrapper for clear/cls.""" - cmd = 'cls' if os.name == 'nt' else 'clear' - proc = subprocess.run(cmd, check=False, shell=True, stderr=subprocess.PIPE) - - # Workaround for live macOS env - if proc.returncode != 0: - print('\033c') - - -def color_string(strings, colors, sep=' '): - """Build colored string using ANSI escapes, returns str.""" - clear_code = COLORS['CLEAR'] - msg = [] - - # Convert to tuples if necessary - if isinstance(strings, (str, pathlib.Path)): - strings = (strings,) - if isinstance(colors, (str, pathlib.Path)): - colors = (colors,) - - # Convert to strings if necessary - try: - iter(strings) - except TypeError: - # Assuming single element passed, convert to string - strings = (str(strings),) - try: - iter(colors) - except TypeError: - # Assuming single element passed, convert to string - colors = (str(colors),) - - # Build new string with color escapes added - for string, color in itertools.zip_longest(strings, colors): - color_code = COLORS.get(color, clear_code) - msg.append(f'{color_code}{string}{clear_code}') - - # Done - return sep.join(msg) - - -def generate_debug_report(): - """Generate debug report, returns str.""" - platform_function_list = ( - 'architecture', - 'machine', - 'platform', - 'python_version', - ) - report = [] - - # Logging data - log_path = get_log_filepath() - if log_path: - report.append('------ Start Log -------') - report.append('') - with open(log_path, 'r', encoding='utf-8') as log_file: - report.extend(log_file.read().splitlines()) - report.append('') - report.append('------- End Log --------') - - # System - report.append('--- Start debug info ---') - report.append('') - report.append('[System]') - report.append(f' {"FQDN":<24} {socket.getfqdn()}') - for func in platform_function_list: - func_name = func.replace('_', ' ').capitalize() - func_result = getattr(platform, func)() - report.append(f' {func_name:<24} {func_result}') - report.append(f' {"Python sys.argv":<24} {sys.argv}') - report.append('') - - # Environment - report.append('[Environment Variables]') - for key, value in sorted(os.environ.items()): - report.append(f' {key:<24} {value}') - report.append('') - - # Done - report.append('---- End debug info ----') - return '\n'.join(report) - - -@cache -def get_exception(name): - """Get exception by name, returns exception object. - - [Doctest] - >>> t = TryAndPrint() - >>> t._get_exception('AttributeError') - - >>> t._get_exception('CalledProcessError') - - >>> t._get_exception('GenericError') - - """ - LOG.debug('Getting exception: %s', name) - obj = getattr(sys.modules[__name__], name, None) - if obj: - return obj - - # Try builtin classes - obj = getattr(sys.modules['builtins'], name, None) - if obj: - return obj - - # Try all modules - for _mod in sys.modules.values(): - obj = getattr(_mod, name, None) - if obj: - break - - # Check if not found - if not obj: - raise AttributeError(f'Failed to find exception: {name}') - - # Done - return obj - - -def get_log_filepath(): - """Get the log filepath from the root logger, returns pathlib.Path obj. - - NOTE: This will use the first handler baseFilename it finds (if any). - """ - log_filepath = None - root_logger = logging.getLogger() - - # Check handlers - for handler in root_logger.handlers: - if hasattr(handler, 'baseFilename'): - log_filepath = pathlib.Path(handler.baseFilename).resolve() - break - - # Done - return log_filepath - - -def input_text(prompt='Enter text', allow_empty_response=True): - """Get text from user, returns string.""" - prompt = str(prompt) - response = None - if prompt[-1:] != ' ': - prompt += ' ' - print(prompt, end='', flush=True) - - while response is None: - try: - response = input() - LOG.debug('%s%s', prompt, response) - except EOFError: - # Ignore and try again - LOG.warning('Exception occured', exc_info=True) - print('', flush=True) - if not allow_empty_response: - if response is None or not response.strip(): - # The None check here is used to avoid a TypeError if response is None - print(f'\r{prompt}', end='', flush=True) - response = None - - return response - - -def major_exception(): - """Display traceback, optionally upload detailes, and exit.""" - LOG.critical('Major exception encountered', exc_info=True) - print_error('Major exception', log=False) - print_warning(SUPPORT_MESSAGE) - if ENABLED_UPLOAD_DATA: - print_warning('Also, please run upload-logs to help debugging!') - print(traceback.format_exc()) - - # Done - pause('Press Enter to exit... ') - raise SystemExit(1) - - -def pause(prompt='Press Enter to continue... '): - """Simple pause implementation.""" - input_text(prompt) - - -def print_colored(strings, colors, log=False, sep=' ', **kwargs): - """Prints strings in the colors specified.""" - LOG.debug( - 'strings: %s, colors: %s, sep: %s, kwargs: %s', - strings, colors, sep, kwargs, - ) - msg = color_string(strings, colors, sep=sep) - print_options = { - 'end': kwargs.get('end', '\n'), - 'file': kwargs.get('file', sys.stdout), - 'flush': kwargs.get('flush', False), - } - - print(msg, **print_options) - if log: - LOG.info(strip_colors(msg)) - - -def print_error(msg, log=True, **kwargs): - """Prints message in RED and log as ERROR.""" - if 'file' not in kwargs: - # Only set if not specified - kwargs['file'] = sys.stderr - print_colored(msg, 'RED', **kwargs) - if log: - LOG.error(msg) - - -def print_info(msg, log=True, **kwargs): - """Prints message in BLUE and log as INFO.""" - print_colored(msg, 'BLUE', **kwargs) - if log: - LOG.info(msg) - - -def print_report(report, indent=None, log=True): - """Print report to screen and optionally to log.""" - for line in report: - if indent: - line = f'{" "*indent}{line}' - print(line) - if log: - LOG.info(strip_colors(line)) - - -def print_standard(msg, log=True, **kwargs): - """Prints message and log as INFO.""" - print(msg, **kwargs) - if log: - LOG.info(msg) - - -def print_success(msg, log=True, **kwargs): - """Prints message in GREEN and log as INFO.""" - print_colored(msg, 'GREEN', **kwargs) - if log: - LOG.info(msg) - - -def print_warning(msg, log=True, **kwargs): - """Prints message in YELLOW and log as WARNING.""" - if 'file' not in kwargs: - # Only set if not specified - kwargs['file'] = sys.stderr - print_colored(msg, 'YELLOW', **kwargs) - if log: - LOG.warning(msg) - - -def save_pickles(obj_dict, out_path=None): - """Save dict of objects using pickle.""" - LOG.info('Saving pickles') - - # Set path - if not out_path: - out_path = pathlib.Path(f'{get_root_logger_path().parent}/debug') - - # Save pickles - try: - for name, obj in obj_dict.copy().items(): - if name.startswith('__') or inspect.ismodule(obj): - continue - with open(f'{out_path}/{name}.pickle', 'wb') as _f: - pickle.dump(obj, _f, protocol=pickle.HIGHEST_PROTOCOL) - except Exception: - LOG.error('Failed to save all the pickles', exc_info=True) - - -def set_title(title): - """Set window title.""" - LOG.debug('title: %s', title) - if os.name == 'nt': - os.system(f'title {title}') - else: - print_error('Setting the title is only supported under Windows.') - - -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), - colors, - log=True, - sep='', - ) - - -def sleep(seconds=2): +def sleep(seconds: int | float = 2) -> None: """Simple wrapper for time.sleep.""" time.sleep(seconds) -def string_to_bytes(size, assume_binary=False): - """Convert human-readable size str to bytes and return an int.""" +def string_to_bytes(size: float | int | str, assume_binary: bool = False) -> int: + """Convert human-readable size to bytes and return an int.""" LOG.debug('size: %s, assume_binary: %s', size, assume_binary) scale = 1000 size = str(size) @@ -1053,55 +130,5 @@ def string_to_bytes(size, assume_binary=False): return size -def strip_colors(string): - """Strip known ANSI color escapes from string, returns str.""" - LOG.debug('string: %s', string) - for color in COLORS.values(): - string = string.replace(color, '') - return string - - -def upload_debug_report(report, compress=True, reason='DEBUG'): - """Upload debug report to CRASH_SERVER as specified in wk.cfg.main.""" - LOG.info('Uploading debug report to %s', CRASH_SERVER.get('Name', '?')) - headers = CRASH_SERVER.get('Headers', {'X-Requested-With': 'XMLHttpRequest'}) - if compress: - headers['Content-Type'] = 'application/octet-stream' - - # Check if the required server details are available - if not all(CRASH_SERVER.get(key, False) for key in ('Name', 'Url', 'User')): - msg = 'Server details missing, aborting upload.' - print_error(msg) - raise RuntimeError(msg) - - # Set filename (based on the logging config if possible) - filename = 'Unknown' - log_path = get_log_filepath() - if log_path: - # Strip everything but the prefix - filename = re.sub(r'^(.*)_(\d{4}-\d{2}-\d{2}.*)', r'\1', log_path.name) - filename = f'{filename}_{reason}_{time.strftime("%Y-%m-%d_%H%M%S%z")}.log' - LOG.debug('filename: %s', filename) - - # Compress report - if compress: - filename += '.xz' - xz_report = lzma.compress(report.encode('utf8')) - - # Upload report - url = f'{CRASH_SERVER["Url"]}/{filename}' - response = requests.put( - url, - auth=(CRASH_SERVER['User'], CRASH_SERVER.get('Pass', '')), - data=xz_report if compress else report, - headers=headers, - timeout=60, - ) - - # Check response - if not response.ok: - raise RuntimeError('Failed to upload report') - - if __name__ == '__main__': print("This file is not meant to be called directly.") diff --git a/scripts/wk/ui/__init__.py b/scripts/wk/ui/__init__.py new file mode 100644 index 00000000..75a335fc --- /dev/null +++ b/scripts/wk/ui/__init__.py @@ -0,0 +1,6 @@ +"""WizardKit: ui module init""" + +from . import ansi +from . import cli +from . import tmux +from . import tui diff --git a/scripts/wk/ui/ansi.py b/scripts/wk/ui/ansi.py new file mode 100644 index 00000000..b65fe2c9 --- /dev/null +++ b/scripts/wk/ui/ansi.py @@ -0,0 +1,70 @@ +"""WizardKit: ANSI control/escape functions""" +# vim: sts=2 sw=2 ts=2 + +import itertools +import logging + +from typing import Iterable + +# STATIC VARIABLES +LOG = logging.getLogger(__name__) +COLORS = { + 'CLEAR': '\033[0m', + 'RED': '\033[31m', + 'RED_BLINK': '\033[31;5m', + 'ORANGE': '\033[31;1m', + 'ORANGE_RED': '\033[1;31;41m', + 'GREEN': '\033[32m', + 'YELLOW': '\033[33m', + 'YELLOW_BLINK': '\033[33;5m', + 'BLUE': '\033[34m', + 'PURPLE': '\033[35m', + 'CYAN': '\033[36m', + } + + +# Functions +def clear_screen() -> None: + """Clear screen using ANSI escape.""" + print('\033c', end='', flush=True) + + +def color_string( + strings: Iterable[str] | str, + colors: Iterable[str | None] | str, + sep=' ', + ) -> str: + """Build colored string using ANSI escapes, returns str.""" + data = {'strings': strings, 'colors': colors} + msg = [] + + # Convert input to tuples of strings + for k, v in data.items(): + if isinstance(v, str): + # Avoid splitting string into a list of characters + data[k] = (v,) + try: + iter(v) + except TypeError: + # Assuming single element passed, convert to string + data[k] = (str(v),) + + # Build new string with color escapes added + for string, color in itertools.zip_longest(data['strings'], data['colors']): + color_code = COLORS.get(str(color), COLORS['CLEAR']) + msg.append(f'{color_code}{string}{COLORS["CLEAR"]}') + + # Done + return sep.join(msg) + + +def strip_colors(string: str) -> str: + """Strip known ANSI color escapes from string, returns str.""" + LOG.debug('string: %s', string) + for color in COLORS.values(): + string = string.replace(color, '') + return string + + +if __name__ == '__main__': + print("This file is not meant to be called directly.") diff --git a/scripts/wk/ui/cli.py b/scripts/wk/ui/cli.py new file mode 100644 index 00000000..922d204a --- /dev/null +++ b/scripts/wk/ui/cli.py @@ -0,0 +1,924 @@ +"""WizardKit: CLI functions""" +# vim: sts=2 sw=2 ts=2 + +import logging +import os +import platform +import re +import subprocess +import sys +import traceback + +from typing import Any, Callable, Iterable + +from prompt_toolkit import prompt +from prompt_toolkit.document import Document +from prompt_toolkit.validation import Validator, ValidationError + +try: + from functools import cache +except ImportError: + # Assuming Python is < 3.9 + from functools import lru_cache as cache + +from wk.cfg.main import ( + ENABLED_UPLOAD_DATA, + INDENT, + SUPPORT_MESSAGE, + WIDTH, + ) +from wk.std import (sleep, GenericWarning) +from wk.ui.ansi import clear_screen, color_string, strip_colors + +# STATIC VARIABLES +LOG = logging.getLogger(__name__) +PLATFORM = platform.system() + + +# Classes +class InputChoiceValidator(Validator): + """Validate that input is one of the provided choices.""" + def __init__(self, choices: Iterable[str], allow_empty: bool = False): + self.allow_empty: bool = allow_empty + self.choices: list[str] = [str(c).upper() for c in choices] + super().__init__() + + def validate(self, document: Document) -> None: + text = document.text + if not (text or self.allow_empty): + raise ValidationError( + message='This input is required!', + cursor_position=len(text), + ) + if text and text.upper() not in self.choices: + raise ValidationError( + message='Invalid selection', + cursor_position=len(text), + ) + +class InputNotEmptyValidator(Validator): + """Validate that input is not empty.""" + def validate(self, document: Document) -> None: + text = document.text + if not text: + raise ValidationError( + message='This input is required!', + cursor_position=len(text), + ) + +class InputTicketIDValidator(Validator): + """Validate that input resembles a ticket ID.""" + def __init__(self, allow_empty: bool = False): + self.allow_empty: bool = allow_empty + super().__init__() + + def validate(self, document: Document) -> None: + text = document.text + if not (text or self.allow_empty): + raise ValidationError( + message='This input is required!', + cursor_position=len(text), + ) + if text and not re.match(r'^\d', text): + raise ValidationError( + message='Ticket ID should start with a number!', + cursor_position=len(text), + ) + +class InputYesNoValidator(Validator): + """Validate that input is a yes or no.""" + def __init__(self, allow_empty: bool = False): + self.allow_empty: bool = allow_empty + super().__init__() + + def validate(self, document: Document) -> None: + text = document.text + if not (text or self.allow_empty): + raise ValidationError( + message='This input is required!', + cursor_position=len(text), + ) + if text and not re.match(r'^(y(es|up|)|n(o|ope|))$', text, re.IGNORECASE): + raise ValidationError( + message='Please answer "yes" or "no"', + cursor_position=len(text), + ) + +class Menu(): + """Object for tracking menu specific data and methods. + + ASSUMPTIONS: + 1. All entry names are unique. + 2. All action entry names start with different letters. + """ + def __init__(self, title: str = '[Untitled Menu]'): + self.actions: dict[str, dict[Any, Any]] = {} + self.options: dict[str, dict[Any, Any]] = {} + self.sets: dict[str, dict[Any, Any]] = {} + self.toggles: dict[str, dict[Any, Any]] = {} + self.disabled_str: str = 'Disabled' + self.separator: str = '─' + self.title: str = title + + def _generate_menu_text(self) -> str: + """Generate menu text, returns str.""" + separator_string = self._get_separator_string() + menu_lines = [self.title, separator_string] if self.title else [] + + # Sets & toggles + for section in (self.sets, self.toggles): + for details in section.values(): + if details.get('Hidden', False): + continue + if details.get('Separator', False): + menu_lines.append(separator_string) + menu_lines.append(details['Display Name']) + if self.sets or self.toggles: + menu_lines.append(separator_string) + + # Options + for details in self.options.values(): + if details.get('Hidden', False): + continue + if details.get('Separator', False): + menu_lines.append(separator_string) + menu_lines.append(details['Display Name']) + if self.options: + menu_lines.append(separator_string) + + # Actions + for details in self.actions.values(): + if details.get('Hidden', False): + continue + if details.get('Separator', False): + menu_lines.append(separator_string) + menu_lines.append(details['Display Name']) + + # Show menu + menu_lines.append('') + menu_lines = [str(line) for line in menu_lines] + return '\n'.join(menu_lines) + + def _get_display_name( + self, name, details, + index=None, no_checkboxes=True, setting_item=False) -> str: + """Format display name based on details and args, returns str.""" + disabled = details.get('Disabled', False) + if setting_item and not details['Selected']: + # Display item in YELLOW + disabled = True + checkmark = '*' + if 'DISPLAY' in os.environ or PLATFORM == 'Darwin': + checkmark = '✓' + display_name = f'{index if index else name[:1].upper()}: ' + if not (index and index >= 10): + display_name = f' {display_name}' + if setting_item and 'Value' in details: + name = f'{name} = {details["Value"]}' + + # Add enabled status if necessary + if not no_checkboxes: + display_name += f'[{checkmark if details["Selected"] else " "}] ' + + # Add name + if disabled: + display_name += color_string(f'{name} ({self.disabled_str})', 'YELLOW') + else: + display_name += name + + # Done + return display_name + + def _get_separator_string(self) -> str: + """Format separator length based on name lengths, returns str.""" + separator_length = 0 + + # Check title line(s) + if self.title: + for line in self.title.split('\n'): + separator_length = max(separator_length, len(strip_colors(line))) + + # Loop over all item names + for section in (self.actions, self.options, self.sets, self.toggles): + for details in section.values(): + if details.get('Hidden', False): + # Skip hidden lines + continue + line = strip_colors(details['Display Name']) + separator_length = max(separator_length, len(line)) + separator_length += 1 + + # Done + return self.separator * separator_length + + def _get_valid_answers(self) -> list[str]: + """Get valid answers based on menu items, returns list.""" + valid_answers = [] + + # Numbered items + index = 0 + for section in (self.sets, self.toggles, self.options): + for details in section.values(): + if details.get('Hidden', False): + # Don't increment index or add to valid_answers + continue + index += 1 + if not details.get('Disabled', False): + valid_answers.append(str(index)) + + # Action items + for name, details in self.actions.items(): + if not details.get('Disabled', False): + valid_answers.append(name[:1].upper()) + + # Done + return valid_answers + + def _resolve_selection(self, selection: str) -> tuple[str, dict[Any, Any]]: + """Get menu item based on user selection, returns tuple.""" + offset = 1 + resolved_selection = tuple() + if selection.isnumeric(): + # Enumerate over numbered entries + entries = [ + *self.sets.items(), + *self.toggles.items(), + *self.options.items(), + ] + for _i, details in enumerate(entries): + if details[1].get('Hidden', False): + offset -= 1 + elif str(_i+offset) == selection: + # TODO: Fix this typo! + # It was discovered after being in production for SEVERAL YEARS! + # Extra testing is needed to verify any calls to this function still + # depend on this functionality + resolved_selection = (details) + break + else: + # Just check actions + for action, details in self.actions.items(): + if action.lower().startswith(selection.lower()): + resolved_selection = (action, details) + break + + # Done + return resolved_selection + + def _update(self, single_selection: bool = True, settings_mode: bool = False) -> None: + """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: str, toggle: bool = True, status: bool = False) -> None: + """Update entry selection status either directly or by toggling.""" + if entry in self.sets: + # Update targets not the set itself + new_status = not self.sets[entry]['Selected'] if toggle else status + targets = self.sets[entry]['Targets'] + self._update_set_selection_status(targets, new_status) + for section in (self.toggles, self.options, self.actions): + if entry in section: + if toggle: + section[entry]['Selected'] = not section[entry]['Selected'] + else: + section[entry]['Selected'] = status + + def _update_set_selection_status(self, targets: Iterable[str], status: bool) -> None: + """Select or deselect options based on targets and status.""" + for option, details in self.options.items(): + # If (new) status is True and this option is a target then select + # Otherwise deselect + details['Selected'] = status and option in targets + + def _user_select(self, prompt_msg: str) -> str: + """Show menu and select an entry, returns str.""" + menu_text = self._generate_menu_text() + valid_answers = self._get_valid_answers() + + # Menu loop + while True: + clear_screen() + print(menu_text) + sleep(0.01) + answer = input_text(prompt_msg).strip() + if answer.upper() in valid_answers: + break + + # Done + return answer + + def add_action(self, name: str, details: dict[Any, Any] | None = None) -> None: + """Add action to menu.""" + details = details if details else {} + details['Selected'] = details.get('Selected', False) + self.actions[name] = details + + def add_option(self, name: str, details: dict[Any, Any] | None = None) -> None: + """Add option to menu.""" + details = details if details else {} + details['Selected'] = details.get('Selected', False) + self.options[name] = details + + def add_set(self, name: str, details: dict[Any, Any] | None = None) -> None: + """Add set to menu.""" + details = details if details else {} + details['Selected'] = details.get('Selected', False) + + # Safety check + if 'Targets' not in details: + raise KeyError('Menu set has no targets') + + # Add set + self.sets[name] = details + + def add_toggle(self, name: str, details: dict[Any, Any] | None = None) -> None: + """Add toggle to menu.""" + details = details if details else {} + details['Selected'] = details.get('Selected', False) + self.toggles[name] = details + + def advanced_select( + self, + prompt_msg: str = 'Please make a selection: ', + ) -> tuple[str, dict[Any, Any]]: + """Display menu and make multiple selections, returns tuple. + + NOTE: Menu is displayed until an action entry is selected. + """ + while True: + self._update(single_selection=False) + user_selection = self._user_select(prompt_msg) + selected_entry = self._resolve_selection(user_selection) + if user_selection.isnumeric(): + # Update selection(s) + self._update_entry_selection_status(selected_entry[0]) + else: + # Action selected + break + + # Done + return selected_entry + + def settings_select( + self, + prompt_msg: str = 'Please make a selection: ', + ) -> tuple[str, dict[Any, Any]]: + """Display menu and make multiple selections, returns tuple. + + NOTE: Menu is displayed until an action entry is selected. + """ + choice_kwargs = { + 'prompt_msg': 'Toggle or change value?', + 'choices': ['T', 'C'], + } + + while True: + self._update(single_selection=True, settings_mode=True) + user_selection = self._user_select(prompt_msg) + selected_entry = self._resolve_selection(user_selection) + if user_selection.isnumeric(): + if 'Value' in selected_entry[-1] and choice(**choice_kwargs) == 'C': + # Change + selected_entry[-1]['Value'] = input_text('Enter new value: ') + else: + # Toggle + self._update_entry_selection_status(selected_entry[0]) + else: + # Action selected + break + + # Done + return selected_entry + + def simple_select( + self, + prompt_msg: str = 'Please make a selection: ', + update: bool = True, + ) -> tuple[str, dict[Any, Any]]: + """Display menu and make a single selection, returns tuple.""" + if update: + self._update() + user_selection = self._user_select(prompt_msg) + return self._resolve_selection(user_selection) + + def update(self) -> None: + """Update menu with default settings.""" + self._update() + +class TryAndPrint(): + """Object used to standardize running functions and returning the result. + + The errors and warning attributes are used to allow fine-tuned results + based on exception names. + """ + def __init__(self, msg_bad: str = 'FAILED', msg_good: str = 'SUCCESS'): + self.catch_all : bool = True + self.indent: int = INDENT + self.list_errors: list[str] = ['GenericError'] + self.list_warnings: list[str] = ['GenericWarning'] + self.msg_bad: str = msg_bad + self.msg_good: str = msg_good + self.verbose : bool = False + self.width: int = WIDTH + + def _format_exception_message(self, _exception: Exception) -> str: + """Format using the exception's args or name, returns str.""" + LOG.debug( + 'Formatting exception: %s, %s', + _exception.__class__.__name__, + _exception, + ) + message = '' + + # Format message string from _exception + try: + if isinstance(_exception, subprocess.CalledProcessError): + message = _exception.stderr + if not isinstance(message, str): + message = message.decode('utf-8') + message = message.strip() + elif isinstance(_exception, ZeroDivisionError): + # Skip and just use exception name below + pass + else: + message = str(_exception) + except Exception: + # Just use the exception name instead + pass + + # Prepend exception name + if _exception.__class__.__name__ not in ('GenericError', 'GenericWarning'): + try: + message = f'{_exception.__class__.__name__}: {message}' + except Exception: + message = f'UNKNOWN ERROR: {message}' + + # Fix multi-line messages + if '\n' in message: + try: + lines = [ + f'{" "*(self.indent+self.width)}{line.strip()}' + for line in message.splitlines() if line.strip() + ] + lines[0] = lines[0].strip() + message = '\n'.join(lines) + except Exception: + pass + + # Done + return message + + def _format_function_output( + self, + output: list | subprocess.CompletedProcess, + msg_good: str, + ) -> str: + """Format function output for use in try_and_print(), returns str.""" + LOG.debug('Formatting output: %s', output) + + if not output: + raise GenericWarning('No output') + + # Ensure we're working with a list + if isinstance(output, subprocess.CompletedProcess): + stdout = output.stdout + if not isinstance(stdout, str): + stdout = stdout.decode('utf8') + output = stdout.strip().splitlines() + if not output: + # Going to treat these as successes (for now) + LOG.warning('Program output was empty, assuming good result.') + return color_string(msg_good, 'GREEN') + else: + try: + output = list(output) + except TypeError: + output = [output] + + # Safety check + if not output: + # Going to ignore empty function output for now + LOG.error('Output is empty') + return 'UNKNOWN' + + # Build result_msg + result_msg = f'{output.pop(0)}' + if output: + output = [f'{" "*(self.indent+self.width)}{line}' for line in output] + result_msg += '\n' + '\n'.join(output) + + # Done + return result_msg + + def _log_result(self, message: str, result_msg: str) -> None: + """Log result text without color formatting.""" + log_text = f'{" "*self.indent}{message:<{self.width}}{result_msg}' + for line in log_text.splitlines(): + line = strip_colors(line) + LOG.info(line) + + def add_error(self, exception_name: str) -> None: + """Add exception name to error list.""" + if exception_name not in self.list_errors: + self.list_errors.append(exception_name) + + def add_warning(self, exception_name: str) -> None: + """Add exception name to warning list.""" + if exception_name not in self.list_warnings: + self.list_warnings.append(exception_name) + + def run( + self, + message: str, + function: Callable, + *args: Iterable[Any], + catch_all: bool | None = None, + msg_good: str | None = None, + verbose: bool | None = None, + **kwargs, + ) -> dict[str, Any]: + """Run a function and print the results, returns results as dict. + + If catch_all is True then (nearly) all exceptions will be caught. + Otherwise if an exception occurs that wasn't specified it will be + re-raised. + + If the function returns data it will be used instead of msg_good, + msg_bad, or exception text. + The output should be a list or a subprocess.CompletedProcess object. + + If msg_good is passed it will override self.msg_good. + + If verbose is True then exception names or messages will be used for + the result message. Otherwise it will simply be set to result_bad. + + If catch_all and/or verbose are passed it will override + self.catch_all and/or self.verbose for this call. + + args and kwargs are passed to the function. + """ + LOG.debug('function: %s.%s', function.__module__, function.__name__) + LOG.debug('args: %s', args) + LOG.debug('kwargs: %s', kwargs) + LOG.debug( + 'catch_all: %s, msg_good: %s, verbose: %s', + catch_all, + msg_good, + verbose, + ) + f_exception = None + catch_all = catch_all if catch_all is not None else self.catch_all + msg_good = msg_good if msg_good is not None else self.msg_good + output = None + result_msg = 'UNKNOWN' + verbose = verbose if verbose is not None else self.verbose + + # Build exception tuples + e_exceptions: tuple = tuple(get_exception(e) for e in self.list_errors) + w_exceptions: tuple = tuple(get_exception(e) for e in self.list_warnings) + + # Run function and catch exceptions + print(f'{" "*self.indent}{message:<{self.width}}', end='', flush=True) + LOG.debug('Running function: %s.%s', function.__module__, function.__name__) + try: + output = function(*args, **kwargs) + except w_exceptions as _exception: + # Warnings + result_msg = self._format_exception_message(_exception) + print_warning(result_msg, log=False) + f_exception = _exception + except e_exceptions as _exception: + # Exceptions + result_msg = self._format_exception_message(_exception) + print_error(result_msg, log=False) + f_exception = _exception + except Exception as _exception: + # Unexpected exceptions + if verbose: + result_msg = self._format_exception_message(_exception) + else: + result_msg = self.msg_bad + print_error(result_msg, log=False) + f_exception = _exception + if not catch_all: + # Re-raise error as necessary + raise + else: + # Success + if output: + result_msg = self._format_function_output(output, msg_good) + print(result_msg) + else: + result_msg = msg_good + print_success(result_msg, log=False) + + # Done + self._log_result(message, result_msg) + return { + 'Exception': f_exception, + 'Failed': bool(f_exception), + 'Message': result_msg, + 'Output': output, + } + + +# Functions +def abort( + prompt_msg: str = 'Aborted.', + show_prompt_msg: bool = True, + return_code: int = 1, + ) -> None: + """Abort script.""" + print_warning(prompt_msg) + if show_prompt_msg: + sleep(0.5) + pause(prompt_msg='Press Enter to exit... ') + sys.exit(return_code) + + +def ask(prompt_msg: str) -> bool: + """Prompt the user with a Y/N question, returns bool.""" + validator = InputYesNoValidator() + + # Show prompt + response = input_text(f'{prompt_msg} [Y/N]: ', validator=validator) + if response.upper().startswith('Y'): + LOG.info('%s Yes', prompt_msg) + return True + if response.upper().startswith('N'): + LOG.info('%s No', prompt_msg) + return False + + # This shouldn't ever be reached + raise ValueError(f'Invalid answer given: {response}') + + +def beep(repeat: int = 1) -> None: + """Play system bell with optional repeat.""" + while repeat >= 1: + # Print bell char without a newline + print('\a', end='', flush=True) + sleep(0.5) + repeat -= 1 + + +def choice(prompt_msg: str, choices: Iterable[str]) -> str: + """Choose an option from a provided list, returns str. + + Choices provided will be converted to uppercase and returned as such. + Similar to the commands choice (Windows) and select (Linux). + """ + LOG.debug('prompt_msg: %s, choices: %s', prompt_msg, choices) + choices = [str(c).upper()[:1] for c in choices] + prompt_msg = f'{prompt_msg} [{"/".join(choices)}]' + + # Show prompt + response = input_text(prompt_msg, validator=InputChoiceValidator(choices)) + + # Done + LOG.info('%s %s', prompt_msg, response) + return response.upper() + + +def fix_prompt(message: str) -> str: + """Fix prompt, returns str.""" + if not message: + message = 'Input text: ' + message = str(message) + if message[-1:] != ' ': + message += ' ' + return message + + +@cache +def get_exception(name: str) -> Exception: + """Get exception by name, returns exception object. + + [Doctest] + >>> t = TryAndPrint() + >>> t._get_exception('AttributeError') + + >>> t._get_exception('CalledProcessError') + + >>> t._get_exception('GenericError') + + """ + LOG.debug('Getting exception: %s', name) + obj = getattr(sys.modules[__name__], name, None) + if obj: + return obj + + # Try builtin classes + obj = getattr(sys.modules['builtins'], name, None) + if obj: + return obj + + # Try all modules + for _mod in sys.modules.values(): + obj = getattr(_mod, name, None) + if obj: + break + + # Check if not found + if not obj: + raise AttributeError(f'Failed to find exception: {name}') + + # Done + return obj + + +def get_ticket_id() -> str: + """Get ticket ID, returns str.""" + prompt_msg = 'Please enter ticket ID:' + validator = InputTicketIDValidator() + + # Show prompt + ticket_id = input_text(prompt_msg, validator=validator) + + # Done + return ticket_id + + +def input_text( + prompt_msg: str = 'Enter text: ', + allow_empty: bool = False, + validator: Validator | None = None, + ) -> str: + """Get input from user, returns str.""" + prompt_msg = fix_prompt(prompt_msg) + + # Accept empty responses? + if not (allow_empty or validator): + validator = InputNotEmptyValidator() + + # Show prompt + result = None + while result is None: + try: + result = prompt(prompt_msg, validator=validator) + except KeyboardInterrupt: + # Ignore CTRL+c + pass + + # Done + return result + + +def major_exception() -> None: + """Display traceback, optionally upload detailes, and exit.""" + LOG.critical('Major exception encountered', exc_info=True) + print_error('Major exception', log=False) + print_warning(SUPPORT_MESSAGE) + if ENABLED_UPLOAD_DATA: + print_warning('Also, please run upload-logs to help debugging!') + print(traceback.format_exc()) + + # Done + pause('Press Enter to exit... ') + raise SystemExit(1) + + +def pause(prompt_msg: str = 'Press Enter to continue... ') -> None: + """Simple pause implementation.""" + input_text(prompt_msg, allow_empty=True) + + +def print_colored( + strings: Iterable[str] | str, + colors: Iterable[str | None] | str, + log: bool = False, + sep: str = ' ', + **kwargs, + ) -> None: + """Prints strings in the colors specified.""" + LOG.debug( + 'strings: %s, colors: %s, sep: %s, kwargs: %s', + strings, colors, sep, kwargs, + ) + msg = color_string(strings, colors, sep=sep) + print_options = { + 'end': kwargs.get('end', '\n'), + 'file': kwargs.get('file', sys.stdout), + 'flush': kwargs.get('flush', False), + } + + print(msg, **print_options) + if log: + LOG.info(strip_colors(msg)) + + +def print_error(msg: str, log: bool = True, **kwargs) -> None: + """Prints message in RED and log as ERROR.""" + if 'file' not in kwargs: + # Only set if not specified + kwargs['file'] = sys.stderr + print_colored(msg, 'RED', **kwargs) + if log: + LOG.error(msg) + + +def print_info(msg: str, log: bool = True, **kwargs) -> None: + """Prints message in BLUE and log as INFO.""" + print_colored(msg, 'BLUE', **kwargs) + if log: + LOG.info(msg) + + +def print_report(report: list[str], indent=None, log: bool = True) -> None: + """Print report to screen and optionally to log.""" + for line in report: + if indent: + line = f'{" "*indent}{line}' + print(line) + if log: + LOG.info(strip_colors(line)) + + +def print_standard(msg: str, log: bool = True, **kwargs) -> None: + """Prints message and log as INFO.""" + print(msg, **kwargs) + if log: + LOG.info(msg) + + +def print_success(msg: str, log: bool = True, **kwargs) -> None: + """Prints message in GREEN and log as INFO.""" + print_colored(msg, 'GREEN', **kwargs) + if log: + LOG.info(msg) + + +def print_warning(msg: str, log: bool = True, **kwargs) -> None: + """Prints message in YELLOW and log as WARNING.""" + if 'file' not in kwargs: + # Only set if not specified + kwargs['file'] = sys.stderr + print_colored(msg, 'YELLOW', **kwargs) + if log: + LOG.warning(msg) + + +def set_title(title: str) -> None: + """Set window title.""" + LOG.debug('title: %s', title) + if os.name == 'nt': + os.system(f'title {title}') + else: + print_error('Setting the title is only supported under Windows.') + + +def show_data( + message: str, + data: Any, + color: str | None = None, + indent: int | None = None, + width: int | None = None, + ) -> None: + """Display info using default or provided indent and width.""" + indent = INDENT if indent is None else indent + width = WIDTH if width is None else width + print_colored( + (f'{" "*indent}{message:<{width}}', data), + (None, color if color else None), + log=True, + sep='', + ) + + +if __name__ == '__main__': + print("This file is not meant to be called directly.") diff --git a/scripts/wk/tmux.py b/scripts/wk/ui/tmux.py similarity index 52% rename from scripts/wk/tmux.py rename to scripts/wk/ui/tmux.py index 49fb3241..943b6aa4 100644 --- a/scripts/wk/tmux.py +++ b/scripts/wk/ui/tmux.py @@ -4,6 +4,8 @@ import logging import pathlib +from typing import Any + from wk.exe import run_program from wk.std import PLATFORM @@ -13,7 +15,7 @@ LOG = logging.getLogger(__name__) # Functions -def capture_pane(pane_id=None): +def capture_pane(pane_id: str | None = None) -> str: """Capture text from current or target pane, returns str.""" cmd = ['tmux', 'capture-pane', '-p'] if pane_id: @@ -24,49 +26,126 @@ def capture_pane(pane_id=None): return proc.stdout.strip() -def clear_pane(pane_id=None): +def clear_pane(pane_id: str | None = None) -> None: """Clear pane buffer for current or target pane.""" - cmd = ['tmux', 'send-keys', '-R'] + commands = [ + ['tmux', 'send-keys', '-R'], + ['tmux', 'clear-history'], + ] if pane_id: - cmd.extend(['-t', pane_id]) + commands = [[*cmd, '-t', pane_id] for cmd in commands] # Clear pane - run_program(cmd, check=False) + for cmd in commands: + run_program(cmd, check=False) -def fix_layout(panes, layout, forced=False): +def fix_layout( + layout: dict[str, dict[str, Any]], + clear_on_resize: bool = False, + forced: bool = False, + ) -> None: """Fix pane sizes based on layout.""" - if not (forced or layout_needs_fixed(panes, layout)): + resize_kwargs = [] + + # Bail early + if not (forced or layout_needs_fixed(layout)): # Layout should be fine return - # Update panes - for name, data in layout.items(): - # Skip missing panes - if name not in panes: + # Clear current pane if needed + if clear_on_resize: + clear_pane() + + # Remove closed panes + for data in layout.values(): + data['Panes'] = [pane for pane in data['Panes'] if poll_pane(pane)] + + # Calc height for "floating" row + # NOTE: We start with height +1 to account for the splits (i.e. splits = num rows - 1) + floating_height = 1 + get_window_size()[1] + for group in ('Title', 'Info', 'Current', 'Workers'): + if layout[group]['Panes']: + group_height = 1 + layout[group].get('height', 0) + if group == 'Workers': + group_height *= len(layout[group]['Panes']) + floating_height -= group_height + + # Update main panes + for section, data in layout.items(): + # "Floating" pane(s) + if 'height' not in data and section in ('Info', 'Current', 'Workers'): + for pane_id in data['Panes']: + resize_kwargs.append({'pane_id': pane_id, 'height': floating_height}) + + # Rest of the panes + if section == 'Workers': + # Skip for now continue + if 'height' in data: + for pane_id in data['Panes']: + resize_kwargs.append({'pane_id': pane_id, 'height': data['height']}) + if 'width' in data: + for pane_id in data['Panes']: + resize_kwargs.append({'pane_id': pane_id, 'width': data['width']}) + for kwargs in resize_kwargs: + try: + resize_pane(**kwargs) + except RuntimeError: + # Assuming pane was closed just before resizing + pass - # Resize pane(s) - pane_list = panes[name] - if isinstance(pane_list, str): - pane_list = [pane_list] - for pane_id in pane_list: - if name == 'Current': - pane_id = None - try: - resize_pane(pane_id, **data) - except RuntimeError: - # Assuming pane was closed just before resizing - pass + # Update "group" panes widths + for group in ('Title', 'Info'): + num_panes = len(layout[group]['Panes']) + if num_panes <= 1: + continue + width = int( (get_pane_size()[0] - (1 - num_panes)) / num_panes ) + for pane_id in layout[group]['Panes']: + resize_pane(pane_id, width=width) + if group == 'Title': + # (re)fix Started pane + resize_pane(layout['Started']['Panes'][0], width=layout['Started']['width']) + + # Bail early + if not ( + layout['Workers']['Panes'] + and 'height' in layout['Workers'] + and floating_height > 0 + ): + return + + # Update worker heights + for worker in reversed(layout['Workers']['Panes']): + resize_pane(worker, height=layout['Workers']['height']) -def get_pane_size(pane_id=None): +def get_pane_size(pane_id: str | None = None) -> tuple[int, int]: """Get current or target pane size, returns tuple.""" cmd = ['tmux', 'display', '-p'] if pane_id: cmd.extend(['-t', pane_id]) cmd.append('#{pane_width} #{pane_height}') + # Get resolution + proc = run_program(cmd, check=False) + try: + width, height = proc.stdout.strip().split() + except ValueError: + # Assuming this is a race condition as it usually happens inside the + # background fix layout loop + return 0, 0 + width = int(width) + height = int(height) + + # Done + return (width, height) + + +def get_window_size() -> tuple[int, int]: + """Get current window size, returns tuple.""" + cmd = ['tmux', 'display', '-p', '#{window_width} #{window_height}'] + # Get resolution proc = run_program(cmd, check=False) width, height = proc.stdout.strip().split() @@ -77,7 +156,7 @@ def get_pane_size(pane_id=None): return (width, height) -def kill_all_panes(pane_id=None): +def kill_all_panes(pane_id: str | None = None) -> None: """Kill all panes except for the current or target pane.""" cmd = ['tmux', 'kill-pane', '-a'] if pane_id: @@ -87,7 +166,7 @@ def kill_all_panes(pane_id=None): run_program(cmd, check=False) -def kill_pane(*pane_ids): +def kill_pane(*pane_ids: str) -> None: """Kill pane(s) by id.""" cmd = ['tmux', 'kill-pane', '-t'] @@ -96,40 +175,26 @@ def kill_pane(*pane_ids): run_program(cmd+[pane_id], check=False) -def layout_needs_fixed(panes, layout): +def layout_needs_fixed(layout: dict[str, dict[str, Any]]) -> bool: """Check if layout needs fixed, returns bool.""" needs_fixed = False # Check panes - for name, data in layout.items(): - # Skip unpredictably sized panes - if not data.get('Check', False): - continue - - # Skip missing panes - if name not in panes: - continue - - # Check pane size(s) - pane_list = panes[name] - if isinstance(pane_list, str): - pane_list = [pane_list] - for pane_id in pane_list: - try: - width, height = get_pane_size(pane_id) - except ValueError: - # Pane may have disappeared during this loop - continue - if data.get('width', False) and data['width'] != width: - needs_fixed = True - if data.get('height', False) and data['height'] != height: - needs_fixed = True + for data in layout.values(): + if 'height' in data: + needs_fixed = needs_fixed or any( + get_pane_size(pane)[1] != data['height'] for pane in data['Panes'] + ) + if 'width' in data: + needs_fixed = needs_fixed or any( + get_pane_size(pane)[0] != data['width'] for pane in data['Panes'] + ) # Done return needs_fixed -def poll_pane(pane_id): +def poll_pane(pane_id: str) -> bool: """Check if pane exists, returns bool.""" cmd = ['tmux', 'list-panes', '-F', '#D'] @@ -142,7 +207,12 @@ def poll_pane(pane_id): def prep_action( - cmd=None, working_dir=None, text=None, watch_file=None, watch_cmd='cat'): + cmd: str | None = None, + working_dir: pathlib.Path | str | None = None, + text: str | None = None, + watch_file: pathlib.Path | str | None = None, + watch_cmd: str = 'cat', + ) -> list[str]: """Prep action to perform during a tmux call, returns list. This will prep for running a basic command, displaying text on screen, @@ -192,7 +262,7 @@ def prep_action( return action_cmd -def prep_file(path): +def prep_file(path: pathlib.Path | str) -> None: """Check if file exists and create empty file if not.""" path = pathlib.Path(path).resolve() try: @@ -202,7 +272,11 @@ def prep_file(path): pass -def resize_pane(pane_id=None, width=None, height=None, **kwargs): +def resize_pane( + pane_id: str | None = None, + width: int | None = None, + height: int | None = None, + ) -> None: """Resize current or target pane. NOTE: kwargs is only here to make calling this function easier @@ -227,10 +301,22 @@ def resize_pane(pane_id=None, width=None, height=None, **kwargs): run_program(cmd, check=False) +def respawn_pane(pane_id: str, **action) -> None: + """Respawn pane with action.""" + cmd = ['tmux', 'respawn-pane', '-k', '-t', pane_id] + cmd.extend(prep_action(**action)) + + # Respawn + run_program(cmd, check=False) + + def split_window( - lines=None, percent=None, - behind=False, vertical=False, - target_id=None, **action): + lines: int | None = None, + percent: int | None = None, + behind: bool = False, + vertical: bool = False, + target_id: str | None = None, + **action) -> str: """Split tmux window, run action, and return pane_id as str.""" cmd = ['tmux', 'split-window', '-d', '-PF', '#D'] @@ -263,16 +349,7 @@ def split_window( return proc.stdout.strip() -def respawn_pane(pane_id, **action): - """Respawn pane with action.""" - cmd = ['tmux', 'respawn-pane', '-k', '-t', pane_id] - cmd.extend(prep_action(**action)) - - # Respawn - run_program(cmd, check=False) - - -def zoom_pane(pane_id=None): +def zoom_pane(pane_id: str | None = None) -> None: """Toggle zoom status for current or target pane.""" cmd = ['tmux', 'resize-pane', '-Z'] if pane_id: diff --git a/scripts/wk/ui/tui.py b/scripts/wk/ui/tui.py new file mode 100644 index 00000000..0cea76ad --- /dev/null +++ b/scripts/wk/ui/tui.py @@ -0,0 +1,377 @@ +"""WizardKit: TUI functions""" +# vim: sts=2 sw=2 ts=2 + +import atexit +import logging +import time + +from copy import deepcopy +from os import environ +from typing import Any + +from wk.exe import start_thread +from wk.std import sleep +from wk.ui import ansi, tmux + +# STATIC VARIABLES +LOG = logging.getLogger(__name__) +TMUX_SIDE_WIDTH = 21 +TMUX_TITLE_HEIGHT = 2 +TMUX_LAYOUT = { # NOTE: This needs to be in order from top to bottom + 'Title': {'Panes': [], 'height': TMUX_TITLE_HEIGHT}, + 'Info': {'Panes': []}, + 'Current': {'Panes': [environ.get('TMUX_PANE', None)]}, + 'Workers': {'Panes': []}, + 'Started': {'Panes': [], 'width': TMUX_SIDE_WIDTH}, + 'Progress': {'Panes': [], 'width': TMUX_SIDE_WIDTH}, + } + + +# Classes +class TUI(): + """Object for tracking TUI elements.""" + def __init__(self, title_text: str | None = None): + self.clear_on_resize = False + self.layout: dict[str, dict[str, Any]] = deepcopy(TMUX_LAYOUT) + self.side_width: int = TMUX_SIDE_WIDTH + self.title_text: str = title_text if title_text else 'Title Text' + self.title_text_line2: str = '' + self.title_colors: list[str] = ['BLUE', ''] + + # Init tmux and start a background process to maintain layout + self.init_tmux() + start_thread(self.fix_layout_loop) + + # Close all panes at exit + atexit.register(tmux.kill_all_panes) + + def add_info_pane( + self, + lines: int | None = None, + percent: int = 0, + update_layout: bool = True, + **tmux_args, + ) -> None: + """Add info pane.""" + if not (lines or percent): + # Bail early + raise RuntimeError('Neither lines nor percent specified.') + + # Calculate lines if needed + if not lines: + lines = int(tmux.get_pane_size()[1] * (percent/100)) + + # Set tmux split args + tmux_args.update({ + 'behind': True, + 'lines': lines, + 'target_id': None, + 'vertical': True, + }) + if self.layout['Info']['Panes']: + tmux_args.update({ + 'behind': False, + 'percent': 50, + 'target_id': self.layout['Info']['Panes'][-1], + 'vertical': False, + }) + tmux_args.pop('lines') + + # Update layout + if update_layout: + self.layout['Info']['height'] = lines + + # Add pane + self.layout['Info']['Panes'].append(tmux.split_window(**tmux_args)) + + def add_title_pane( + self, + line1: str, + line2: str | None = None, + colors: list[str] | None = None, + ) -> None: + """Add pane to title row.""" + lines = [line1, line2] + colors = colors if colors else self.title_colors.copy() + if not line2: + lines.pop() + colors.pop() + tmux_args = { + 'behind': True, + 'lines': TMUX_TITLE_HEIGHT, + 'target_id': None, + 'text': ansi.color_string(lines, colors, sep='\n'), + 'vertical': True, + } + if self.layout['Title']['Panes']: + tmux_args.update({ + 'behind': False, + 'percent': 50, + 'target_id': self.layout['Title']['Panes'][-1], + 'vertical': False, + }) + tmux_args.pop('lines') + + # Add pane + self.layout['Title']['Panes'].append(tmux.split_window(**tmux_args)) + + def add_worker_pane( + self, + lines: int | None = None, + percent: int = 0, + update_layout: bool = True, + **tmux_args, + ) -> None: + """Add worker pane.""" + height = lines + + # Bail early + if not (lines or percent): + raise RuntimeError('Neither lines nor percent specified.') + + # Calculate height if needed + if not height: + height = int(tmux.get_pane_size()[1] * (percent/100)) + + # Set tmux split args + tmux_args.update({ + 'behind': False, + 'lines': lines, + 'percent': percent if percent else None, + 'target_id': None, + 'vertical': True, + }) + + # Update layout + if update_layout: + self.layout['Workers']['height'] = height + + # Add pane + self.layout['Workers']['Panes'].append(tmux.split_window(**tmux_args)) + + def clear_current_pane(self) -> None: + """Clear screen and history for current pane.""" + tmux.clear_pane() + + def clear_current_pane_height(self) -> None: + """Clear current pane height and update layout.""" + self.layout['Current'].pop('height', None) + + def fix_layout(self, forced: bool = True) -> None: + """Fix tmux layout based on self.layout.""" + try: + tmux.fix_layout(self.layout, clear_on_resize=self.clear_on_resize, forced=forced) + except RuntimeError: + # Assuming self.panes changed while running + pass + + def fix_layout_loop(self) -> None: + """Fix layout on a loop. + + NOTE: This should be called as a thread. + """ + while True: + self.fix_layout(forced=False) + sleep(1) + + def init_tmux(self) -> None: + """Initialize tmux layout.""" + tmux.kill_all_panes() + self.layout.clear() + self.layout.update(deepcopy(TMUX_LAYOUT)) + + # Title + self.layout['Title']['Panes'].append(tmux.split_window( + behind=True, + lines=2, + vertical=True, + text=ansi.color_string( + [self.title_text, self.title_text_line2], + self.title_colors, + sep = '\n', + ), + )) + + # Started + self.layout['Started']['Panes'].append(tmux.split_window( + lines=TMUX_SIDE_WIDTH, + target_id=self.layout['Title']['Panes'][0], + text=ansi.color_string( + ['Started', time.strftime("%Y-%m-%d %H:%M %Z")], + ['BLUE', None], + sep='\n', + ), + )) + + # Progress + self.layout['Progress']['Panes'].append(tmux.split_window( + lines=TMUX_SIDE_WIDTH, + text=' ', + )) + + def remove_all_info_panes(self) -> None: + """Remove all info panes and update layout.""" + self.layout['Info'].pop('height', None) + panes = self.layout['Info']['Panes'].copy() + self.layout['Info']['Panes'].clear() + tmux.kill_pane(*panes) + + def remove_all_worker_panes(self) -> None: + """Remove all worker panes and update layout.""" + self.layout['Workers'].pop('height', None) + panes = self.layout['Workers']['Panes'].copy() + self.layout['Workers']['Panes'].clear() + tmux.kill_pane(*panes) + + def reset_title_pane( + self, + line1: str = 'Title Text', + line2: str = '', + colors: list[str] | None = None, + ) -> None: + """Remove all extra title panes, reset main title pane, and update layout.""" + colors = self.title_colors if colors is None else colors + panes = self.layout['Title']['Panes'].copy() + if len(panes) > 1: + tmux.kill_pane(*panes[1:]) + self.layout['Title']['Panes'] = panes[:1] + self.set_title(line1, line2, colors) + + def set_current_pane_height(self, height: int) -> None: + """Set current pane height and update layout.""" + self.layout['Current']['height'] = height + tmux.resize_pane(height=height) + + def set_progress_file(self, progress_file: str) -> None: + """Set the file to use for the progresse pane.""" + tmux.respawn_pane( + pane_id=self.layout['Progress']['Panes'][0], + watch_file=progress_file, + ) + + def set_title( + self, + line1: str, + line2: str | None = None, + colors: list[str] | None = None, + ) -> None: + """Set title text.""" + self.title_text = line1 + self.title_text_line2 = line2 if line2 else '' + if colors: + self.title_colors = colors + + # Update pane (if present) + if self.layout['Title']['Panes']: + tmux.respawn_pane( + pane_id=self.layout['Title']['Panes'][0], + text=ansi.color_string( + [self.title_text, self.title_text_line2], + self.title_colors, + sep = '\n', + ), + ) + + def update_clock(self) -> None: + """Update 'Started' pane following clock sync.""" + tmux.respawn_pane( + pane_id=self.layout['Started']['Panes'][0], + text=ansi.color_string( + ['Started', time.strftime("%Y-%m-%d %H:%M %Z")], + ['BLUE', None], + sep='\n', + ), + ) + + +# Functions +def fix_layout(layout, forced: bool = False) -> None: + """Fix pane sizes based on layout.""" + resize_kwargs = [] + + # Bail early + if not (forced or layout_needs_fixed(layout)): + # Layout should be fine + return + + # Remove closed panes + for data in layout.values(): + data['Panes'] = [pane for pane in data['Panes'] if tmux.poll_pane(pane)] + + # Update main panes + for section, data in layout.items(): + if section == 'Workers': + # Skip for now + continue + + if 'height' in data: + for pane_id in data['Panes']: + resize_kwargs.append({'pane_id': pane_id, 'height': data['height']}) + if 'width' in data: + for pane_id in data['Panes']: + resize_kwargs.append({'pane_id': pane_id, 'width': data['width']}) + for kwargs in resize_kwargs: + try: + tmux.resize_pane(**kwargs) + except RuntimeError: + # Assuming pane was closed just before resizing + pass + + # Update "group" panes widths + for group in ('Title', 'Info'): + num_panes = len(layout[group]['Panes']) + if num_panes <= 1: + continue + width = int( (tmux.get_pane_size()[0] - (1 - num_panes)) / num_panes ) + for pane_id in layout[group]['Panes']: + tmux.resize_pane(pane_id, width=width) + if group == 'Title': + # (re)fix Started pane + tmux.resize_pane(layout['Started']['Panes'][0], width=TMUX_SIDE_WIDTH) + + # Bail early + if not (layout['Workers']['Panes'] and layout['Workers']['height']): + return + + # Update worker heights + worker_height = layout['Workers']['height'] + workers = layout['Workers']['Panes'].copy() + num_workers = len(workers) + avail_height = sum(tmux.get_pane_size(pane)[1] for pane in workers) + avail_height += tmux.get_pane_size()[1] # Current pane + # Check if window is too small + if avail_height < (worker_height*num_workers) + 3: + # Just leave things as-is + return + # Resize current pane + tmux.resize_pane(height=avail_height-(worker_height*num_workers)) + # Resize bottom pane + tmux.resize_pane(workers.pop(0), height=worker_height) + # Resize the rest of the panes by adjusting the ones above them + while len(workers) > 1: + next_height = sum(tmux.get_pane_size(pane)[1] for pane in workers[:2]) + next_height -= worker_height + tmux.resize_pane(workers[1], height=next_height) + workers.pop(0) + +def layout_needs_fixed(layout) -> bool: + """Check if layout needs fixed, returns bool.""" + needs_fixed = False + + # Check panes + for data in layout.values(): + if 'height' in data: + needs_fixed = needs_fixed or any( + tmux.get_pane_size(pane)[1] != data['height'] for pane in data['Panes'] + ) + if 'width' in data: + needs_fixed = needs_fixed or any( + tmux.get_pane_size(pane)[0] != data['width'] for pane in data['Panes'] + ) + + # Done + return needs_fixed + + +if __name__ == '__main__': + print("This file is not meant to be called directly.") diff --git a/scripts/wk_debug.py b/scripts/wk_debug.py index 86900004..1808be89 100755 --- a/scripts/wk_debug.py +++ b/scripts/wk_debug.py @@ -17,7 +17,7 @@ OPTIONS = { def get_debug_prefix() -> str: """Ask what we're debugging, returns log dir prefix.""" - menu = wk.std.Menu(title=f'{wk.cfg.main.KIT_NAME_FULL}: Debugging Menu\n') + menu = wk.ui.cli.Menu(title=f'{wk.cfg.main.KIT_NAME_FULL}: Debugging Menu\n') for name, prefix in OPTIONS.items(): menu.add_option(name, {'Prefix': prefix}) selection = menu.simple_select() @@ -38,14 +38,14 @@ def get_debug_path() -> pathlib.Path: # Safety check if not debug_paths: - wk.std.abort('No logs found, aborting.') + wk.ui.cli.abort('No logs found, aborting.') # Use latest option - if wk.std.ask('Use latest session?'): + if wk.ui.cli.ask('Use latest session?'): return debug_paths[-1] # Select from list - menu = wk.std.Menu(title=f'{wk.cfg.main.KIT_NAME_FULL}: Debugging Menu\n') + menu = wk.ui.cli.Menu(title=f'{wk.cfg.main.KIT_NAME_FULL}: Debugging Menu\n') for item in debug_paths: menu.add_option(item.parent.name, {'Path': item}) selection = menu.simple_select() diff --git a/setup/linux/packages/base b/setup/linux/packages/base index a2c839d7..b51cc82d 100644 --- a/setup/linux/packages/base +++ b/setup/linux/packages/base @@ -100,6 +100,7 @@ python-docopt python-gnuplot python-mariadb-connector python-packaging +python-prompt_toolkit python-psutil python-pytz python-requests diff --git a/setup/windows/build.ps1 b/setup/windows/build.ps1 index 31e6de9d..6479f027 100644 --- a/setup/windows/build.ps1 +++ b/setup/windows/build.ps1 @@ -116,9 +116,9 @@ if ($MyInvocation.InvocationName -ne ".") { $Url = FindDynamicUrl $DownloadPage $RegEx DownloadFile -Path $Temp -Name "psutil64.whl" -Url $Url - # Python: pytz, requests, & dependencies + # Python: prompt_toolkit, pytz, requests, & dependancies $RegEx = "href=.*.py3-none-any.whl" - foreach ($Module in @("certifi", "chardet", "idna", "packaging", "pytz", "requests", "urllib3")) { + foreach ($Module in @("certifi", "chardet", "idna", "packaging", "prompt_toolkit", "Pygments", "pytz", "requests", "urllib3", "wcwidth")) { $DownloadPage = "https://pypi.org/project/$Module/" $Name = "$Module.whl" $Url = FindDynamicUrl -SourcePage $DownloadPage -RegEx $RegEx