2Shirt/WizardKit
2Shirt/WizardKit
1
0
Fork 0
Code Issues 2 Pull requests Projects Releases 26 Packages Wiki Activity Actions

Merge branch 'ui-split' into dev

Browse source
This commit is contained in:
2Shirt 2023-05-27 19:50:49 -07:00
commit 0126452bf1
Signed by: 2Shirt
GPG key ID: 152FAC923B0E132C
48 changed files with 2088 additions and 1705 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
scripts/auto_repairs.py
View file

@ -5,7 +5,7 @@ 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):
@ -87,8 +87,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()

4
scripts/auto_setup.py
View file

@ -154,8 +154,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()

2
scripts/build-ufd
View file

@ -11,4 +11,4 @@ if __name__ == '__main__':
except SystemExit:
raise
except: # noqa: E722
wk.std.major_exception()
wk.ui.cli.major_exception()

4
scripts/build_kit_windows.py
View file

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

4
scripts/ddrescue-tui.py
View file

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

2
scripts/embedded_python_env.py
View file

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

4
scripts/hw-diags.py
View file

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

6
scripts/hw-sensors
View file

@ -11,11 +11,11 @@ def main():
"""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()

14
scripts/launch_sdio.py
View file

@ -22,10 +22,10 @@ SDIO_REMOTE_PATH = wk.io.get_path_obj(
# Functions
def try_again():
"""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
@ -47,7 +47,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 +57,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 +66,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 +76,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

18
scripts/mount-all-volumes
View file

@ -10,32 +10,32 @@ import wk
# Functions
def main():
"""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()

6
scripts/mount-backup-shares
View file

@ -8,7 +8,7 @@ import wk
# Functions
def main():
"""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()

14
scripts/msword-search
View file

@ -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

1
scripts/pyproject.toml
View file

@ -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
"wk/borrowed/acpi.py" = ["E501", "F841"]

6
scripts/unmount-backup-shares
View file

@ -8,14 +8,14 @@ import wk
# Functions
def main():
"""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.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()

14
scripts/upload-logs
View file

@ -28,21 +28,21 @@ if PLATFORM not in ('macOS', 'Linux'):
def main():
"""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

11
scripts/wk/__init__.py
View file

@ -1,7 +1,7 @@
"""WizardKit: wk module init"""
# vim: sts=2 sw=2 ts=2
from sys import version_info as version
from sys import stderr, version_info
from . import cfg
from . import clone
@ -17,21 +17,22 @@ from . import os
from . import repairs
from . import setup
from . import std
from . import tmux
from . import ui
# Check env
if version < (3, 7):
if version_info < (3, 7):
# 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__':

14
scripts/wk/cfg/hw.py
View file

@ -3,8 +3,6 @@
import re
from collections import OrderedDict
# STATIC VARIABLES
ATTRIBUTE_COLORS = (
@ -161,18 +159,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

264
scripts/wk/clone/ddrescue.py
View file

@ -20,7 +20,7 @@ from docopt import docopt
import psutil
import pytz
from wk import cfg, debug, exe, io, log, net, std, tmux
from wk import cfg, debug, exe, io, log, net, std
from wk.cfg.ddrescue import (
DDRESCUE_MAP_TEMPLATE,
DDRESCUE_SETTINGS,
@ -33,6 +33,8 @@ from wk.hw.smart import (
smart_status_ok,
update_smart_details,
)
from wk.ui import cli as ui
from wk.ui import ansi, tmux
# STATIC VARIABLES
@ -89,8 +91,8 @@ REGEX_REMAINING_TIME = re.compile(
LOG = logging.getLogger(__name__)
MENU_ACTIONS = (
'Start',
f'Change settings {std.color_string("(experts only)", "YELLOW")}',
f'Detect drives {std.color_string("(experts only)", "YELLOW")}',
f'Change settings {ansi.color_string("(experts only)", "YELLOW")}',
f'Detect drives {ansi.color_string("(experts only)", "YELLOW")}',
'Quit')
MENU_TOGGLES = {
'Auto continue (if recovery % over threshold)': True,
@ -296,7 +298,7 @@ class BlockPair():
# Check destination size if cloning
if not self.destination.is_file() and dest_size < self.size:
std.print_error(f'Invalid destination: {self.destination}')
ui.print_error(f'Invalid destination: {self.destination}')
raise std.GenericAbort()
def set_initial_status(self):
@ -420,7 +422,7 @@ class State():
self.panes['Started'] = tmux.split_window(
lines=cfg.ddrescue.TMUX_SIDE_WIDTH,
target_id=self.panes['Source'],
text=std.color_string(
text=ansi.color_string(
['Started', time.strftime("%Y-%m-%d %H:%M %Z")],
['BLUE', None],
sep='\n',
@ -442,7 +444,7 @@ class State():
settings = json.loads(_f.read())
except (OSError, json.JSONDecodeError) as err:
LOG.error('Failed to load clone settings')
std.print_error('Invalid clone settings detected.')
ui.print_error('Invalid clone settings detected.')
raise std.GenericAbort() from err
# Check settings
@ -454,10 +456,10 @@ class State():
bail = False
for key in ('model', 'serial'):
if settings['Source'][key] != getattr(self.source, key):
std.print_error(f"Clone settings don't match source {key}")
ui.print_error(f"Clone settings don't match source {key}")
bail = True
if settings['Destination'][key] != getattr(self.destination, key):
std.print_error(f"Clone settings don't match destination {key}")
ui.print_error(f"Clone settings don't match destination {key}")
bail = True
if bail:
raise std.GenericAbort()
@ -488,7 +490,7 @@ class State():
with open(settings_file, 'w', encoding='utf-8') as _f:
json.dump(settings, _f)
except OSError as err:
std.print_error('Failed to save clone settings')
ui.print_error('Failed to save clone settings')
raise std.GenericAbort() from err
def add_clone_block_pairs(self):
@ -522,9 +524,9 @@ class State():
# New run, use new settings file
settings['Needs Format'] = True
offset = 0
user_choice = std.choice(
['G', 'M', 'S'],
user_choice = ui.choice(
'Format clone using GPT, MBR, or match Source type?',
['G', 'M', 'S'],
)
if user_choice == 'G':
settings['Table Type'] = 'GPT'
@ -533,7 +535,7 @@ class State():
else:
# Match source type
settings['Table Type'] = get_table_type(self.source.path)
if std.ask('Create an empty Windows boot partition on the clone?'):
if ui.ask('Create an empty Windows boot partition on the clone?'):
settings['Create Boot Partition'] = True
offset = 2 if settings['Table Type'] == 'GPT' else 1
@ -561,19 +563,19 @@ class State():
bp_dest = self.destination
self._add_block_pair(part, bp_dest)
def confirm_selections(self, prompt, source_parts=None):
def confirm_selections(self, prompt_msg, source_parts=None):
"""Show selection details and prompt for confirmation."""
report = []
# Source
report.append(std.color_string('Source', 'GREEN'))
report.append(ansi.color_string('Source', 'GREEN'))
report.extend(build_object_report(self.source))
report.append(' ')
# Destination
report.append(std.color_string('Destination', 'GREEN'))
report.append(ansi.color_string('Destination', 'GREEN'))
if self.mode == 'Clone':
report[-1] += std.color_string(' (ALL DATA WILL BE DELETED)', 'RED')
report[-1] += ansi.color_string(' (ALL DATA WILL BE DELETED)', 'RED')
report.extend(build_object_report(self.destination))
report.append(' ')
@ -581,12 +583,12 @@ class State():
# NOTE: The check for block_pairs is to limit this section
# to the second confirmation
if self.mode == 'Clone' and self.block_pairs:
report.append(std.color_string('WARNING', 'YELLOW'))
report.append(ansi.color_string('WARNING', 'YELLOW'))
report.append(
'All data will be deleted from the destination listed above.',
)
report.append(
std.color_string(
ansi.color_string(
['This is irreversible and will lead to', 'DATA LOSS.'],
['YELLOW', 'RED'],
),
@ -605,18 +607,18 @@ class State():
# Map dir
if self.working_dir:
report.append(std.color_string('Map Save Directory', 'GREEN'))
report.append(ansi.color_string('Map Save Directory', 'GREEN'))
report.append(f'{self.working_dir}/')
report.append(' ')
if not fstype_is_ok(self.working_dir, map_dir=True):
report.append(
std.color_string(
ansi.color_string(
'Map file(s) are being saved to a non-recommended filesystem.',
'YELLOW',
),
)
report.append(
std.color_string(
ansi.color_string(
['This is strongly discouraged and may lead to', 'DATA LOSS'],
[None, 'RED'],
),
@ -625,11 +627,11 @@ class State():
# Source part(s) selected
if source_parts:
report.append(std.color_string('Source Part(s) selected', 'GREEN'))
report.append(ansi.color_string('Source Part(s) selected', 'GREEN'))
if self.source.path.samefile(source_parts[0].path):
report.append('Whole Disk')
else:
report.append(std.color_string(f'{"NAME":<9} SIZE', 'BLUE'))
report.append(ansi.color_string(f'{"NAME":<9} SIZE', 'BLUE'))
for part in source_parts:
report.append(
f'{part.path.name:<9} '
@ -638,9 +640,9 @@ class State():
report.append(' ')
# Prompt user
std.clear_screen()
std.print_report(report)
if not std.ask(prompt):
ui.clear_screen()
ui.print_report(report)
if not ui.ask(prompt_msg):
raise std.GenericAbort()
def generate_report(self):
@ -661,7 +663,7 @@ class State():
error_size = self.get_error_size()
error_size_str = std.bytes_to_string(error_size, decimals=2)
if error_size > 0:
error_size_str = std.color_string(error_size_str, 'YELLOW')
error_size_str = ansi.color_string(error_size_str, 'YELLOW')
percent = self.get_percent_recovered()
percent = format_status_string(percent, width=0)
report.append(f'Overall rescued: {percent}, error size: {error_size_str}')
@ -673,7 +675,7 @@ class State():
error_size = pair.get_error_size()
error_size_str = std.bytes_to_string(error_size, decimals=2)
if error_size > 0:
error_size_str = std.color_string(error_size_str, 'YELLOW')
error_size_str = ansi.color_string(error_size_str, 'YELLOW')
pair_size = std.bytes_to_string(pair.size, decimals=2)
percent = pair.get_percent_recovered()
percent = format_status_string(percent, width=0)
@ -704,7 +706,7 @@ class State():
def init_recovery(self, docopt_args):
"""Select source/dest and set env."""
std.clear_screen()
ui.clear_screen()
source_parts = []
# Set log
@ -744,7 +746,7 @@ class State():
# Confirmation #1
self.confirm_selections(
prompt='Are these selections correct?',
prompt_msg='Are these selections correct?',
source_parts=source_parts,
)
@ -795,8 +797,8 @@ class State():
try:
exe.run_program(cmd)
except subprocess.CalledProcessError:
std.print_error('Failed to unmount source and/or destination')
std.abort()
ui.print_error('Failed to unmount source and/or destination')
ui.abort()
# Prep destination
if self.mode == 'Clone':
@ -928,7 +930,7 @@ class State():
check=False,
)
if proc.returncode != 0:
std.print_error('Error(s) encoundtered while formatting destination')
ui.print_error('Error(s) encoundtered while formatting destination')
raise std.GenericAbort()
# Update settings
@ -969,13 +971,13 @@ class State():
# Check for critical errors
if not smart_status_ok(self.destination):
std.print_error(
ui.print_error(
f'Critical error(s) detected for: {self.destination.path}',
)
# Check for minor errors
if not check_attributes(self.destination, only_blocking=False):
std.print_warning(
ui.print_warning(
f'Attribute error(s) detected for: {self.destination.path}',
)
@ -1026,7 +1028,7 @@ class State():
destination_size *= 1.05
error_msg = 'Not enough free space on the destination'
if required_size > destination_size:
std.print_error(error_msg)
ui.print_error(error_msg)
raise std.GenericAbort()
def save_debug_reports(self):
@ -1037,7 +1039,7 @@ class State():
debug_dir.mkdir()
# State (self)
std.save_pickles({'state': self}, debug_dir)
debug.save_pickles({'state': self}, debug_dir)
with open(f'{debug_dir}/state.report', 'a', encoding='utf-8') as _f:
_f.write('[Debug report]\n')
_f.write('\n'.join(debug.generate_object_report(self)))
@ -1064,10 +1066,10 @@ class State():
width = cfg.ddrescue.TMUX_SIDE_WIDTH
# Status
report.append(std.color_string(f'{"Status":^{width}}', 'BLUE'))
report.append(ansi.color_string(f'{"Status":^{width}}', 'BLUE'))
if 'NEEDS ATTENTION' in overall_status:
report.append(
std.color_string(f'{overall_status:^{width}}', 'YELLOW_BLINK'),
ansi.color_string(f'{overall_status:^{width}}', 'YELLOW_BLINK'),
)
else:
report.append(f'{overall_status:^{width}}')
@ -1077,12 +1079,12 @@ class State():
if self.block_pairs:
total_rescued = self.get_rescued_size()
percent = self.get_percent_recovered()
report.append(std.color_string('Overall Progress', 'BLUE'))
report.append(ansi.color_string('Overall Progress', 'BLUE'))
report.append(
f'Rescued: {format_status_string(percent, width=width-9)}',
)
report.append(
std.color_string(
ansi.color_string(
[f'{std.bytes_to_string(total_rescued, decimals=2):>{width}}'],
[get_percent_color(percent)],
),
@ -1091,7 +1093,7 @@ class State():
# Block pair progress
for pair in self.block_pairs:
report.append(std.color_string(pair.source, 'BLUE'))
report.append(ansi.color_string(pair.source, 'BLUE'))
for name, status in pair.status.items():
name = name.title()
report.append(
@ -1103,9 +1105,9 @@ class State():
if overall_status in ('Active', 'NEEDS ATTENTION'):
etoc = get_etoc()
report.append(separator)
report.append(std.color_string('Estimated Pass Finish', 'BLUE'))
report.append(ansi.color_string('Estimated Pass Finish', 'BLUE'))
if overall_status == 'NEEDS ATTENTION' or etoc == 'N/A':
report.append(std.color_string('N/A', 'YELLOW'))
report.append(ansi.color_string('N/A', 'YELLOW'))
else:
report.append(etoc)
@ -1166,7 +1168,7 @@ class State():
source_str = _format_string(self.source, width)
tmux.respawn_pane(
self.panes['Source'],
text=std.color_string(
text=ansi.color_string(
['Source', '' if source_exists else ' (Missing)', '\n', source_str],
['BLUE', 'RED', None, None],
sep='',
@ -1181,7 +1183,7 @@ class State():
percent=50,
vertical=False,
target_id=self.panes['Source'],
text=std.color_string(
text=ansi.color_string(
['Destination', '' if dest_exists else ' (Missing)', '\n', dest_str],
['BLUE', 'RED', None, None],
sep='',
@ -1195,7 +1197,7 @@ def build_block_pair_report(block_pairs, settings):
report = []
notes = []
if block_pairs:
report.append(std.color_string('Block Pairs', 'GREEN'))
report.append(ansi.color_string('Block Pairs', 'GREEN'))
else:
# Bail early
return report
@ -1214,7 +1216,7 @@ def build_block_pair_report(block_pairs, settings):
if settings:
if not settings['First Run']:
notes.append(
std.color_string(
ansi.color_string(
['NOTE:', 'Clone settings loaded from previous run.'],
['BLUE', None],
),
@ -1222,14 +1224,14 @@ def build_block_pair_report(block_pairs, settings):
if settings['Needs Format'] and settings['Table Type']:
msg = f'Destination will be formatted using {settings["Table Type"]}'
notes.append(
std.color_string(
ansi.color_string(
['NOTE:', msg],
['BLUE', None],
),
)
if any(pair.get_rescued_size() > 0 for pair in block_pairs):
notes.append(
std.color_string(
ansi.color_string(
['NOTE:', 'Resume data loaded from map file(s).'],
['BLUE', None],
),
@ -1311,12 +1313,12 @@ def build_directory_report(path):
for line in proc.stdout.splitlines():
line = line.replace('\n', '')
if 'FSTYPE' in line:
line = std.color_string(f'{"PATH":<{width}}{line}', 'BLUE')
line = ansi.color_string(f'{"PATH":<{width}}{line}', 'BLUE')
else:
line = f'{path:<{width}}{line}'
report.append(line)
else:
report.append(std.color_string('PATH', 'BLUE'))
report.append(ansi.color_string('PATH', 'BLUE'))
report.append(str(path))
# Done
@ -1352,7 +1354,7 @@ def build_disk_report(dev):
# Partition details
report.append(
std.color_string(
ansi.color_string(
(
f'{"NAME":<{widths["name"]}}'
f'{" " if dev.children else ""}'
@ -1397,8 +1399,8 @@ def build_disk_report(dev):
def build_main_menu():
"""Build main menu, returns wk.std.Menu."""
menu = std.Menu(title=std.color_string('ddrescue TUI: Main Menu', 'GREEN'))
"""Build main menu, returns wk.ui.cli.Menu."""
menu = ui.Menu(title=ansi.color_string('ddrescue TUI: Main Menu', 'GREEN'))
menu.separator = ' '
# Add actions, options, etc
@ -1429,23 +1431,25 @@ def build_object_report(obj):
def build_settings_menu(silent=True):
"""Build settings menu, returns wk.std.Menu."""
"""Build settings menu, returns wk.ui.cli.Menu."""
title_text = [
std.color_string('ddrescue TUI: Expert Settings', 'GREEN'),
ansi.color_string('ddrescue TUI: Expert Settings', 'GREEN'),
' ',
std.color_string(
ansi.color_string(
['These settings can cause', 'MAJOR DAMAGE', 'to drives'],
['YELLOW', 'RED', 'YELLOW'],
),
'Please read the manual before making changes',
]
menu = std.Menu(title='\n'.join(title_text))
menu = ui.Menu(title='\n'.join(title_text))
menu.separator = ' '
preset = 'Default'
if not silent:
# Ask which preset to use
print(f'Available ddrescue presets: {" / ".join(SETTING_PRESETS)}')
preset = std.choice(SETTING_PRESETS, 'Please select a preset:')
ui.print_standard(
f'Available ddrescue presets: {" / ".join(SETTING_PRESETS)}'
)
preset = ui.choice('Please select a preset:', SETTING_PRESETS)
# Fix selection
for _p in SETTING_PRESETS:
@ -1494,7 +1498,7 @@ def build_sfdisk_partition_line(table_type, dev_path, size, details):
# Safety Check
if not dest_type:
std.print_error(f'Failed to determine partition type for: {dev_path}')
ui.print_error(f'Failed to determine partition type for: {dev_path}')
raise std.GenericAbort()
# Add extra details
@ -1576,7 +1580,7 @@ def format_status_string(status, width):
# Add color if necessary
if color:
status_str = std.color_string(status_str, color)
status_str = ansi.color_string(status_str, color)
# Done
return status_str
@ -1698,8 +1702,8 @@ def get_object(path):
# Child/Parent check
if obj.parent:
std.print_warning(f'"{obj.path}" is a child device')
if std.ask(f'Use parent device "{obj.parent}" instead?'):
ui.print_warning(f'"{obj.path}" is a child device')
if ui.ask(f'Use parent device "{obj.parent}" instead?'):
obj = hw_disk.Disk(obj.parent)
elif path.is_dir():
obj = path
@ -1710,7 +1714,7 @@ def get_object(path):
# Abort if obj not set
if not obj:
std.print_error(f'Invalid source/dest path: {path}')
ui.print_error(f'Invalid source/dest path: {path}')
raise std.GenericAbort()
# Done
@ -1775,7 +1779,7 @@ def get_table_type(disk_path):
# Check type
if table_type not in ('GPT', 'MBR'):
std.print_error(f'Unsupported partition table type: {table_type}')
ui.print_error(f'Unsupported partition table type: {table_type}')
raise std.GenericAbort()
# Done
@ -1784,30 +1788,20 @@ def get_table_type(disk_path):
def get_working_dir(mode, destination, force_local=False):
"""Get working directory using mode and destination, returns path."""
ticket_id = None
ticket_id = ui.get_ticket_id()
working_dir = None
# Set ticket ID
while ticket_id is None:
ticket_id = std.input_text(
prompt='Please enter ticket ID:',
allow_empty_response=False,
)
ticket_id = ticket_id.replace(' ', '_')
if not re.match(r'^\d+', ticket_id):
ticket_id = None
# Use preferred path if possible
if mode == 'Image':
try:
path = pathlib.Path(destination).resolve()
except TypeError as err:
std.print_error(f'Invalid destination: {destination}')
ui.print_error(f'Invalid destination: {destination}')
raise std.GenericAbort() from err
if path.exists() and fstype_is_ok(path, map_dir=False):
working_dir = path
elif mode == 'Clone' and not force_local:
std.print_info('Mounting backup shares...')
ui.print_info('Mounting backup shares...')
net.mount_backup_shares(read_write=True)
for server in cfg.net.BACKUP_SERVERS:
path = pathlib.Path(
@ -1851,11 +1845,11 @@ def is_missing_source_or_destination(state):
if hasattr(item, 'path'):
if not item.path.exists():
missing = True
std.print_error(f'{name} disappeared')
ui.print_error(f'{name} disappeared')
elif hasattr(item, 'exists'):
if not item.exists():
missing = True
std.print_error(f'{name} disappeared')
ui.print_error(f'{name} disappeared')
else:
LOG.error('Unknown %s type: %s', name, item)
@ -1887,7 +1881,7 @@ def source_or_destination_changed(state):
# Done
if changed:
std.print_error('Source and/or Destination changed')
ui.print_error('Source and/or Destination changed')
return changed
@ -1910,7 +1904,7 @@ def main():
state.init_recovery(args)
except (FileNotFoundError, std.GenericAbort):
is_missing_source_or_destination(state)
std.abort()
ui.abort()
# Show menu
while True:
@ -1928,18 +1922,18 @@ def main():
# Detect drives
if 'Detect drives' in selection[0]:
std.clear_screen()
std.print_warning(DETECT_DRIVES_NOTICE)
if std.ask('Are you sure you proceed?'):
std.print_standard('Forcing controllers to rescan for devices...')
ui.clear_screen()
ui.print_warning(DETECT_DRIVES_NOTICE)
if ui.ask('Are you sure you proceed?'):
ui.print_standard('Forcing controllers to rescan for devices...')
cmd = 'echo "- - -" | sudo tee /sys/class/scsi_host/host*/scan'
exe.run_program(cmd, check=False, shell=True)
if source_or_destination_changed(state):
std.abort()
ui.abort()
# Start recovery
if 'Start' in selection:
std.clear_screen()
ui.clear_screen()
run_recovery(state, main_menu, settings_menu, dry_run=args['--dry-run'])
# Quit
@ -1949,14 +1943,14 @@ def main():
break
# Recovey < 100%
std.print_warning('Recovery is less than 100%')
if std.ask('Are you sure you want to quit?'):
ui.print_warning('Recovery is less than 100%')
if ui.ask('Are you sure you want to quit?'):
break
# Save results to log
LOG.info('')
for line in state.generate_report():
LOG.info(' %s', std.strip_colors(line))
LOG.info(' %s', ansi.strip_colors(line))
def mount_raw_image(path):
@ -1970,7 +1964,7 @@ def mount_raw_image(path):
# Check
if not loopback_path:
std.print_error(f'Failed to mount image: {path}')
ui.print_error(f'Failed to mount image: {path}')
# Register unmount atexit
atexit.register(unmount_loopback_device, loopback_path)
@ -2037,7 +2031,7 @@ def run_ddrescue(state, block_pair, pass_name, settings, dry_run=True):
cmd = build_ddrescue_cmd(block_pair, pass_name, settings)
poweroff_source_after_idle = True
state.update_progress_pane('Active')
std.clear_screen()
ui.clear_screen()
warning_message = ''
def _poweroff_source_drive(idle_minutes):
@ -2056,8 +2050,8 @@ def run_ddrescue(state, block_pair, pass_name, settings, dry_run=True):
return
if i % 600 == 0 and i > 0:
if i == 600:
std.print_standard(' ', flush=True)
std.print_warning(
ui.print_standard(' ', flush=True)
ui.print_warning(
f'Powering off source in {int((idle_minutes*60-i)/60)} minutes...',
)
std.sleep(5)
@ -2067,10 +2061,10 @@ def run_ddrescue(state, block_pair, pass_name, settings, dry_run=True):
cmd = ['sudo', 'hdparm', '-Y', source_dev]
proc = exe.run_program(cmd, check=False)
if proc.returncode:
std.print_error(f'Failed to poweroff source {source_dev}')
ui.print_error(f'Failed to poweroff source {source_dev}')
else:
std.print_warning(f'Powered off source {source_dev}')
std.print_standard(
ui.print_warning(f'Powered off source {source_dev}')
ui.print_standard(
'Press Enter to return to main menu...', end='', flush=True,
)
@ -2080,7 +2074,7 @@ def run_ddrescue(state, block_pair, pass_name, settings, dry_run=True):
now = datetime.datetime.now(tz=TIMEZONE).strftime('%Y-%m-%d %H:%M %Z')
with open(f'{state.log_dir}/smart.out', 'w', encoding='utf-8') as _f:
_f.write(
std.color_string(
ansi.color_string(
['SMART Attributes', f'Updated: {now}\n'],
['BLUE', 'YELLOW'],
sep='\t\t',
@ -2113,7 +2107,7 @@ def run_ddrescue(state, block_pair, pass_name, settings, dry_run=True):
if warning_message:
# Error detected on destination, stop recovery
exe.stop_process(proc)
std.print_error(warning_message)
ui.print_error(warning_message)
break
if _i % 60 == 0:
@ -2159,19 +2153,19 @@ def run_ddrescue(state, block_pair, pass_name, settings, dry_run=True):
warning_message = 'Error(s) encountered, see message above'
state.update_top_panes()
if warning_message:
print(' ')
print(' ')
std.print_error('DDRESCUE PROCESS HALTED')
print(' ')
std.print_warning(warning_message)
ui.print_standard(' ')
ui.print_standard(' ')
ui.print_error('DDRESCUE PROCESS HALTED')
ui.print_standard(' ')
ui.print_warning(warning_message)
# Needs attention?
if str(proc.poll()) != '0':
state.update_progress_pane('NEEDS ATTENTION')
std.pause('Press Enter to return to main menu...')
ui.pause('Press Enter to return to main menu...')
# Stop source poweroff countdown
std.print_standard('Stopping device poweroff countdown...', flush=True)
ui.print_standard('Stopping device poweroff countdown...', flush=True)
poweroff_source_after_idle = False
poweroff_thread.join()
@ -2187,12 +2181,12 @@ def run_recovery(state, main_menu, settings_menu, dry_run=True):
# Bail early
if is_missing_source_or_destination(state):
std.print_standard('')
std.pause('Press Enter to return to main menu...')
ui.print_standard('')
ui.pause('Press Enter to return to main menu...')
return
if source_or_destination_changed(state):
std.print_standard('')
std.abort()
ui.print_standard('')
ui.abort()
# Get settings
for name, details in main_menu.toggles.items():
@ -2248,9 +2242,9 @@ def run_recovery(state, main_menu, settings_menu, dry_run=True):
# Show warning if nothing was done
if not attempted_recovery:
std.print_warning('No actions performed')
std.print_standard(' ')
std.pause('Press Enter to return to main menu...')
ui.print_warning('No actions performed')
ui.print_standard(' ')
ui.pause('Press Enter to return to main menu...')
# Done
state.save_debug_reports()
@ -2258,12 +2252,12 @@ def run_recovery(state, main_menu, settings_menu, dry_run=True):
state.update_progress_pane('Idle')
def select_disk(prompt, skip_disk=None):
def select_disk(prompt_msg, skip_disk=None):
"""Select disk from list, returns Disk()."""
std.print_info('Scanning disks...')
ui.print_info('Scanning disks...')
disks = hw_disk.get_disks()
menu = std.Menu(
title=std.color_string(f'ddrescue TUI: {prompt} Selection', 'GREEN'),
menu = ui.Menu(
title=ansi.color_string(f'ddrescue TUI: {prompt_msg} Selection', 'GREEN'),
)
menu.disabled_str = 'Already selected'
menu.separator = ' '
@ -2303,11 +2297,11 @@ def select_disk(prompt, skip_disk=None):
return selected_disk
def select_disk_parts(prompt, disk):
def select_disk_parts(prompt_msg, disk):
"""Select disk parts from list, returns list of Disk()."""
title = std.color_string('ddrescue TUI: Partition Selection', 'GREEN')
title = ansi.color_string('ddrescue TUI: Partition Selection', 'GREEN')
title += f'\n\nDisk: {disk.path} {disk.description}'
menu = std.Menu(title)
menu = ui.Menu(title)
menu.separator = ' '
menu.add_action('All')
menu.add_action('None')
@ -2319,7 +2313,7 @@ def select_disk_parts(prompt, disk):
"""Loop over selection menu until at least one partition selected."""
while True:
selection = menu.advanced_select(
f'Please select the parts to {prompt.lower()}: ',
f'Please select the parts to {prompt_msg.lower()}: ',
)
if 'All' in selection:
for option in menu.options.values():
@ -2356,7 +2350,7 @@ def select_disk_parts(prompt, disk):
if not menu.options:
menu.add_option(whole_disk_str, {'Selected': True, 'Path': disk.path})
menu.title += '\n\n'
menu.title += std.color_string(' No partitions detected.', 'YELLOW')
menu.title += ansi.color_string(' No partitions detected.', 'YELLOW')
# Get selection
_select_parts(menu)
@ -2369,25 +2363,25 @@ def select_disk_parts(prompt, disk):
# Check if whole disk selected
if len(object_list) == len(disk.children):
# NOTE: This is not true if the disk has no partitions
msg = f'Preserve partition table and unused space in {prompt.lower()}?'
if std.ask(msg):
msg = f'Preserve partition table and unused space in {prompt_msg.lower()}?'
if ui.ask(msg):
# Replace part list with whole disk obj
object_list = [disk.path]
# Convert object_list to hw_disk.Disk() objects
print(' ')
std.print_info('Getting disk/partition details...')
ui.print_standard(' ')
ui.print_info('Getting disk/partition details...')
object_list = [hw_disk.Disk(path) for path in object_list]
# Done
return object_list
def select_path(prompt):
def select_path(prompt_msg):
"""Select path, returns pathlib.Path."""
invalid = False
menu = std.Menu(
title=std.color_string(f'ddrescue TUI: {prompt} Path Selection', 'GREEN'),
menu = ui.Menu(
title=ansi.color_string(f'ddrescue TUI: {prompt_msg} Path Selection', 'GREEN'),
)
menu.separator = ' '
menu.add_action('Quit')
@ -2400,7 +2394,7 @@ def select_path(prompt):
if 'Current directory' in selection:
path = os.getcwd()
elif 'Enter manually' in selection:
path = std.input_text('Please enter path: ')
path = ui.input_text('Please enter path: ')
elif 'Quit' in selection:
raise std.GenericAbort()
@ -2410,7 +2404,7 @@ def select_path(prompt):
except TypeError:
invalid = True
if invalid or not path.is_dir():
std.print_error(f'Invalid path: {path}')
ui.print_error(f'Invalid path: {path}')
raise std.GenericAbort()
# Done
@ -2429,7 +2423,7 @@ def set_mode(docopt_args):
# Ask user if necessary
if not mode:
answer = std.choice(['C', 'I'], 'Are we cloning or imaging?')
answer = ui.choice('Are we cloning or imaging?', ['C', 'I'])
if answer == 'C':
mode = 'Clone'
else:

121
scripts/wk/debug.py
View file

@ -1,6 +1,21 @@
"""WizardKit: Debug Functions"""
# vim: sts=2 sw=2 ts=2
import inspect
import logging
import lzma
import os
import pickle
import platform
import re
import socket
import sys
import time
import requests
from wk.cfg.net import CRASH_SERVER
from wk.log import get_log_filepath, get_root_logger_path
# Classes
class Debug():
@ -10,11 +25,55 @@ class Debug():
# STATIC VARIABLES
LOG = logging.getLogger(__name__)
DEBUG_CLASS = Debug()
METHOD_TYPE = type(DEBUG_CLASS.method)
# Functions
def generate_debug_report():
"""Generate debug report, returns str."""
platform_function_list = (
'architecture',
'machine',
'platform',
'python_version',
)
report = []
# Logging data
log_path = get_log_filepath()
if log_path:
report.append('------ Start Log -------')
report.append('')
with open(log_path, 'r', encoding='utf-8') as log_file:
report.extend(log_file.read().splitlines())
report.append('')
report.append('------- End Log --------')
# System
report.append('--- Start debug info ---')
report.append('')
report.append('[System]')
report.append(f' {"FQDN":<24} {socket.getfqdn()}')
for func in platform_function_list:
func_name = func.replace('_', ' ').capitalize()
func_result = getattr(platform, func)()
report.append(f' {func_name:<24} {func_result}')
report.append(f' {"Python sys.argv":<24} {sys.argv}')
report.append('')
# Environment
report.append('[Environment Variables]')
for key, value in sorted(os.environ.items()):
report.append(f' {key:<24} {value}')
report.append('')
# Done
report.append('---- End debug info ----')
return '\n'.join(report)
def generate_object_report(obj, indent=0):
"""Generate debug report for obj, returns list."""
report = []
@ -46,5 +105,67 @@ def generate_object_report(obj, indent=0):
return report
def save_pickles(obj_dict, out_path=None):
"""Save dict of objects using pickle."""
LOG.info('Saving pickles')
# Set path
if not out_path:
out_path = 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, compress=True, reason='DEBUG'):
"""Upload debug report to CRASH_SERVER as specified in wk.cfg.main."""
LOG.info('Uploading debug report to %s', CRASH_SERVER.get('Name', '?'))
headers = CRASH_SERVER.get('Headers', {'X-Requested-With': 'XMLHttpRequest'})
if compress:
headers['Content-Type'] = 'application/octet-stream'
# Check if the required server details are available
if not all(CRASH_SERVER.get(key, False) for key in ('Name', 'Url', 'User')):
msg = 'Server details missing, aborting upload.'
print(msg)
raise RuntimeError(msg)
# Set filename (based on the logging config if possible)
filename = 'Unknown'
log_path = get_log_filepath()
if log_path:
# Strip everything but the prefix
filename = re.sub(r'^(.*)_(\d{4}-\d{2}-\d{2}.*)', r'\1', log_path.name)
filename = f'{filename}_{reason}_{time.strftime("%Y-%m-%d_%H%M%S%z")}.log'
LOG.debug('filename: %s', filename)
# Compress report
if compress:
filename += '.xz'
xz_report = lzma.compress(report.encode('utf8'))
# Upload report
url = f'{CRASH_SERVER["Url"]}/{filename}'
response = requests.put(
url,
auth=(CRASH_SERVER['User'], CRASH_SERVER.get('Pass', '')),
data=xz_report if compress else report,
headers=headers,
timeout=60,
)
# Check response
if not response.ok:
raise RuntimeError('Failed to upload report')
if __name__ == '__main__':
print("This file is not meant to be called directly.")

6
scripts/wk/exe.py
View file

@ -1,5 +1,5 @@
"""WizardKit: Execution functions"""
#vim: sts=2 sw=2 ts=2
# vim: sts=2 sw=2 ts=2
import json
import logging
@ -106,8 +106,8 @@ def build_cmd_kwargs(cmd, minimized=False, pipe=True, shell=False, **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

14
scripts/wk/graph.py
View file

@ -3,7 +3,7 @@
import logging
from wk.std import color_string
from wk.ui import ansi
# STATIC VARIABLES
@ -52,27 +52,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:
@ -128,7 +128,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}',

13
scripts/wk/hw/benchmark.py
View file

@ -21,11 +21,8 @@ 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
# STATIC VARIABLES
@ -116,7 +113,7 @@ def check_io_results(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)
@ -154,7 +151,7 @@ def run_io_test(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:
@ -162,7 +159,7 @@ def run_io_test(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

28
scripts/wk/hw/cpu.py
View file

@ -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,7 +111,7 @@ def start_mprime(working_dir, log_path) -> subprocess.Popen:
stdin=proc_mprime.stdout,
stdout=subprocess.PIPE,
)
proc_mprime.stdout.close()
proc_mprime.stdout.close() # type: ignore[reportOptionalMemberAccess]
save_nsbr = exe.NonBlockingStreamReader(proc_grep.stdout)
exe.start_thread(
save_nsbr.save_to_file,
@ -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)

317
scripts/wk/hw/diags.py
View file

@ -10,7 +10,7 @@ import time
from docopt import docopt
from wk import cfg, debug, exe, log, std, tmux
from wk import cfg, debug, exe, log, std
from wk.cfg.hw import STATUS_COLORS
from wk.hw import benchmark as hw_benchmark
from wk.hw import cpu as hw_cpu
@ -25,6 +25,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
@ -81,37 +83,25 @@ class State():
"""Object for tracking hardware diagnostic data."""
def __init__(self, test_mode=False):
self.disks = []
self.layout = cfg.hw.TMUX_LAYOUT.copy()
self.log_dir = None
self.panes = {}
self.progress_file = None
self.system = None
self.test_groups = []
self.top_text = std.color_string('Hardware Diagnostics', 'GREEN')
self.title_text = 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(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)."""
@ -128,30 +118,11 @@ class State():
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()
# Set log
@ -166,15 +137,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()
@ -216,35 +185,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')
@ -253,7 +193,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)))
@ -291,27 +231,15 @@ 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)}}'],
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='',
))
@ -320,19 +248,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:
@ -398,7 +324,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')
@ -410,25 +336,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)
@ -446,17 +371,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,
)
@ -467,16 +392,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)
@ -493,18 +421,18 @@ 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)
# 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:
@ -522,7 +450,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(
@ -532,14 +460,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=' ',
)
state.ui.set_current_pane_height(10)
for test in test_objects:
if test.disabled:
# Skip
@ -551,16 +475,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(test, test_log, test_mode=test_mode)
except KeyboardInterrupt:
@ -569,20 +495,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:
@ -594,13 +521,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:
@ -614,12 +540,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_cmd='tail', 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):
@ -634,10 +558,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:
@ -691,13 +613,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)
@ -713,20 +634,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
@ -737,13 +651,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:
@ -761,7 +673,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'])
@ -787,15 +698,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:
@ -819,7 +730,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 print_countdown(proc, seconds) -> None:
@ -858,8 +769,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
# Run tests
@ -870,13 +781,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['Skip USB Benchmarks']['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
@ -897,48 +808,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_report())
std.print_standard(' ')
cli.print_success('CPU:')
cli.print_report(state.system.generate_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:

13
scripts/wk/hw/disk.py
View file

@ -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
@ -73,7 +74,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 +83,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 +99,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}')

5
scripts/wk/hw/keyboard.py
View file

@ -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:

14
scripts/wk/hw/network.py
View file

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

6
scripts/wk/hw/screensavers.py
View file

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

9
scripts/wk/hw/sensors.py
View file

@ -12,7 +12,8 @@ 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
@ -109,7 +110,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 +426,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 +436,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)

19
scripts/wk/hw/smart.py
View file

@ -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,26 @@ 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')]
report = [ansi.color_string('Self-Test', 'BLUE')]
test_details = get_smart_self_test_details(test_obj.dev)
test_result = test_details.get('status', {}).get('string', 'Unknown')
# 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'))
report.append(ansi.color_string(f' {test_obj.status}', 'YELLOW'))
elif test_obj.status == 'TestInProgress':
report.append(color_string(' Failed to stop previous test', 'RED'))
report.append(ansi.color_string(' Failed to stop previous test', 'RED'))
test_obj.set_status('Failed')
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'))
report.append(ansi.color_string(' Aborted', 'YELLOW'))
test_obj.set_status('Aborted')
elif test_obj.status == 'TimedOut':
report.append(color_string(' TimedOut', 'YELLOW'))
report.append(ansi.color_string(' TimedOut', 'YELLOW'))
# Done
test_obj.report.extend(report)
@ -136,7 +137,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'],
)
@ -324,7 +325,7 @@ def run_smart_self_test(test_obj, log_path) -> None:
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='',

18
scripts/wk/hw/surface_scan.py
View file

@ -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
@ -31,7 +27,7 @@ def check_surface_scan_results(test_obj, log_path) -> None:
"""Check results and set test status."""
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 or BADBLOCKS_SKIP_REGEX.match(line):
# Skip
continue
@ -48,10 +44,10 @@ def check_surface_scan_results(test_obj, log_path) -> None:
test_obj.set_status('Passed')
else:
test_obj.failed = True
test_obj.report.append(f' {color_string(line, "YELLOW")}')
test_obj.report.append(f' {ansi.color_string(line, "YELLOW")}')
test_obj.set_status('Failed')
else:
test_obj.report.append(f' {color_string(line, "YELLOW")}')
test_obj.report.append(f' {ansi.color_string(line, "YELLOW")}')
if not (test_obj.passed or test_obj.failed):
test_obj.set_status('Unknown')
@ -65,7 +61,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
@ -84,7 +80,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='',

12
scripts/wk/hw/system.py
View file

@ -12,12 +12,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
@ -41,11 +37,11 @@ class System:
def generate_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

32
scripts/wk/kit/build_win.py
View file

@ -21,16 +21,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
@ -245,7 +237,7 @@ def 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_neutron():
@ -315,7 +307,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)
@ -459,13 +451,13 @@ def build_kit():
"""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'):
@ -496,15 +488,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,
@ -514,7 +506,7 @@ def build_kit():
# Done
print('')
print('Done.')
pause('Press Enter to exit...')
ui.pause('Press Enter to exit...')
if __name__ == '__main__':

90
scripts/wk/kit/ufd.py
View file

@ -10,7 +10,7 @@ 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
@ -93,10 +95,10 @@ def build_ufd():
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
@ -104,9 +106,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'])
@ -118,7 +120,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,
@ -170,8 +172,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,
@ -187,8 +189,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}...',
@ -203,8 +205,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,
@ -235,8 +237,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,
@ -245,30 +247,30 @@ 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):
"""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):
@ -481,12 +483,12 @@ def show_selections(args, sources, ufd_dev, ufd_sources, extra_images):
"""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'],
)
@ -498,15 +500,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',
@ -514,14 +516,14 @@ 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):
@ -621,11 +623,11 @@ 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()
ui.print_error(f'ERROR: {label} not found: {s_path}')
ui.abort()
if not is_valid_path(s_path_obj, data['Type']):
std.print_error(f'ERROR: Invalid {label} source: {s_path}')
std.abort()
ui.print_error(f'ERROR: Invalid {label} source: {s_path}')
ui.abort()
sources[label] = s_path_obj
return sources
@ -638,12 +640,12 @@ def verify_ufd(dev_path):
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

18
scripts/wk/log.py
View file

@ -61,6 +61,24 @@ def format_log_path(
return log_path
def get_log_filepath():
"""Get the log filepath from the root logger, returns pathlib.Path obj.
NOTE: This will use the first handler baseFilename it finds (if any).
"""
log_filepath = None
root_logger = logging.getLogger()
# Check handlers
for handler in root_logger.handlers:
if hasattr(handler, 'baseFilename'):
log_filepath = pathlib.Path(handler.baseFilename).resolve()
break
# Done
return log_filepath
def get_root_logger_path():
"""Get path to log file from root logger, returns pathlib.Path obj."""
log_path = None

5
scripts/wk/net.py
View file

@ -8,9 +8,10 @@ import re
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
@ -199,7 +200,7 @@ def show_valid_addresses():
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():

9
scripts/wk/os/linux.py
View file

@ -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
@ -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],
)

16
scripts/wk/os/win.py
View file

@ -30,10 +30,10 @@ from wk.std import (
GenericError,
GenericWarning,
bytes_to_string,
color_string,
input_text,
sleep,
)
from wk.ui import cli as ui
from wk.ui import ansi
# STATIC VARIABLES
@ -183,7 +183,7 @@ def check_4k_alignment(show_alert=False):
continue
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)})'
,
@ -199,7 +199,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
@ -212,7 +212,7 @@ def export_bitlocker_info():
]
# Get filename
file_name = input_text(prompt='Enter filename', allow_empty_response=False)
file_name = ui.input_text(prompt_msg='Enter filename')
file_path = pathlib.Path(f'../../Bitlocker_{file_name}.txt').resolve()
# Save info
@ -251,13 +251,13 @@ 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
@ -364,7 +364,7 @@ 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

79
scripts/wk/repairs/win.py
View file

@ -58,21 +58,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
@ -87,7 +76,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(
@ -105,7 +94,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'):
@ -116,7 +105,7 @@ for error in ('CalledProcessError', 'FileNotFoundError'):
def build_menus(base_menus, title, presets):
"""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']:
@ -125,7 +114,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')
@ -135,7 +124,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)
@ -167,7 +156,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)
@ -271,7 +260,7 @@ 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)
@ -316,7 +305,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
@ -392,20 +381,20 @@ def load_settings(menus):
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):
"""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
@ -423,21 +412,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()
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():
@ -446,19 +435,19 @@ def run_auto_repairs(base_menus, presets):
try:
run_group(group, menu)
except KeyboardInterrupt:
abort()
ui.abort()
# Done
end_session()
print_info('Done')
pause('Press Enter to exit...')
ui.print_info('Done')
ui.pause('Press Enter to exit...')
def run_group(group, menu):
"""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)
@ -472,7 +461,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
@ -482,7 +471,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,
)
@ -490,7 +479,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
@ -513,7 +502,7 @@ def save_selection_settings(menus):
def save_settings(group, name, result=None, **kwargs):
"""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:
@ -527,7 +516,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):
@ -891,7 +880,7 @@ def backup_all_browser_profiles(use_try_print=False):
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)
@ -905,7 +894,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,
)
@ -933,7 +922,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,
)
@ -1319,7 +1308,7 @@ def kill_explorer():
def reboot(timeout=10):
"""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)

82
scripts/wk/setup/win.py
View file

@ -60,22 +60,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
@ -91,7 +79,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)',
@ -103,7 +91,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'):
@ -114,7 +102,7 @@ for error in ('CalledProcessError', 'FileNotFoundError'):
def build_menus(base_menus, title, presets):
"""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']:
@ -124,7 +112,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')
@ -153,7 +141,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)
@ -170,26 +158,26 @@ def check_os_and_set_menu_title(title):
"""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):
@ -215,10 +203,10 @@ 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' and ask('Install LibreOffice?'):
if selection[0] == 'Default' and ui.ask('Install LibreOffice?'):
menus['Install Software'].options['LibreOffice']['Selected'] = True
# Re-enable Main Menu action if disabled
@ -235,11 +223,11 @@ def run_auto_setup(base_menus, presets):
"""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)
@ -254,10 +242,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():
@ -266,22 +254,22 @@ 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):
"""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
@ -418,7 +406,7 @@ 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)
@ -804,8 +792,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()
@ -906,7 +894,7 @@ def get_storage_status():
"""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

976
scripts/wk/std.py
View file

File diff suppressed because it is too large Load diff

6
scripts/wk/ui/__init__.py Normal file
View file

@ -0,0 +1,6 @@
"""WizardKit: ui module init"""
from . import ansi
from . import cli
from . import tmux
from . import tui

72
scripts/wk/ui/ansi.py Normal file
View file

@ -0,0 +1,72 @@
"""WizardKit: ANSI control/escape functions"""
# vim: sts=2 sw=2 ts=2
import itertools
import logging
import pathlib
# 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():
"""Clear screen using ANSI escape."""
print('\033c', end='', flush=True)
def color_string(strings, colors, sep=' '):
"""Build colored string using ANSI escapes, returns str."""
clear_code = COLORS['CLEAR']
msg = []
# Convert to tuples if necessary
if isinstance(strings, (str, pathlib.Path)):
strings = (strings,)
if isinstance(colors, (str, pathlib.Path)):
colors = (colors,)
# Convert to strings if necessary
try:
iter(strings)
except TypeError:
# Assuming single element passed, convert to string
strings = (str(strings),)
try:
iter(colors)
except TypeError:
# Assuming single element passed, convert to string
colors = (str(colors),)
# Build new string with color escapes added
for string, color in itertools.zip_longest(strings, colors):
color_code = COLORS.get(color, clear_code)
msg.append(f'{color_code}{string}{clear_code}')
# Done
return sep.join(msg)
def strip_colors(string):
"""Strip known ANSI color escapes from string, returns str."""
LOG.debug('string: %s', string)
for color in COLORS.values():
string = string.replace(color, '')
return string
if __name__ == '__main__':
print("This file is not meant to be called directly.")

880
scripts/wk/ui/cli.py Normal file
View file

@ -0,0 +1,880 @@
"""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 collections import OrderedDict
from prompt_toolkit import prompt
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, allow_empty=False):
self.allow_empty = allow_empty
self.choices = [str(c).upper() for c in choices]
super().__init__()
def validate(self, document):
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):
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=False):
self.allow_empty = allow_empty
super().__init__()
def validate(self, document):
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=False):
self.allow_empty = allow_empty
super().__init__()
def validate(self, document):
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.
Menu items are added to an OrderedDict so the order is preserved.
ASSUMPTIONS:
1. All entry names are unique.
2. All action entry names start with different letters.
"""
def __init__(self, title='[Untitled Menu]'):
self.actions = OrderedDict()
self.options = OrderedDict()
self.sets = OrderedDict()
self.toggles = OrderedDict()
self.disabled_str = 'Disabled'
self.separator = '─'
self.title = title
def _generate_menu_text(self):
"""Generate menu text, returns str."""
separator_string = self._get_separator_string()
menu_lines = [self.title, separator_string] if self.title else []
# Sets & toggles
for section in (self.sets, self.toggles):
for details in section.values():
if details.get('Hidden', False):
continue
if details.get('Separator', False):
menu_lines.append(separator_string)
menu_lines.append(details['Display Name'])
if self.sets or self.toggles:
menu_lines.append(separator_string)
# Options
for details in self.options.values():
if details.get('Hidden', False):
continue
if details.get('Separator', False):
menu_lines.append(separator_string)
menu_lines.append(details['Display Name'])
if self.options:
menu_lines.append(separator_string)
# Actions
for details in self.actions.values():
if details.get('Hidden', False):
continue
if details.get('Separator', False):
menu_lines.append(separator_string)
menu_lines.append(details['Display Name'])
# Show menu
menu_lines.append('')
menu_lines = [str(line) for line in menu_lines]
return '\n'.join(menu_lines)
def _get_display_name(
self, name, details,
index=None, no_checkboxes=True, setting_item=False):
"""Format display name based on details and args, returns str."""
disabled = details.get('Disabled', False)
if setting_item and not details['Selected']:
# Display item in YELLOW
disabled = True
checkmark = '*'
if 'DISPLAY' in os.environ or PLATFORM == 'Darwin':
checkmark = '✓'
display_name = f'{index if index else name[:1].upper()}: '
if not (index and index >= 10):
display_name = f' {display_name}'
if setting_item and 'Value' in details:
name = f'{name} = {details["Value"]}'
# Add enabled status if necessary
if not no_checkboxes:
display_name += f'[{checkmark if details["Selected"] else " "}] '
# Add name
if disabled:
display_name += color_string(f'{name} ({self.disabled_str})', 'YELLOW')
else:
display_name += name
# Done
return display_name
def _get_separator_string(self):
"""Format separator length based on name lengths, returns str."""
separator_length = 0
# Check title line(s)
if self.title:
for line in self.title.split('\n'):
separator_length = max(separator_length, len(strip_colors(line)))
# Loop over all item names
for section in (self.actions, self.options, self.sets, self.toggles):
for details in section.values():
if details.get('Hidden', False):
# Skip hidden lines
continue
line = strip_colors(details['Display Name'])
separator_length = max(separator_length, len(line))
separator_length += 1
# Done
return self.separator * separator_length
def _get_valid_answers(self):
"""Get valid answers based on menu items, returns list."""
valid_answers = []
# Numbered items
index = 0
for section in (self.sets, self.toggles, self.options):
for details in section.values():
if details.get('Hidden', False):
# Don't increment index or add to valid_answers
continue
index += 1
if not details.get('Disabled', False):
valid_answers.append(str(index))
# Action items
for name, details in self.actions.items():
if not details.get('Disabled', False):
valid_answers.append(name[:1].upper())
# Done
return valid_answers
def _resolve_selection(self, selection):
"""Get menu item based on user selection, returns tuple."""
offset = 1
resolved_selection = None
if selection.isnumeric():
# Enumerate over numbered entries
entries = [
*self.sets.items(),
*self.toggles.items(),
*self.options.items(),
]
for _i, details in enumerate(entries):
if details[1].get('Hidden', False):
offset -= 1
elif str(_i+offset) == selection:
resolved_selection = (details)
break
else:
# Just check actions
for action, details in self.actions.items():
if action.lower().startswith(selection.lower()):
resolved_selection = (action, details)
break
# Done
return resolved_selection
def _update(self, single_selection=True, settings_mode=False):
"""Update menu items in preparation for printing to screen."""
index = 0
# Fix selection status for sets
for set_details in self.sets.values():
set_selected = True
set_targets = set_details['Targets']
for option, option_details in self.options.items():
if option in set_targets and not option_details['Selected']:
set_selected = False
elif option not in set_targets and option_details['Selected']:
set_selected = False
set_details['Selected'] = set_selected
# Numbered sections
for section in (self.sets, self.toggles, self.options):
for name, details in section.items():
if details.get('Hidden', False):
# Skip hidden lines and don't increment index
continue
index += 1
details['Display Name'] = self._get_display_name(
name,
details,
index=index,
no_checkboxes=single_selection,
setting_item=settings_mode,
)
# Actions
for name, details in self.actions.items():
details['Display Name'] = self._get_display_name(
name,
details,
no_checkboxes=True,
)
def _update_entry_selection_status(self, entry, toggle=True, status=None):
"""Update entry selection status either directly or by toggling."""
if entry in self.sets:
# Update targets not the set itself
new_status = not self.sets[entry]['Selected'] if toggle else status
targets = self.sets[entry]['Targets']
self._update_set_selection_status(targets, new_status)
for section in (self.toggles, self.options, self.actions):
if entry in section:
if toggle:
section[entry]['Selected'] = not section[entry]['Selected']
else:
section[entry]['Selected'] = status
def _update_set_selection_status(self, targets, status):
"""Select or deselect options based on targets and status."""
for option, details in self.options.items():
# If (new) status is True and this option is a target then select
# Otherwise deselect
details['Selected'] = status and option in targets
def _user_select(self, prompt_msg):
"""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, details=None):
"""Add action to menu."""
details = details if details else {}
details['Selected'] = details.get('Selected', False)
self.actions[name] = details
def add_option(self, name, details=None):
"""Add option to menu."""
details = details if details else {}
details['Selected'] = details.get('Selected', False)
self.options[name] = details
def add_set(self, name, details=None):
"""Add set to menu."""
details = details if details else {}
details['Selected'] = details.get('Selected', False)
# Safety check
if 'Targets' not in details:
raise KeyError('Menu set has no targets')
# Add set
self.sets[name] = details
def add_toggle(self, name, details=None):
"""Add toggle to menu."""
details = details if details else {}
details['Selected'] = details.get('Selected', False)
self.toggles[name] = details
def advanced_select(self, prompt_msg='Please make a selection: '):
"""Display menu and make multiple selections, returns tuple.
NOTE: Menu is displayed until an action entry is selected.
"""
while True:
self._update(single_selection=False)
user_selection = self._user_select(prompt_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='Please make a selection: '):
"""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='Please make a selection: ', update=True):
"""Display menu and make a single selection, returns tuple."""
if update:
self._update()
user_selection = self._user_select(prompt_msg)
return self._resolve_selection(user_selection)
def update(self):
"""Update menu with default settings."""
self._update()
class TryAndPrint():
"""Object used to standardize running functions and returning the result.
The errors and warning attributes are used to allow fine-tuned results
based on exception names.
"""
def __init__(self, msg_bad='FAILED', msg_good='SUCCESS'):
self.catch_all = True
self.indent = INDENT
self.list_errors = ['GenericError']
self.list_warnings = ['GenericWarning']
self.msg_bad = msg_bad
self.msg_good = msg_good
self.verbose = False
self.width = WIDTH
def _format_exception_message(self, _exception):
"""Format using the exception's args or name, returns str."""
LOG.debug(
'Formatting exception: %s, %s',
_exception.__class__.__name__,
_exception,
)
message = ''
# Format message string from _exception
try:
if isinstance(_exception, subprocess.CalledProcessError):
message = _exception.stderr
if not isinstance(message, str):
message = message.decode('utf-8')
message = message.strip()
elif isinstance(_exception, ZeroDivisionError):
# Skip and just use exception name below
pass
else:
message = str(_exception)
except Exception:
# Just use the exception name instead
pass
# Prepend exception name
if _exception.__class__.__name__ not in ('GenericError', 'GenericWarning'):
try:
message = f'{_exception.__class__.__name__}: {message}'
except Exception:
message = f'UNKNOWN ERROR: {message}'
# Fix multi-line messages
if '\n' in message:
try:
lines = [
f'{" "*(self.indent+self.width)}{line.strip()}'
for line in message.splitlines() if line.strip()
]
lines[0] = lines[0].strip()
message = '\n'.join(lines)
except Exception:
pass
# Done
return message
def _format_function_output(self, output, msg_good):
"""Format function output for use in try_and_print(), returns str."""
LOG.debug('Formatting output: %s', output)
if not output:
raise GenericWarning('No output')
# Ensure we're working with a list
if isinstance(output, subprocess.CompletedProcess):
stdout = output.stdout
if not isinstance(stdout, str):
stdout = stdout.decode('utf8')
output = stdout.strip().splitlines()
if not output:
# Going to treat these as successes (for now)
LOG.warning('Program output was empty, assuming good result.')
return color_string(msg_good, 'GREEN')
else:
try:
output = list(output)
except TypeError:
output = [output]
# Safety check
if not output:
# Going to ignore empty function output for now
LOG.error('Output is empty')
return 'UNKNOWN'
# Build result_msg
result_msg = f'{output.pop(0)}'
if output:
output = [f'{" "*(self.indent+self.width)}{line}' for line in output]
result_msg += '\n' + '\n'.join(output)
# Done
return result_msg
def _log_result(self, message, result_msg):
"""Log result text without color formatting."""
log_text = f'{" "*self.indent}{message:<{self.width}}{result_msg}'
for line in log_text.splitlines():
line = strip_colors(line)
LOG.info(line)
def add_error(self, exception_name):
"""Add exception name to error list."""
if exception_name not in self.list_errors:
self.list_errors.append(exception_name)
def add_warning(self, exception_name):
"""Add exception name to warning list."""
if exception_name not in self.list_warnings:
self.list_warnings.append(exception_name)
def run(
self, message, function, *args,
catch_all=None, msg_good=None, verbose=None, **kwargs):
"""Run a function and print the results, returns results as dict.
If catch_all is True then (nearly) all exceptions will be caught.
Otherwise if an exception occurs that wasn't specified it will be
re-raised.
If the function returns data it will be used instead of msg_good,
msg_bad, or exception text.
The output should be a list or a subprocess.CompletedProcess object.
If msg_good is passed it will override self.msg_good for this call.
If verbose is True then exception names or messages will be used for
the result message. Otherwise it will simply be set to result_bad.
If catch_all and/or verbose are passed it will override
self.catch_all and/or self.verbose for this call.
args and kwargs are passed to the function.
"""
LOG.debug('function: %s.%s', function.__module__, function.__name__)
LOG.debug('args: %s', args)
LOG.debug('kwargs: %s', kwargs)
LOG.debug(
'catch_all: %s, msg_good: %s, verbose: %s',
catch_all,
msg_good,
verbose,
)
f_exception = None
catch_all = catch_all if catch_all is not None else self.catch_all
msg_good = msg_good if msg_good is not None else self.msg_good
output = None
result_msg = 'UNKNOWN'
verbose = verbose if verbose is not None else self.verbose
# Build exception tuples
e_exceptions = tuple(get_exception(e) for e in self.list_errors)
w_exceptions = tuple(get_exception(e) for e in self.list_warnings)
# Run function and catch exceptions
print(f'{" "*self.indent}{message:<{self.width}}', end='', flush=True)
LOG.debug('Running function: %s.%s', function.__module__, function.__name__)
try:
output = function(*args, **kwargs)
except w_exceptions as _exception:
# Warnings
result_msg = self._format_exception_message(_exception)
print_warning(result_msg, log=False)
f_exception = _exception
except e_exceptions as _exception:
# Exceptions
result_msg = self._format_exception_message(_exception)
print_error(result_msg, log=False)
f_exception = _exception
except Exception as _exception:
# Unexpected exceptions
if verbose:
result_msg = self._format_exception_message(_exception)
else:
result_msg = self.msg_bad
print_error(result_msg, log=False)
f_exception = _exception
if not catch_all:
# Re-raise error as necessary
raise
else:
# Success
if output:
result_msg = self._format_function_output(output, msg_good)
print(result_msg)
else:
result_msg = msg_good
print_success(result_msg, log=False)
# Done
self._log_result(message, result_msg)
return {
'Exception': f_exception,
'Failed': bool(f_exception),
'Message': result_msg,
'Output': output,
}
# Functions
def abort(prompt_msg='Aborted.', show_prompt_msg=True, return_code=1):
"""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):
"""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'):
answer = True
elif response.upper().startswith('N'):
answer = False
# Done
LOG.info('%s%s', prompt_msg, 'Yes' if answer else 'No')
return answer
def beep(repeat=1):
"""Play system bell with optional repeat."""
while repeat >= 1:
# Print bell char without a newline
print('\a', end='', flush=True)
sleep(0.5)
repeat -= 1
def choice(prompt_msg, choices):
"""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):
"""Fix prompt, returns str."""
if not message:
message = 'Input text: '
message = str(message)
if message[-1:] != ' ':
message += ' '
return message
@cache
def get_exception(name):
"""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():
"""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='Enter text: ', allow_empty=False, validator=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():
"""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='Press Enter to continue... '):
"""Simple pause implementation."""
input_text(prompt_msg, allow_empty=True)
def print_colored(strings, colors, log=False, sep=' ', **kwargs):
"""Prints strings in the colors specified."""
LOG.debug(
'strings: %s, colors: %s, sep: %s, kwargs: %s',
strings, colors, sep, kwargs,
)
msg = color_string(strings, colors, sep=sep)
print_options = {
'end': kwargs.get('end', '\n'),
'file': kwargs.get('file', sys.stdout),
'flush': kwargs.get('flush', False),
}
print(msg, **print_options)
if log:
LOG.info(strip_colors(msg))
def print_error(msg, log=True, **kwargs):
"""Prints message in RED and log as ERROR."""
if 'file' not in kwargs:
# Only set if not specified
kwargs['file'] = sys.stderr
print_colored(msg, 'RED', **kwargs)
if log:
LOG.error(msg)
def print_info(msg, log=True, **kwargs):
"""Prints message in BLUE and log as INFO."""
print_colored(msg, 'BLUE', **kwargs)
if log:
LOG.info(msg)
def print_report(report, indent=None, log=True):
"""Print report to screen and optionally to log."""
for line in report:
if indent:
line = f'{" "*indent}{line}'
print(line)
if log:
LOG.info(strip_colors(line))
def print_standard(msg, log=True, **kwargs):
"""Prints message and log as INFO."""
print(msg, **kwargs)
if log:
LOG.info(msg)
def print_success(msg, log=True, **kwargs):
"""Prints message in GREEN and log as INFO."""
print_colored(msg, 'GREEN', **kwargs)
if log:
LOG.info(msg)
def print_warning(msg, log=True, **kwargs):
"""Prints message in YELLOW and log as WARNING."""
if 'file' not in kwargs:
# Only set if not specified
kwargs['file'] = sys.stderr
print_colored(msg, 'YELLOW', **kwargs)
if log:
LOG.warning(msg)
def set_title(title):
"""Set window title."""
LOG.debug('title: %s', title)
if os.name == 'nt':
os.system(f'title {title}')
else:
print_error('Setting the title is only supported under Windows.')
def show_data(message, data, color=None, indent=None, width=None):
"""Display info using default or provided indent and width."""
colors = (None, color if color else None)
indent = INDENT if indent is None else indent
width = WIDTH if width is None else width
print_colored(
(f'{" "*indent}{message:<{width}}', data),
colors,
log=True,
sep='',
)
if __name__ == '__main__':
print("This file is not meant to be called directly.")

172
scripts/wk/tmux.py → scripts/wk/ui/tmux.py
View file

@ -26,38 +26,88 @@ def capture_pane(pane_id=None):
def clear_pane(pane_id=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, forced=False):
"""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:
continue
# Remove closed panes
for data in layout.values():
data['Panes'] = [pane for pane in data['Panes'] if poll_pane(pane)]
# 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
# 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
# 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):
@ -77,6 +127,20 @@ def get_pane_size(pane_id=None):
return (width, height)
def get_window_size():
"""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()
width = int(width)
height = int(height)
# Done
return (width, height)
def kill_all_panes(pane_id=None):
"""Kill all panes except for the current or target pane."""
cmd = ['tmux', 'kill-pane', '-a']
@ -96,34 +160,30 @@ def kill_pane(*pane_ids):
run_program(cmd+[pane_id], check=False)
def layout_needs_fixed(panes, layout):
def layout_needs_fixed(layout):
"""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
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']
)
# 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
# TODO: Re-enable?
## Group panes
#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 in layout[group]['Panes']:
# needs_fixed = needs_fixed or abs(get_pane_size(pane)[0] - width) > 2
# Done
return needs_fixed
@ -182,7 +242,7 @@ def prep_action(
'cat',
])
elif watch_cmd == 'tail':
action_cmd.extend(['tail', '-f'])
action_cmd.extend(['tail', '-q', '-f'])
action_cmd.append(watch_file)
else:
LOG.error('No action specified')
@ -202,7 +262,7 @@ def prep_file(path):
pass
def resize_pane(pane_id=None, width=None, height=None, **kwargs):
def resize_pane(pane_id=None, width=None, height=None):
"""Resize current or target pane.
NOTE: kwargs is only here to make calling this function easier
@ -227,6 +287,15 @@ def resize_pane(pane_id=None, width=None, height=None, **kwargs):
run_program(cmd, check=False)
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 split_window(
lines=None, percent=None,
behind=False, vertical=False,
@ -263,15 +332,6 @@ 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):
"""Toggle zoom status for current or target pane."""
cmd = ['tmux', 'resize-pane', '-Z']

355
scripts/wk/ui/tui.py Normal file
View file

@ -0,0 +1,355 @@
"""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 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=None) -> None:
self.layout = deepcopy(TMUX_LAYOUT)
self.side_width = TMUX_SIDE_WIDTH
self.title_text = title_text if title_text else 'Title Text'
self.title_text_line2 = ''
self.title_colors = ['BLUE', None]
# 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=None, percent=None, update_layout=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, line2=None, colors=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=None, percent=None, update_layout=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,
'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=True) -> None:
"""Fix tmux layout based on self.layout."""
try:
tmux.fix_layout(self.layout, 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 set_current_pane_height(self, height) -> 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) -> 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, line2=None, colors=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=False):
"""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):
"""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
def test():
"""TODO: Deleteme"""
ui = TUI()
ui.add_info_pane(lines=10, text='Info One')
ui.add_info_pane(lines=10, text='Info Two')
ui.add_info_pane(lines=10, text='Info Three')
ui.add_worker_pane(lines=3, text='Work One')
ui.add_worker_pane(lines=3, text='Work Two')
ui.add_worker_pane(lines=3, text='Work Three')
ui.fix_layout()
return ui
if __name__ == '__main__':
print("This file is not meant to be called directly.")

8
scripts/wk_debug.py
View file

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

4
setup/windows/build.ps1
View file

@ -107,9 +107,9 @@ if ($MyInvocation.InvocationName -ne ".") {
$Url = FindDynamicUrl $DownloadPage $RegEx
DownloadFile -Path $Temp -Name "psutil64.whl" -Url $Url
# Python: pytz, requests, & dependancies
# Python: prompt_toolkit, pytz, requests, & dependancies
$RegEx = "href=.*.py3-none-any.whl"
foreach ($Module in @("chardet", "certifi", "idna", "pytz", "urllib3", "requests")) {
foreach ($Module in @("chardet", "certifi", "idna", "prompt_toolkit", "Pygments", "pytz", "requests", "urllib3", "wcwidth")) {
$DownloadPage = "https://pypi.org/project/$Module/"
$Name = "$Module.whl"
$Url = FindDynamicUrl -SourcePage $DownloadPage -RegEx $RegEx