Merge remote-tracking branch 'upstream/dev' into dev

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()

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()

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()

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()

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()

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='',

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()

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()

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'

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

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():

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()

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()

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

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"]

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()

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()

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__':

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 = {

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

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 = {

scripts/wk/clone/__init__.py

@@ -1,3 +1,4 @@
 """WizardKit: ddrescue-tui module init"""
 
 from . import ddrescue
+from . import menus

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.")

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.")

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)

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}',

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
 

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)

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:

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}')
 

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:

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__':

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__':

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)
 
 
 

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')

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='',

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

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

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)

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__':

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."""

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',

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)

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 = []
 

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.

scripts/wk/os/mac.py

@@ -16,7 +16,7 @@ REGEX_FANS = re.compile(r'^.*\(bytes (?P<bytes>.*)\)$')
 
 
 # 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

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)

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'])
 

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)

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

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.")

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')
+  <class 'AttributeError'>
+  >>> t._get_exception('CalledProcessError')
+  <class 'subprocess.CalledProcessError'>
+  >>> t._get_exception('GenericError')
+  <class 'wk.std.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.")

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:

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.")

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()

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

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