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

This commit is contained in:
2Shirt 2020-01-08 17:07:33 -07:00
commit 8932242a86
Signed by: 2Shirt
GPG key ID: 152FAC923B0E132C
73 changed files with 8882 additions and 1433 deletions

48
.gitignore vendored
View file

@ -1,41 +1,7 @@
**/__pycache__/*
*.bak
*.exe
*.swp
.bin/7-Zip/
.bin/AIDA64/
.bin/BleachBit/
.bin/ClassicStartSkin/
.bin/ConEmu/
.bin/Erunt/
.bin/Everything/
.bin/FastCopy/
.bin/HWiNFO/HWiNFO*.ini
.bin/NotepadPlusPlus/
.bin/ProcessKiller/
.bin/ProduKey/
.bin/Python/
.bin/Tmp/
.bin/XMPlay/
.bin/_Drivers/SDIO/
.cbin/*.7z
.cbin/AIDA64/
.cbin/Autoruns/
.cbin/BleachBit-Portable/
.cbin/BlueScreenView/
.cbin/Caffeine/
.cbin/Du/
.cbin/Everything/
.cbin/FirefoxExtensions/
.cbin/IObitUninstallerPortable/
.cbin/ProduKey/
.cbin/TestDisk/
.cbin/TreeSizeFree-Portable/
.cbin/XMPlay/
.cbin/XYplorerFree/
.cbin/_Drivers/
.cbin/_Office/
.cbin/_vcredists/
.cbin/wimlib/
BUILD*/
OUT*/
**/__pycache__
**/*.7z
**/*.bak
**/*.exe
**/*.swp
setup/BUILD*
setup/OUT*

6
scripts/README.md Normal file
View file

@ -0,0 +1,6 @@
## pylint ##
These scripts use two spaces per indent instead of the default four. As such you will need to update your pylintrc file or run like this:
`pylint --indent-after-paren=2 --indent-string=' ' wk`

31
scripts/activate.py Normal file
View file

@ -0,0 +1,31 @@
"""Wizard Kit: Activate Windows using a BIOS key"""
# vim: sts=2 sw=2 ts=2
import wk
def main():
"""Attempt to activate Windows and show result."""
title = f'{wk.cfg.main.KIT_NAME_FULL}: Activation Tool'
try_print = wk.std.TryAndPrint()
wk.std.clear_screen()
wk.std.set_title(title)
wk.std.print_info(title)
print('')
# Attempt activation
try_print.run('Attempting activation...', wk.os.win.activate_with_bios)
# Done
print('')
print('Done.')
wk.std.pause('Press Enter to exit...')
if __name__ == '__main__':
try:
main()
except SystemExit:
raise
except: #pylint: disable=bare-except
wk.std.major_exception()

14
scripts/build-ufd Executable file
View file

@ -0,0 +1,14 @@
#!/usr/bin/env python3
"""Wizard Kit: Build UFD Tool"""
# vim: sts=2 sw=2 ts=2
import wk
if __name__ == '__main__':
try:
wk.kit.ufd.build_ufd()
except SystemExit:
raise
except: #pylint: disable=bare-except
wk.std.major_exception()

49
scripts/check_disk.py Normal file
View file

@ -0,0 +1,49 @@
"""Wizard Kit: Check or repair the %SYSTEMDRIVE% filesystem via CHKDSK"""
# vim: sts=2 sw=2 ts=2
import os
import wk
def main():
"""Run or schedule CHKDSK and show result."""
title = f'{wk.cfg.main.KIT_NAME_FULL}: Check Disk Tool'
menu = wk.std.Menu(title=title)
try_print = wk.std.TryAndPrint()
wk.std.clear_screen()
wk.std.set_title(title)
print('')
# Add menu entries
menu.add_option('Offline scan')
menu.add_option('Online scan')
# Show menu and make selection
selection = menu.simple_select()
# Run or schedule scan
if 'Offline' in selection[0]:
function = wk.os.win.run_chkdsk_offline
msg_good = 'Scheduled'
else:
function = wk.os.win.run_chkdsk_online
msg_good = 'No issues detected'
try_print.run(
message=f'CHKDSK ({os.environ.get("SYSTEMDRIVE")})...',
function=function,
msg_good=msg_good,
)
# Done
print('')
print('Done.')
wk.std.pause('Press Enter to exit...')
if __name__ == '__main__':
try:
main()
except SystemExit:
raise
except: #pylint: disable=bare-except
wk.std.major_exception()

21
scripts/ddrescue-tui Executable file
View file

@ -0,0 +1,21 @@
#!/bin/bash
#
## Wizard Kit: ddrescue TUI Launcher
# Check if running under Linux
os_name="$(uname -s)"
if [[ "$os_name" == "Darwin" ]]; then
os_name="macOS"
fi
if [[ "$os_name" != "Linux" ]]; then
echo "This script is not supported under $os_name." 1>&2
exit 1
fi
source ./launch-in-tmux
SESSION_NAME="ddrescue-tui"
WINDOW_NAME="ddrescue TUI"
TMUX_CMD="./ddrescue-tui.py"
launch_in_tmux "$@"

14
scripts/ddrescue-tui.py Executable file
View file

@ -0,0 +1,14 @@
#!/usr/bin/env python3
"""Wizard Kit: ddrescue TUI"""
# vim: sts=2 sw=2 ts=2
import wk
if __name__ == '__main__':
try:
wk.hw.ddrescue.main()
except SystemExit:
raise
except: #pylint: disable=bare-except
wk.std.major_exception()

View file

@ -2,10 +2,10 @@
#
## Wizard Kit: HW Diagnostics Launcher
source launch-in-tmux
source ./launch-in-tmux
SESSION_NAME="hw-diags"
WINDOW_NAME="Hardware Diagnostics"
TMUX_CMD="hw-diags-menu"
TMUX_CMD="./hw-diags.py"
launch_in_tmux "$@"

14
scripts/hw-diags.py Executable file
View file

@ -0,0 +1,14 @@
#!/usr/bin/env python3
"""Wizard Kit: Hardware Diagnostics"""
# vim: sts=2 sw=2 ts=2
import wk
if __name__ == '__main__':
try:
wk.hw.diags.main()
except SystemExit:
raise
except: #pylint: disable=bare-except
wk.std.major_exception()

49
scripts/hw-drive-info Executable file
View file

@ -0,0 +1,49 @@
#!/bin/bash
#
BLUE='\033[34m'
CLEAR='\033[0m'
IFS=$'\n'
# Check if running under Linux
os_name="$(uname -s)"
if [[ "$os_name" == "Darwin" ]]; then
os_name="macOS"
fi
if [[ "$os_name" != "Linux" ]]; then
echo "This script is not supported under $os_name." 1>&2
exit 1
fi
# List devices
for line in $(lsblk -do NAME,TRAN,SIZE,VENDOR,MODEL,SERIAL); do
if [[ "${line:0:4}" == "NAME" ]]; then
echo -e "${BLUE}${line}${CLEAR}"
else
echo "${line}"
fi
done
echo ""
# List loopback devices
if [[ "$(losetup -l | wc -l)" > 0 ]]; then
for line in $(losetup -lO NAME,PARTSCAN,RO,BACK-FILE); do
if [[ "${line:0:4}" == "NAME" ]]; then
echo -e "${BLUE}${line}${CLEAR}"
else
echo "${line}" | sed -r 's#/dev/(loop[0-9]+)#\1 #'
fi
done
echo ""
fi
# List partitions
for line in $(lsblk -o NAME,SIZE,FSTYPE,LABEL,MOUNTPOINT); do
if [[ "${line:0:4}" == "NAME" ]]; then
echo -e "${BLUE}${line}${CLEAR}"
else
echo "${line}"
fi
done
echo ""

View file

@ -25,6 +25,16 @@ function print_dmi_value() {
print_in_columns "$name: $value"
}
# Check if running under Linux
os_name="$(uname -s)"
if [[ "$os_name" == "Darwin" ]]; then
os_name="macOS"
fi
if [[ "$os_name" != "Linux" ]]; then
echo "This script is not supported under $os_name." 1>&2
exit 1
fi
# System
echo -e "${BLUE}System Information${CLEAR}"
print_dmi_value "Vendor" "sys_vendor"
@ -86,7 +96,7 @@ echo ""
# Audio
echo -e "${BLUE}Audio${CLEAR}"
while read -r line; do
if [[ "$line" =~ .*no.soundcards.found.* ]]; then
if [[ "$line" = .*no.soundcards.found.* ]]; then
echo " No soundcards found"
else
print_in_columns "$line"

46
scripts/hw-sensors Executable file
View file

@ -0,0 +1,46 @@
#!/usr/bin/env python3
"""Wizard Kit: Hardware Sensors"""
# vim: sts=2 sw=2 ts=2
import platform
import wk
def main():
"""Show sensor data on screen."""
sensors = wk.hw.sensors.Sensors()
if platform.system() == 'Darwin':
wk.std.clear_screen()
while True:
print('\033[100A', end='')
sensors.update_sensor_data()
wk.std.print_report(sensors.generate_report('Current', 'Max'))
wk.std.sleep(1)
elif platform.system() == 'Linux':
proc = wk.exe.run_program(cmd=['mktemp'])
sensors.start_background_monitor(
out_path=proc.stdout.strip(),
exit_on_thermal_limit=False,
temp_labels=('Current', 'Max'),
)
watch_cmd = [
'watch',
'--color',
'--exec',
'--no-title',
'--interval', '1',
'cat',
proc.stdout.strip(),
]
wk.exe.run_program(watch_cmd, check=False, pipe=False)
if __name__ == '__main__':
try:
main()
except KeyboardInterrupt:
pass
except SystemExit:
raise
except: #pylint: disable=bare-except
wk.std.major_exception()

View file

@ -13,16 +13,16 @@ function ask() {
done
}
die () {
function err () {
echo "$0:" "$@" >&2
exit 1
return 1
}
function launch_in_tmux() {
# Check for required vars
[[ -n "${SESSION_NAME:-}" ]] || die "Required variable missing (SESSION_NAME)"
[[ -n "${WINDOW_NAME:-}" ]] || die "Required variable missing (WINDOW_NAME)"
[[ -n "${TMUX_CMD:-}" ]] || die "Required variable missing (TMUX_CMD)"
[[ -n "${SESSION_NAME:-}" ]] || return $(err "Required variable missing (SESSION_NAME)")
[[ -n "${WINDOW_NAME:-}" ]] || return $(err "Required variable missing (WINDOW_NAME)")
[[ -n "${TMUX_CMD:-}" ]] || return $(err "Required variable missing (TMUX_CMD)")
# Check for running session
if tmux list-session | grep -q "$SESSION_NAME"; then
@ -33,33 +33,34 @@ function launch_in_tmux() {
# Running inside TMUX, switch to session
tmux switch-client -t "$SESSION_NAME"
if ! jobs %% >/dev/null 2>&1; then
# No running jobs, try exiting abandoned tmux session
exit 0
fi
else
# Running outside TMUX, attach to session
tmux attach-session -t "$SESSION_NAME"
fi
exit 0
return 0
elif ask "Kill current session and start new session?"; then
tmux kill-session -t "$SESSION_NAME" || \
die "Failed to kill session: $SESSION_NAME"
else
echo "Aborted."
echo ""
echo -n "Press Enter to exit... "
read -r
exit 0
return 1
fi
fi
# Start/Rename session
# Start session
if [[ -n "${TMUX:-}" ]]; then
# Running inside TMUX, rename session/window and open the menu
# Running inside TMUX, save current session/window names
ORIGINAL_SESSION_NAME="$(tmux display-message -p '#S')"
ORIGINAL_WINDOW_NAME="$(tmux display-message -p '#W')"
tmux rename-session "$SESSION_NAME"
tmux rename-window "$WINDOW_NAME"
"$TMUX_CMD" "$@"
tmux rename-session "${SESSION_NAME}_DONE"
tmux rename-window "${WINDOW_NAME}_DONE"
# Restore previous session/window names
tmux rename-session "${ORIGINAL_SESSION_NAME}"
tmux rename-window "${ORIGINAL_WINDOW_NAME}"
else
# Running outside TMUX, start/attach to session
tmux new-session -s "$SESSION_NAME" -n "$WINDOW_NAME" "$TMUX_CMD" "$@"

33
scripts/mount-all-volumes Executable file
View file

@ -0,0 +1,33 @@
#!/usr/bin/env python3
"""Wizard Kit: Mount all volumes"""
# vim: sts=2 sw=2 ts=2
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(' ')
# Mount volumes and get report
wk.std.print_standard('Mounting volumes...')
report = wk.os.linux.mount_volumes()
# Show results
wk.std.print_info('Results')
wk.std.print_report(report, indent=2)
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()
try:
main()
except SystemExit:
raise
except: #pylint: disable=bare-except
wk.std.major_exception()

30
scripts/mount-backup-shares Executable file
View file

@ -0,0 +1,30 @@
#!/usr/bin/env python3
"""Wizard Kit: Mount Backup Shares"""
# pylint: disable=invalid-name
# vim: sts=2 sw=2 ts=2
import wk
# Functions
def main():
"""Attempt to mount backup shares and print report."""
wk.std.print_info('Mounting Backup Shares')
report = wk.net.mount_backup_shares()
for line in report:
color = 'GREEN'
line = f' {line}'
if 'Failed' in line:
color = 'RED'
elif 'Already' in line:
color = 'YELLOW'
print(wk.std.color_string(line, color))
if __name__ == '__main__':
try:
main()
except SystemExit:
raise
except: #pylint: disable=bare-except
wk.std.major_exception()

View file

@ -1,63 +0,0 @@
# Wizard Kit: Activate Windows using various methods
import os
import sys
# Init
sys.path.append(os.path.dirname(os.path.realpath(__file__)))
from functions.activation import *
init_global_vars()
os.system('title {}: Windows Activation Tool'.format(KIT_NAME_FULL))
if __name__ == '__main__':
try:
stay_awake()
clear_screen()
print_info('{}: Windows Activation Tool\n'.format(KIT_NAME_FULL))
# Bail early if already activated
if windows_is_activated():
print_info('This system is already activated')
sleep(5)
exit_script()
other_results = {
'Error': {
'BIOSKeyNotFoundError': 'BIOS key not found.',
}}
# Determine activation method
activation_methods = [
{'Name': 'Activate with BIOS key', 'Function': activate_with_bios},
]
if global_vars['OS']['Version'] not in ('8', '8.1', '10'):
activation_methods[0]['Disabled'] = True
actions = [
{'Name': 'Quit', 'Letter': 'Q'},
]
while True:
selection = menu_select(
'{}: Windows Activation Menu'.format(KIT_NAME_FULL),
main_entries=activation_methods, action_entries=actions)
if (selection.isnumeric()):
result = try_and_print(
message = activation_methods[int(selection)-1]['Name'],
function = activation_methods[int(selection)-1]['Function'],
other_results=other_results)
if result['CS']:
break
else:
sleep(2)
elif selection == 'Q':
exit_script()
# Done
print_success('\nDone.')
pause("Press Enter to exit...")
exit_script()
except SystemExit as sys_exit:
exit_script(sys_exit.code)
except:
major_exception()
# vim: sts=2 sw=2 ts=2

View file

@ -1,149 +0,0 @@
#!/bin/env python3
#
# pylint: disable=no-name-in-module,wildcard-import,wrong-import-position
# vim: sts=2 sw=2 ts=2
"""Wizard Kit: UFD build tool"""
import os
import sys
# Init
sys.path.append(os.path.dirname(os.path.realpath(__file__)))
from docopt import docopt
from functions.common import *
from functions.ufd import *
from settings.ufd import *
init_global_vars(silent=True)
# Main section
if __name__ == '__main__':
# pylint: disable=invalid-name
# Set log
try:
global_vars['LogDir'] = '{}/Logs'.format(
get_user_home(get_user_name()))
set_log_file('Build UFD ({Date-Time}).log'.format(**global_vars))
except: # pylint: disable=bare-except
major_exception()
# Header
print_success(KIT_NAME_FULL)
print_standard('UFD Build Tool')
print_standard(' ')
# Check if running as root
if not running_as_root():
print_error('ERROR: This script is meant to be run as root.')
abort(False)
# Docopt
try:
args = docopt(DOCSTRING)
except SystemExit as sys_exit:
# Catch docopt exits
exit_script(sys_exit.code)
except: # pylint: disable=bare-except
major_exception()
try:
# Verify selections
ufd_dev = verify_ufd(args['--ufd-device'])
sources = verify_sources(args, UFD_SOURCES)
show_selections(args, sources, ufd_dev, UFD_SOURCES)
if not args['--force']:
confirm_selections(args)
# Prep UFD
if not args['--update']:
print_info('Prep UFD')
prep_device(ufd_dev, UFD_LABEL, use_mbr=args['--use-mbr'])
# Mount UFD
try_and_print(
indent=2,
message='Mounting UFD...',
function=mount,
mount_source=find_first_partition(ufd_dev),
mount_point='/mnt/UFD',
read_write=True,
)
# Remove Arch folder
if args['--update']:
try_and_print(
indent=2,
message='Removing Linux...',
function=remove_arch,
)
# Copy sources
print_standard(' ')
print_info('Copy Sources')
for s_label, s_path in sources.items():
try_and_print(
indent=2,
message='Copying {}...'.format(s_label),
function=copy_source,
source=s_path,
items=ITEMS[s_label],
overwrite=True,
)
# Update boot entries
print_standard(' ')
print_info('Boot Setup')
try_and_print(
indent=2,
message='Updating boot entries...',
function=update_boot_entries,
boot_entries=BOOT_ENTRIES,
boot_files=BOOT_FILES,
iso_label=ISO_LABEL,
ufd_label=UFD_LABEL,
)
# Install syslinux (to partition)
try_and_print(
indent=2,
message='Syslinux (partition)...',
function=install_syslinux_to_partition,
partition=find_first_partition(ufd_dev),
)
# Unmount UFD
try_and_print(
indent=2,
message='Unmounting UFD...',
function=unmount,
mount_point='/mnt/UFD',
)
# Install syslinux (to device)
try_and_print(
indent=2,
message='Syslinux (device)...',
function=install_syslinux_to_dev,
ufd_dev=ufd_dev,
use_mbr=args['--use-mbr'],
)
# Hide items
print_standard(' ')
print_info('Final Touches')
try_and_print(
indent=2,
message='Hiding items...',
function=hide_items,
ufd_dev=ufd_dev,
items=ITEMS_HIDDEN,
)
# Done
if not args['--force']:
print_standard('\nDone.')
pause('Press Enter to exit...')
exit_script()
except SystemExit as sys_exit:
exit_script(sys_exit.code)
except: # pylint: disable=bare-except
major_exception()

View file

@ -1,43 +0,0 @@
# Wizard Kit: Backup CBS Logs and prep CBS temp data for deletion
import os
import sys
# Init
sys.path.append(os.path.dirname(os.path.realpath(__file__)))
from functions.cleanup import *
from functions.data import *
init_global_vars()
os.system('title {}: CBS Cleanup'.format(KIT_NAME_FULL))
set_log_file('CBS Cleanup.log')
if __name__ == '__main__':
try:
# Prep
stay_awake()
clear_screen()
folder_path = r'{}\Backups'.format(KIT_NAME_SHORT)
dest = select_destination(folder_path=folder_path,
prompt='Which disk are we using for temp data and backup?')
# Show details
print_info('{}: CBS Cleanup Tool\n'.format(KIT_NAME_FULL))
show_data('Backup / Temp path:', dest)
print_standard('\n')
if (not ask('Proceed with CBS cleanup?')):
abort()
# Run Cleanup
try_and_print(message='Running cleanup...', function=cleanup_cbs,
cs='Done', dest_folder=dest)
# Done
print_standard('\nDone.')
pause("Press Enter to exit...")
exit_script()
except SystemExit as sys_exit:
exit_script(sys_exit.code)
except:
major_exception()
# vim: sts=2 sw=2 ts=2

View file

@ -1,11 +0,0 @@
#!/bin/bash
#
## Wizard Kit: ddrescue TUI Launcher
source launch-in-tmux
SESSION_NAME="ddrescue-tui"
WINDOW_NAME="ddrescue TUI"
TMUX_CMD="ddrescue-tui-menu"
launch_in_tmux "$@"

View file

@ -1,64 +0,0 @@
#!/bin/python3
#
## Wizard Kit: TUI for ddrescue cloning and imaging
import os
import sys
# Init
sys.path.append(os.path.dirname(os.path.realpath(__file__)))
from functions.ddrescue import *
from functions.hw_diags import *
init_global_vars()
if __name__ == '__main__':
try:
# Prep
clear_screen()
args = list(sys.argv)
run_mode = ''
source_path = None
dest_path = None
# Parse args
try:
script_name = os.path.basename(args.pop(0))
run_mode = str(args.pop(0)).lower()
source_path = args.pop(0)
dest_path = args.pop(0)
except IndexError:
# We'll set the missing paths later
pass
# Show usage
if re.search(r'-+(h|help)', str(sys.argv), re.IGNORECASE):
show_usage(script_name)
exit_script()
# Start cloning/imaging
if run_mode in ('clone', 'image'):
menu_ddrescue(source_path, dest_path, run_mode)
else:
if not re.search(r'^-*(h|help\?)', run_mode, re.IGNORECASE):
print_error('Invalid mode.')
# Done
print_standard('\nDone.')
pause("Press Enter to exit...")
tmux_switch_client()
exit_script()
except GenericAbort:
abort()
except GenericError as ge:
msg = 'Generic Error'
if str(ge):
msg = str(ge)
print_error(msg)
abort()
except SystemExit as sys_exit:
tmux_switch_client()
exit_script(sys_exit.code)
except:
major_exception()
# vim: sts=2 sw=2 ts=2

View file

@ -1,42 +0,0 @@
#!/bin/python3
#
## Wizard Kit: HW Diagnostics - Audio
import os
import sys
# Init
sys.path.append(os.path.dirname(os.path.realpath(__file__)))
from functions.common import *
init_global_vars()
if __name__ == '__main__':
try:
# Prep
clear_screen()
print_standard('Hardware Diagnostics: Audio\n')
# Set volume
try:
run_program('amixer -q set "Master" 80% unmute'.split())
run_program('amixer -q set "PCM" 90% unmute'.split())
except subprocess.CalledProcessError:
print_error('Failed to set volume')
# Run tests
for mode in ['pink', 'wav']:
run_program(
cmd = 'speaker-test -c 2 -l 1 -t {}'.format(mode).split(),
check = False,
pipe = False)
# Done
#print_standard('\nDone.')
#pause("Press Enter to exit...")
exit_script()
except SystemExit as sys_exit:
exit_script(sys_exit.code)
except:
major_exception()
# vim: sts=2 sw=2 ts=2

View file

@ -1,18 +0,0 @@
#!/bin/bash
#
## Wizard Kit: HW Diagnostics - Benchmarks
function usage {
echo "Usage: ${0} device log-file"
echo " e.g. ${0} /dev/sda /tmp/tmp.XXXXXXX/benchmarks.log"
}
# Bail early
if [ ! -b "${1}" ]; then
usage
exit 1
fi
# Run Benchmarks
echo 3 | sudo tee -a /proc/sys/vm/drop_caches >/dev/null 2>&1
sudo dd bs=4M if="${1}" of=/dev/null status=progress 2>&1 | tee -a "${2}"

View file

@ -1,66 +0,0 @@
#!/bin/python3
#
## Wizard Kit: HW Diagnostics - Menu
import os
import sys
# Init
sys.path.append(os.path.dirname(os.path.realpath(__file__)))
from functions.hw_diags import *
from functions.tmux import *
init_global_vars()
if __name__ == '__main__':
# Show menu
try:
state = State()
menu_diags(state, sys.argv)
except KeyboardInterrupt:
print_standard(' ')
print_warning('Aborted')
print_standard(' ')
sleep(1)
pause('Press Enter to exit...')
except SystemExit as sys_exit:
tmux_switch_client()
exit_script(sys_exit.code)
except:
# Cleanup
tmux_kill_all_panes()
if DEBUG_MODE:
# Custom major exception
print_standard(' ')
print_error('Major exception')
print_warning(SUPPORT_MESSAGE)
print(traceback.format_exc())
print_log(traceback.format_exc())
# Save debug reports and upload data
try_and_print(
message='Saving debug reports...',
function=save_debug_reports,
state=state, global_vars=global_vars)
question = 'Upload crash details to {}?'.format(CRASH_SERVER['Name'])
if ENABLED_UPLOAD_DATA and ask(question):
try_and_print(
message='Uploading Data...',
function=upload_logdir,
global_vars=global_vars)
# Done
sleep(1)
pause('Press Enter to exit...')
exit_script(1)
else:
# "Normal" major exception
major_exception()
# Done
tmux_kill_all_panes()
tmux_switch_client()
exit_script()
# vim: sts=2 sw=2 ts=2

View file

@ -1,48 +0,0 @@
#!/bin/python3
#
## Wizard Kit: HW Diagnostics - Network
import os
import sys
# Init
sys.path.append(os.path.dirname(os.path.realpath(__file__)))
from functions.network import *
def check_connection():
if not is_connected():
# Raise to cause NS in try_and_print()
raise Exception
if __name__ == '__main__':
try:
# Prep
clear_screen()
print_standard('Hardware Diagnostics: Network\n')
# Connect
print_standard('Initializing...')
connect_to_network()
# Tests
try_and_print(
message='Network connection:', function=check_connection, cs='OK')
show_valid_addresses()
try_and_print(message='Internet connection:', function=ping,
addr='8.8.8.8', cs='OK')
try_and_print(message='DNS Resolution:', function=ping, cs='OK')
try_and_print(message='Speedtest:', function=speedtest,
print_return=True)
# Done
print_standard('\nDone.')
#pause("Press Enter to exit...")
exit_script()
except SystemExit as sys_exit:
exit_script(sys_exit.code)
except:
major_exception()
# vim: sts=2 sw=2 ts=2

View file

@ -1,19 +0,0 @@
#!/bin/bash
#
## Wizard Kit: HW Diagnostics - Prime95
function usage {
echo "Usage: $0 log-dir"
echo " e.g. $0 /tmp/tmp.7Mh5f1RhSL9001"
}
# Bail early
if [ ! -d "$1" ]; then
usage
exit 1
fi
# Run Prime95
cd "$1"
mprime -t | grep -iv --line-buffered 'stress.txt' | tee -a "prime.log"

View file

@ -1,39 +0,0 @@
#!/bin/bash
#
BLUE='\033[34m'
CLEAR='\033[0m'
IFS=$'\n'
# List devices
for line in $(lsblk -do NAME,TRAN,SIZE,VENDOR,MODEL,SERIAL); do
if [[ "${line:0:4}" == "NAME" ]]; then
echo -e "${BLUE}${line}${CLEAR}"
else
echo "${line}"
fi
done
echo ""
# List loopback devices
if [[ "$(losetup -l | wc -l)" > 0 ]]; then
for line in $(losetup -lO NAME,PARTSCAN,RO,BACK-FILE); do
if [[ "${line:0:4}" == "NAME" ]]; then
echo -e "${BLUE}${line}${CLEAR}"
else
echo "${line}" | sed -r 's#/dev/(loop[0-9]+)#\1 #'
fi
done
echo ""
fi
# List partitions
for line in $(lsblk -o NAME,SIZE,FSTYPE,LABEL,MOUNTPOINT); do
if [[ "${line:0:4}" == "NAME" ]]; then
echo -e "${BLUE}${line}${CLEAR}"
else
echo "${line}"
fi
done
echo ""

View file

@ -1,10 +0,0 @@
#!/bin/bash
#
## Wizard Kit: Sensor monitoring tool
WINDOW_NAME="Hardware Sensors"
MONITOR="hw-sensors-monitor"
# Start session
tmux new-session -n "$WINDOW_NAME" "$MONITOR"

View file

@ -1,36 +0,0 @@
#!/bin/python3
#
## Wizard Kit: Sensor monitoring tool
import os
import sys
# Init
sys.path.append(os.path.dirname(os.path.realpath(__file__)))
from functions.sensors import *
from functions.tmux import *
init_global_vars(silent=True)
if __name__ == '__main__':
background = False
try:
if len(sys.argv) > 1 and os.path.exists(sys.argv[1]):
background = True
monitor_file = sys.argv[1]
monitor_pane = None
else:
result = run_program(['mktemp'])
monitor_file = result.stdout.decode().strip()
if not background:
monitor_pane = tmux_split_window(
percent=1, vertical=True, watch=monitor_file)
cmd = ['tmux', 'resize-pane', '-Z', '-t', monitor_pane]
run_program(cmd, check=False)
monitor_sensors(monitor_pane, monitor_file)
exit_script()
except SystemExit as sys_exit:
exit_script(sys_exit.code)
except:
major_exception()
# vim: sts=2 sw=2 ts=2

View file

@ -1,38 +0,0 @@
#!/bin/python3
#
## Wizard Kit: Volume mount tool
import os
import sys
# Init
sys.path.append(os.path.dirname(os.path.realpath(__file__)))
from functions.data import *
init_global_vars()
if __name__ == '__main__':
try:
# Prep
clear_screen()
print_standard('{}: Volume mount tool'.format(KIT_NAME_FULL))
# Mount volumes
report = mount_volumes(all_devices=True)
# Print report
print_info('\nResults')
for vol_name, vol_data in sorted(report.items()):
show_data(indent=4, width=20, **vol_data['show_data'])
# Done
print_standard('\nDone.')
if 'gui' in sys.argv:
pause("Press Enter to exit...")
popen_program(['nohup', 'thunar', '/media'], pipe=True)
exit_script()
except SystemExit as sys_exit:
exit_script(sys_exit.code)
except:
major_exception()
# vim: sts=2 sw=2 ts=2

View file

@ -1,35 +0,0 @@
#!/bin/python3
#
## Wizard Kit: Backup share mount tool
import os
import sys
# Init
sys.path.append(os.path.dirname(os.path.realpath(__file__)))
from functions.data import *
from functions.network import *
init_global_vars()
if __name__ == '__main__':
try:
# Prep
clear_screen()
# Mount
if is_connected():
mount_backup_shares(read_write=True)
else:
# Couldn't connect
print_error('ERROR: No network connectivity.')
# Done
print_standard('\nDone.')
#pause("Press Enter to exit...")
exit_script()
except SystemExit as sys_exit:
exit_script(sys_exit.code)
except:
major_exception()
# vim: sts=2 sw=2 ts=2

View file

@ -1,39 +0,0 @@
# Wizard Kit: Enter SafeMode by editing the BCD
import os
import sys
# Init
sys.path.append(os.path.dirname(os.path.realpath(__file__)))
from functions.safemode import *
init_global_vars()
os.system('title {}: SafeMode Tool'.format(KIT_NAME_FULL))
if __name__ == '__main__':
try:
clear_screen()
print_info('{}: SafeMode Tool\n'.format(KIT_NAME_FULL))
other_results = {
'Error': {'CalledProcessError': 'Unknown Error'},
'Warning': {}}
if not ask('Enable booting to SafeMode (with Networking)?'):
abort()
# Configure SafeMode
try_and_print(message='Set BCD option...',
function=enable_safemode, other_results=other_results)
try_and_print(message='Enable MSI in SafeMode...',
function=enable_safemode_msi, other_results=other_results)
# Done
print_standard('\nDone.')
pause('Press Enter to reboot...')
reboot()
exit_script()
except SystemExit as sys_exit:
exit_script(sys_exit.code)
except:
major_exception()
# vim: sts=2 sw=2 ts=2

View file

@ -1,39 +0,0 @@
# Wizard Kit: Exit SafeMode by editing the BCD
import os
import sys
# Init
sys.path.append(os.path.dirname(os.path.realpath(__file__)))
from functions.safemode import *
init_global_vars()
os.system('title {}: SafeMode Tool'.format(KIT_NAME_FULL))
if __name__ == '__main__':
try:
clear_screen()
print_info('{}: SafeMode Tool\n'.format(KIT_NAME_FULL))
other_results = {
'Error': {'CalledProcessError': 'Unknown Error'},
'Warning': {}}
if not ask('Disable booting to SafeMode?'):
abort()
# Configure SafeMode
try_and_print(message='Remove BCD option...',
function=disable_safemode, other_results=other_results)
try_and_print(message='Disable MSI in SafeMode...',
function=disable_safemode_msi, other_results=other_results)
# Done
print_standard('\nDone.')
pause('Press Enter to reboot...')
reboot()
exit_script()
except SystemExit as sys_exit:
exit_script(sys_exit.code)
except:
major_exception()
# vim: sts=2 sw=2 ts=2

View file

@ -1,40 +0,0 @@
# Wizard Kit: Check, and possibly repair, system file health via SFC
import os
import sys
# Init
sys.path.append(os.path.dirname(os.path.realpath(__file__)))
from functions.repairs import *
init_global_vars()
os.system('title {}: SFC Tool'.format(KIT_NAME_FULL))
set_log_file('SFC Tool.log')
if __name__ == '__main__':
try:
stay_awake()
clear_screen()
print_info('{}: SFC Tool\n'.format(KIT_NAME_FULL))
other_results = {
'Error': {
'CalledProcessError': 'Unknown Error',
},
'Warning': {
'GenericRepair': 'Repaired',
}}
if ask('Run a SFC scan now?'):
try_and_print(message='SFC scan...',
function=run_sfc_scan, other_results=other_results)
else:
abort()
# Done
print_standard('\nDone.')
pause('Press Enter to exit...')
exit_script()
except SystemExit as sys_exit:
exit_script(sys_exit.code)
except:
major_exception()
# vim: sts=2 sw=2 ts=2

37
scripts/safemode_enter.py Normal file
View file

@ -0,0 +1,37 @@
"""Wizard Kit: Enter SafeMode by editing the BCD"""
# vim: sts=2 sw=2 ts=2
import wk
def main():
"""Prompt user to enter safe mode."""
title = f'{wk.cfg.main.KIT_NAME_FULL}: SafeMode Tool'
try_print = wk.std.TryAndPrint()
wk.std.clear_screen()
wk.std.set_title(title)
wk.std.print_info(title)
print('')
# Ask
if not wk.std.ask('Enable booting to SafeMode (with Networking)?'):
wk.std.abort()
print('')
# Configure SafeMode
try_print.run('Set BCD option...', wk.os.win.enable_safemode)
try_print.run('Enable MSI in SafeMode...', wk.os.win.enable_safemode_msi)
# Done
print('Done.')
wk.std.pause('Press Enter to reboot...')
wk.exe.run_program('shutdown -r -t 3'.split(), check=False)
if __name__ == '__main__':
try:
main()
except SystemExit:
raise
except: #pylint: disable=bare-except
wk.std.major_exception()

37
scripts/safemode_exit.py Normal file
View file

@ -0,0 +1,37 @@
"""Wizard Kit: Exit SafeMode by editing the BCD"""
# vim: sts=2 sw=2 ts=2
import wk
def main():
"""Prompt user to exit safe mode."""
title = f'{wk.cfg.main.KIT_NAME_FULL}: SafeMode Tool'
try_print = wk.std.TryAndPrint()
wk.std.clear_screen()
wk.std.set_title(title)
wk.std.print_info(title)
print('')
# Ask
if not wk.std.ask('Disable booting to SafeMode?'):
wk.std.abort()
print('')
# Configure SafeMode
try_print.run('Remove BCD option...', wk.os.win.disable_safemode)
try_print.run('Disable MSI in SafeMode...', wk.os.win.disable_safemode_msi)
# Done
print('Done.')
wk.std.pause('Press Enter to reboot...')
wk.exe.run_program('shutdown -r -t 3'.split(), check=False)
if __name__ == '__main__':
try:
main()
except SystemExit:
raise
except: #pylint: disable=bare-except
wk.std.major_exception()

35
scripts/sfc_scan.py Normal file
View file

@ -0,0 +1,35 @@
"""Wizard Kit: Check, and possibly repair, system file health via SFC"""
# vim: sts=2 sw=2 ts=2
import wk
def main():
"""Run SFC and report result."""
title = f'{wk.cfg.main.KIT_NAME_FULL}: SFC Tool'
try_print = wk.std.TryAndPrint()
wk.std.clear_screen()
wk.std.set_title(title)
wk.std.print_info(title)
print('')
# Ask
if not wk.std.ask('Run a SFC scan now?'):
wk.std.abort()
print('')
# Run
try_print.run('SFC scan...', wk.os.win.run_sfc_scan)
# Done
print('Done')
wk.std.pause('Press Enter to exit...')
if __name__ == '__main__':
try:
main()
except SystemExit:
raise
except: #pylint: disable=bare-except
wk.std.major_exception()

28
scripts/unmount-backup-shares Executable file
View file

@ -0,0 +1,28 @@
#!/usr/bin/env python3
"""Wizard Kit: Unmount Backup Shares"""
# pylint: disable=invalid-name
# vim: sts=2 sw=2 ts=2
import wk
# Functions
def main():
"""Attempt to mount backup shares and print report."""
wk.std.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))
if __name__ == '__main__':
try:
main()
except SystemExit:
raise
except: #pylint: disable=bare-except
wk.std.major_exception()

11
scripts/watch-mac Executable file
View file

@ -0,0 +1,11 @@
#!/bin/zsh
#
## watch-like utility
WATCH_FILE="${1}"
while :; do
echo -n "\e[100A"
cat "${WATCH_FILE}"
sleep 1s
done

View file

@ -1,471 +0,0 @@
"""Wizard Kit: Functions - UFD"""
# pylint: disable=broad-except,wildcard-import
# vim: sts=2 sw=2 ts=2
import os
import re
import shutil
import pathlib
from collections import OrderedDict
from functions.common import *
def case_insensitive_search(path, item):
"""Search path for item case insensitively, returns str."""
regex_match = '^{}$'.format(item)
real_path = ''
# Quick check first
if os.path.exists('{}/{}'.format(path, item)):
real_path = '{}{}{}'.format(
path,
'' if path == '/' else '/',
item,
)
# Check all items in dir
for entry in os.scandir(path):
if re.match(regex_match, entry.name, re.IGNORECASE):
real_path = '{}{}{}'.format(
path,
'' if path == '/' else '/',
entry.name,
)
# Done
if not real_path:
raise FileNotFoundError('{}/{}'.format(path, item))
return real_path
def confirm_selections(args):
"""Ask tech to confirm selections, twice if necessary."""
if not ask('Is the above information correct?'):
abort(False)
## Safety check
if not args['--update']:
print_standard(' ')
print_warning('SAFETY CHECK')
print_standard(
'All data will be DELETED from the disk and partition(s) listed above.')
print_standard(
'This is irreversible and will lead to {RED}DATA LOSS.{CLEAR}'.format(
**COLORS))
if not ask('Asking again to confirm, is this correct?'):
abort(False)
print_standard(' ')
def copy_source(source, items, overwrite=False):
"""Copy source items to /mnt/UFD."""
is_image = source.is_file()
# Mount source if necessary
if is_image:
mount(source, '/mnt/Source')
# Copy items
for i_source, i_dest in items:
i_source = '{}{}'.format(
'/mnt/Source' if is_image else source,
i_source,
)
i_dest = '/mnt/UFD{}'.format(i_dest)
try:
recursive_copy(i_source, i_dest, overwrite=overwrite)
except FileNotFoundError:
# Going to assume (hope) that this is fine
pass
# Unmount source if necessary
if is_image:
unmount('/mnt/Source')
def find_first_partition(dev_path):
"""Find path to first partition of dev, returns str."""
cmd = [
'lsblk',
'--list',
'--noheadings',
'--output', 'name',
'--paths',
dev_path,
]
result = run_program(cmd, encoding='utf-8', errors='ignore')
part_path = result.stdout.splitlines()[-1].strip()
return part_path
def find_path(path):
"""Find path case-insensitively, returns pathlib.Path obj."""
path_obj = pathlib.Path(path).resolve()
# Quick check first
if path_obj.exists():
return path_obj
# Fix case
parts = path_obj.relative_to('/').parts
real_path = '/'
for part in parts:
try:
real_path = case_insensitive_search(real_path, part)
except NotADirectoryError:
# Reclassify error
raise FileNotFoundError(path)
# Raise error if path doesn't exist
path_obj = pathlib.Path(real_path)
if not path_obj.exists():
raise FileNotFoundError(path_obj)
# Done
return path_obj
def get_user_home(user):
"""Get path to user's home dir, returns str."""
home_dir = None
cmd = ['getent', 'passwd', user]
result = run_program(cmd, encoding='utf-8', errors='ignore', check=False)
try:
home_dir = result.stdout.split(':')[5]
except Exception:
# Just use HOME from ENV (or '/root' if that fails)
home_dir = os.environ.get('HOME', '/root')
return home_dir
def get_user_name():
"""Get real user name, returns str."""
user = None
if 'SUDO_USER' in os.environ:
user = os.environ.get('SUDO_USER', 'Unknown')
else:
user = os.environ.get('USER', 'Unknown')
return user
def hide_items(ufd_dev, items):
"""Set FAT32 hidden flag for items."""
# pylint: disable=invalid-name
with open('/root/.mtoolsrc', 'w') as f:
f.write('drive U: file="{}"\n'.format(
find_first_partition(ufd_dev)))
f.write('mtools_skip_check=1\n')
# Hide items
for item in items:
cmd = ['yes | mattrib +h "U:/{}"'.format(item)]
run_program(cmd, check=False, shell=True)
def install_syslinux_to_dev(ufd_dev, use_mbr):
"""Install Syslinux to UFD (dev)."""
cmd = [
'dd',
'bs=440',
'count=1',
'if=/usr/lib/syslinux/bios/{}.bin'.format(
'mbr' if use_mbr else 'gptmbr',
),
'of={}'.format(ufd_dev),
]
run_program(cmd)
def install_syslinux_to_partition(partition):
"""Install Syslinux to UFD (partition)."""
cmd = [
'syslinux',
'--install',
'--directory',
'/arch/boot/syslinux/',
partition,
]
run_program(cmd)
def is_valid_path(path_obj, path_type):
"""Verify path_obj is valid by type, returns bool."""
valid_path = False
if path_type == 'DIR':
valid_path = path_obj.is_dir()
elif path_type == 'KIT':
valid_path = path_obj.is_dir() and path_obj.joinpath('.bin').exists()
elif path_type == 'IMG':
valid_path = path_obj.is_file() and path_obj.suffix.lower() == '.img'
elif path_type == 'ISO':
valid_path = path_obj.is_file() and path_obj.suffix.lower() == '.iso'
elif path_type == 'UFD':
valid_path = path_obj.is_block_device()
return valid_path
def mount(mount_source, mount_point, read_write=False):
"""Mount mount_source on mount_point."""
os.makedirs(mount_point, exist_ok=True)
cmd = [
'mount',
mount_source,
mount_point,
'-o',
'rw' if read_write else 'ro',
]
run_program(cmd)
def prep_device(dev_path, label, use_mbr=False, indent=2):
"""Format device in preparation for applying the WizardKit components
This is done is four steps:
1. Zero-out first 64MB (this deletes the partition table and/or bootloader)
2. Create a new partition table (GPT by default, optionally MBR)
3. Set boot flag
4. Format partition (FAT32, 4K aligned)
"""
# Zero-out first 64MB
cmd = 'dd bs=4M count=16 if=/dev/zero of={}'.format(dev_path).split()
try_and_print(
indent=indent,
message='Zeroing first 64MB...',
function=run_program,
cmd=cmd,
)
# Create partition table
cmd = 'parted {} --script -- mklabel {} mkpart primary fat32 4MiB {}'.format(
dev_path,
'msdos' if use_mbr else 'gpt',
'-1s' if use_mbr else '-4MiB',
).split()
try_and_print(
indent=indent,
message='Creating partition table...',
function=run_program,
cmd=cmd,
)
# Set boot flag
cmd = 'parted {} set 1 {} on'.format(
dev_path,
'boot' if use_mbr else 'legacy_boot',
).split()
try_and_print(
indent=indent,
message='Setting boot flag...',
function=run_program,
cmd=cmd,
)
# Format partition
cmd = [
'mkfs.vfat', '-F', '32',
'-n', label,
find_first_partition(dev_path),
]
try_and_print(
indent=indent,
message='Formatting partition...',
function=run_program,
cmd=cmd,
)
def recursive_copy(source, dest, overwrite=False):
"""Copy source to dest recursively.
NOTE: This uses rsync style source/dest syntax.
If the source has a trailing slash then it's contents are copied,
otherwise the source itself is copied.
Examples assuming "ExDir/ExFile.txt" exists:
recursive_copy("ExDir", "Dest/") results in "Dest/ExDir/ExFile.txt"
recursive_copy("ExDir/", "Dest/") results in "Dest/ExFile.txt"
NOTE 2: dest does not use find_path because it might not exist.
"""
copy_contents = source.endswith('/')
source = find_path(source)
dest = pathlib.Path(dest).resolve().joinpath(source.name)
os.makedirs(dest.parent, exist_ok=True)
if source.is_dir():
if copy_contents:
# Trailing slash syntax
for item in os.scandir(source):
recursive_copy(item.path, dest.parent, overwrite=overwrite)
elif not dest.exists():
# No conflict, copying whole tree (no merging needed)
shutil.copytree(source, dest)
elif not dest.is_dir():
# Refusing to replace file with dir
raise FileExistsError('Refusing to replace file: {}'.format(dest))
else:
# Dest exists and is a dir, merge dirs
for item in os.scandir(source):
recursive_copy(item.path, dest, overwrite=overwrite)
elif source.is_file():
if not dest.exists():
# No conflict, copying file
shutil.copy2(source, dest)
elif not dest.is_file():
# Refusing to replace dir with file
raise FileExistsError('Refusing to replace dir: {}'.format(dest))
elif overwrite:
# Dest file exists, deleting and replacing file
os.remove(dest)
shutil.copy2(source, dest)
else:
# Refusing to delete file when overwrite=False
raise FileExistsError('Refusing to delete file: {}'.format(dest))
def remove_arch():
"""Remove arch dir from UFD.
This ensures a clean installation to the UFD and resets the boot files
"""
shutil.rmtree(find_path('/mnt/UFD/arch'))
def running_as_root():
"""Check if running with effective UID of 0, returns bool."""
return os.geteuid() == 0
def show_selections(args, sources, ufd_dev, ufd_sources):
"""Show selections including non-specified options."""
# Sources
print_info('Sources')
for label in ufd_sources.keys():
if label in sources:
print_standard(' {label:<18} {path}'.format(
label=label+':',
path=sources[label],
))
else:
print_standard(' {label:<18} {YELLOW}Not Specified{CLEAR}'.format(
label=label+':',
**COLORS,
))
print_standard(' ')
# Destination
print_info('Destination')
cmd = [
'lsblk', '--nodeps', '--noheadings', '--paths',
'--output', 'NAME,FSTYPE,TRAN,SIZE,VENDOR,MODEL,SERIAL',
ufd_dev,
]
result = run_program(cmd, check=False, encoding='utf-8', errors='ignore')
print_standard(result.stdout.strip())
cmd = [
'lsblk', '--noheadings', '--paths',
'--output', 'NAME,SIZE,FSTYPE,LABEL,MOUNTPOINT',
ufd_dev,
]
result = run_program(cmd, check=False, encoding='utf-8', errors='ignore')
for line in result.stdout.splitlines()[1:]:
print_standard(line)
# Notes
if args['--update']:
print_warning('Updating kit in-place')
elif args['--use-mbr']:
print_warning('Formatting using legacy MBR')
print_standard(' ')
def unmount(mount_point):
"""Unmount mount_point."""
cmd = ['umount', mount_point]
run_program(cmd)
def update_boot_entries(boot_entries, boot_files, iso_label, ufd_label):
"""Update boot files for UFD usage"""
configs = []
# Find config files
for c_path, c_ext in boot_files.items():
c_path = find_path('/mnt/UFD{}'.format(c_path))
for item in os.scandir(c_path):
if item.name.lower().endswith(c_ext.lower()):
configs.append(item.path)
# Update Linux labels
cmd = [
'sed',
'--in-place',
'--regexp-extended',
's/(eSysRescueLiveCD|{})/{}/'.format(iso_label, ufd_label),
*configs,
]
run_program(cmd)
# Uncomment extra entries if present
for b_path, b_comment in boot_entries.items():
try:
find_path('/mnt/UFD{}'.format(b_path))
except (FileNotFoundError, NotADirectoryError):
# Entry not found, continue to next entry
continue
# Entry found, update config files
cmd = [
'sed',
'--in-place',
's/#{}#//'.format(b_comment),
*configs,
]
run_program(cmd, check=False)
def verify_sources(args, ufd_sources):
"""Check all sources and abort if necessary, returns dict."""
sources = OrderedDict()
for label, data in ufd_sources.items():
s_path = args[data['Arg']]
if s_path:
try:
s_path_obj = find_path(s_path)
except FileNotFoundError:
print_error('ERROR: {} not found: {}'.format(label, s_path))
abort(False)
if not is_valid_path(s_path_obj, data['Type']):
print_error('ERROR: Invalid {} source: {}'.format(label, s_path))
abort(False)
sources[label] = s_path_obj
return sources
def verify_ufd(dev_path):
"""Check that dev_path is a valid UFD, returns pathlib.Path obj."""
ufd_dev = None
try:
ufd_dev = find_path(dev_path)
except FileNotFoundError:
print_error('ERROR: UFD device not found: {}'.format(dev_path))
abort(False)
if not is_valid_path(ufd_dev, 'UFD'):
print_error('ERROR: Invalid UFD device: {}'.format(ufd_dev))
abort(False)
return ufd_dev
if __name__ == '__main__':
print("This file is not meant to be called directly.")

View file

@ -12,7 +12,7 @@ SOURCE_URLS = {
'BlueScreenView32': 'http://www.nirsoft.net/utils/bluescreenview.zip',
'BlueScreenView64': 'http://www.nirsoft.net/utils/bluescreenview-x64.zip',
'Caffeine': 'http://www.zhornsoftware.co.uk/caffeine/caffeine.zip',
'ClassicStartSkin': 'http://www.classicshell.net/forum/download/file.php?id=3001&sid=9a195960d98fd754867dcb63d9315335',
'ClassicStartSkin': 'https://coddec.github.io/Classic-Shell/www.classicshell.net/forum/download/fileb1ba.php?id=3001',
'Du': 'https://download.sysinternals.com/files/DU.zip',
'ERUNT': 'http://www.aumha.org/downloads/erunt.zip',
'ESET AVRemover32': 'https://download.eset.com/com/eset/tools/installers/av_remover/latest/avremover_nt32_enu.exe',

36
scripts/wk/__init__.py Normal file
View file

@ -0,0 +1,36 @@
"""WizardKit: wk module init"""
# vim: sts=2 sw=2 ts=2
from sys import version_info as version
from wk import cfg
from wk import debug
from wk import exe
from wk import graph
from wk import hw
from wk import io
from wk import kit
from wk import log
from wk import net
from wk import os
from wk import std
from wk import sw
from wk import tmux
# Check env
if version < (3, 7):
# Unsupported
raise RuntimeError(
f'This package is unsupported on Python {version.major}.{version.minor}'
)
# Init
try:
log.start()
except UserWarning as err:
std.print_warning(err)
if __name__ == '__main__':
print("This file is not meant to be called directly.")

View file

@ -0,0 +1,8 @@
"""WizardKit: cfg module init"""
from wk.cfg import ddrescue
from wk.cfg import hw
from wk.cfg import log
from wk.cfg import main
from wk.cfg import net
from wk.cfg import ufd

View file

@ -0,0 +1,67 @@
"""WizardKit: Config - ddrescue"""
# pylint: disable=bad-whitespace,line-too-long
# vim: sts=2 sw=2 ts=2
from collections import OrderedDict
# Layout
TMUX_SIDE_WIDTH = 21
TMUX_LAYOUT = OrderedDict({
'Source': {'height': 2, 'Check': True},
'Started': {'width': TMUX_SIDE_WIDTH, 'Check': True},
'Progress': {'width': TMUX_SIDE_WIDTH, 'Check': True},
})
# ddrescue
AUTO_PASS_THRESHOLDS = {
# NOTE: The scrape key is set to infinity to force a break
'read': 95,
'trim': 98,
'scrape': float('inf'),
}
DDRESCUE_SETTINGS = {
'Default': {
'--binary-prefixes': {'Selected': True, 'Hidden': True, },
'--data-preview': {'Selected': True, 'Value': '5', 'Hidden': True, },
'--idirect': {'Selected': True, },
'--odirect': {'Selected': True, },
'--max-error-rate': {'Selected': True, 'Value': '100MiB', },
'--max-read-rate': {'Selected': False, 'Value': '1MiB', },
'--min-read-rate': {'Selected': True, 'Value': '64KiB', },
'--reopen-on-error': {'Selected': True, },
'--retry-passes': {'Selected': True, 'Value': '0', },
'--reverse': {'Selected': False, },
'--test-mode': {'Selected': False, 'Value': 'test.map', },
'--timeout': {'Selected': True, 'Value': '30m', },
'-vvvv': {'Selected': True, 'Hidden': True, },
},
'Fast': {
'--max-error-rate': {'Selected': True, 'Value': '32MiB', },
'--min-read-rate': {'Selected': True, 'Value': '1MiB', },
'--reopen-on-error': {'Selected': False, },
'--timeout': {'Selected': True, 'Value': '5m', },
},
'Safe': {
'--max-read-rate': {'Selected': True, 'Value': '64MiB', },
'--min-read-rate': {'Selected': True, 'Value': '1KiB', },
'--reopen-on-error': {'Selected': True, },
'--timeout': {'Selected': False, 'Value': '30m', },
},
}
PARTITION_TYPES = {
'GPT': {
'NTFS': 'EBD0A0A2-B9E5-4433-87C0-68B6B72699C7', # Basic Data Partition
'VFAT': 'EBD0A0A2-B9E5-4433-87C0-68B6B72699C7', # Basic Data Partition
'EXFAT': 'EBD0A0A2-B9E5-4433-87C0-68B6B72699C7', # Basic Data Partition
},
'MBR': {
'EXFAT': 7, # 0x7
'NTFS': 7, # 0x7
'VFAT': 11, # 0xb
},
}
if __name__ == '__main__':
print("This file is not meant to be called directly.")

140
scripts/wk/cfg/hw.py Normal file
View file

@ -0,0 +1,140 @@
"""WizardKit: Config - Hardware"""
# pylint: disable=bad-whitespace,line-too-long
# vim: sts=2 sw=2 ts=2
import re
from collections import OrderedDict
# STATIC VARIABLES
ATTRIBUTE_COLORS = (
# NOTE: Ordered by ascending importance
('Warning', 'YELLOW'),
('Error', 'RED'),
('Maximum', 'PURPLE'),
)
# NOTE: Force 4K read block size for disks >= 3TB
BADBLOCKS_LARGE_DISK = 3 * 1024**4
CPU_CRITICAL_TEMP = 99
CPU_FAILURE_TEMP = 90
CPU_TEST_MINUTES = 7
KEY_NVME = 'nvme_smart_health_information_log'
KEY_SMART = 'ata_smart_attributes'
KNOWN_DISK_ATTRIBUTES = {
# NVMe
'critical_warning': {'Blocking': True, 'Warning': None, 'Error': 1, 'Maximum': None, },
'media_errors': {'Blocking': False, 'Warning': None, 'Error': 1, 'Maximum': None, },
'power_on_hours': {'Blocking': False, 'Warning': 17532, 'Error': 26298, 'Maximum': 100000,},
'unsafe_shutdowns': {'Blocking': False, 'Warning': 1, 'Error': None, 'Maximum': None, },
# SMART
5: {'Hex': '05', 'Blocking': True, 'Warning': None, 'Error': 1, 'Maximum': None, },
9: {'Hex': '09', 'Blocking': False, 'Warning': 17532, 'Error': 26298, 'Maximum': 100000,},
10: {'Hex': '10', 'Blocking': False, 'Warning': 1, 'Error': 10, 'Maximum': 10000, },
184: {'Hex': 'B8', 'Blocking': False, 'Warning': 1, 'Error': 10, 'Maximum': 10000, },
187: {'Hex': 'BB', 'Blocking': False, 'Warning': 1, 'Error': 10, 'Maximum': 10000, },
188: {'Hex': 'BC', 'Blocking': False, 'Warning': 1, 'Error': 10, 'Maximum': 10000, },
196: {'Hex': 'C4', 'Blocking': False, 'Warning': 1, 'Error': 10, 'Maximum': 10000, },
197: {'Hex': 'C5', 'Blocking': True, 'Warning': None, 'Error': 1, 'Maximum': None, },
198: {'Hex': 'C6', 'Blocking': True, 'Warning': None, 'Error': 1, 'Maximum': None, },
199: {'Hex': 'C7', 'Blocking': False, 'Warning': None, 'Error': 1, 'Maximum': None, },
201: {'Hex': 'C9', 'Blocking': False, 'Warning': None, 'Error': 1, 'Maximum': 10000, },
}
KNOWN_DISK_MODELS = {
# model_regex: model_attributes
r'CT(250|500|1000|2000)MX500SSD(1|4)': {
197: {'Warning': 1, 'Error': 2, 'Note': '(MX500 thresholds)',},
},
}
KNOWN_RAM_VENDOR_IDS = {
# https://github.com/hewigovens/hewigovens.github.com/wiki/Memory-vendor-code
'0x014F': 'Transcend',
'0x2C00': 'Micron',
'0x802C': 'Micron',
'0x80AD': 'Hynix',
'0x80CE': 'Samsung',
'0xAD00': 'Hynix',
'0xCE00': 'Samsung',
}
REGEX_POWER_ON_TIME = re.compile(
r'^(\d+)([Hh].*|\s+\(\d+\s+\d+\s+\d+\).*)'
)
SMC_IDS = {
# Sources: https://github.com/beltex/SMCKit/blob/master/SMCKit/SMC.swift
# http://www.opensource.apple.com/source/net_snmp/
# https://github.com/jedda/OSX-Monitoring-Tools
'TA0P': {'CPU Temp': False, 'Source': 'Ambient temp'},
'TA0S': {'CPU Temp': False, 'Source': 'PCIE Slot 1 Ambient'},
'TA1P': {'CPU Temp': False, 'Source': 'Ambient temp'},
'TA1S': {'CPU Temp': False, 'Source': 'PCIE Slot 1 PCB'},
'TA2S': {'CPU Temp': False, 'Source': 'PCIE Slot 2 Ambient'},
'TA3S': {'CPU Temp': False, 'Source': 'PCIE Slot 2 PCB'},
'TC0C': {'CPU Temp': True, 'Source': 'CPU Core 0'},
'TC0D': {'CPU Temp': True, 'Source': 'CPU die temp'},
'TC0H': {'CPU Temp': True, 'Source': 'CPU heatsink temp'},
'TC0P': {'CPU Temp': True, 'Source': 'CPU Ambient 1'},
'TC1C': {'CPU Temp': True, 'Source': 'CPU Core 1'},
'TC1P': {'CPU Temp': True, 'Source': 'CPU Ambient 2'},
'TC2C': {'CPU Temp': True, 'Source': 'CPU B Core 0'},
'TC2P': {'CPU Temp': True, 'Source': 'CPU B Ambient 1'},
'TC3C': {'CPU Temp': True, 'Source': 'CPU B Core 1'},
'TC3P': {'CPU Temp': True, 'Source': 'CPU B Ambient 2'},
'TCAC': {'CPU Temp': True, 'Source': 'CPU core from PCECI'},
'TCAH': {'CPU Temp': True, 'Source': 'CPU HeatSink'},
'TCBC': {'CPU Temp': True, 'Source': 'CPU B core from PCECI'},
'TCBH': {'CPU Temp': True, 'Source': 'CPU HeatSink'},
'Te1P': {'CPU Temp': False, 'Source': 'PCIE ambient temp'},
'Te1S': {'CPU Temp': False, 'Source': 'PCIE slot 1'},
'Te2S': {'CPU Temp': False, 'Source': 'PCIE slot 2'},
'Te3S': {'CPU Temp': False, 'Source': 'PCIE slot 3'},
'Te4S': {'CPU Temp': False, 'Source': 'PCIE slot 4'},
'TG0C': {'CPU Temp': False, 'Source': 'Mezzanine GPU Core'},
'TG0P': {'CPU Temp': False, 'Source': 'Mezzanine GPU Exhaust'},
'TH0P': {'CPU Temp': False, 'Source': 'Drive Bay 0'},
'TH1P': {'CPU Temp': False, 'Source': 'Drive Bay 1'},
'TH2P': {'CPU Temp': False, 'Source': 'Drive Bay 2'},
'TH3P': {'CPU Temp': False, 'Source': 'Drive Bay 3'},
'TH4P': {'CPU Temp': False, 'Source': 'Drive Bay 4'},
'TM0P': {'CPU Temp': False, 'Source': 'CPU DIMM Exit Ambient'},
'Tp0C': {'CPU Temp': False, 'Source': 'PSU1 Inlet Ambient'},
'Tp0P': {'CPU Temp': False, 'Source': 'PSU1 Inlet Ambient'},
'Tp1C': {'CPU Temp': False, 'Source': 'PSU1 Secondary Component'},
'Tp1P': {'CPU Temp': False, 'Source': 'PSU1 Primary Component'},
'Tp2P': {'CPU Temp': False, 'Source': 'PSU1 Secondary Component'},
'Tp3P': {'CPU Temp': False, 'Source': 'PSU2 Inlet Ambient'},
'Tp4P': {'CPU Temp': False, 'Source': 'PSU2 Primary Component'},
'Tp5P': {'CPU Temp': False, 'Source': 'PSU2 Secondary Component'},
'TS0C': {'CPU Temp': False, 'Source': 'CPU B DIMM Exit Ambient'},
}
TEMP_COLORS = {
float('-inf'): 'CYAN',
00: 'BLUE',
60: 'GREEN',
70: 'YELLOW',
80: 'ORANGE',
90: 'RED',
100: 'ORANGE_RED',
}
# THRESHOLDS: Rates used to determine HDD/SSD pass/fail
THRESH_HDD_MIN = 50 * 1024**2
THRESH_HDD_AVG_HIGH = 75 * 1024**2
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': 3, 'Check': True},
'badblocks': {'height': 5, 'Check': True},
'I/O Benchmark': {'height': 1000, 'Check': False},
})
if __name__ == '__main__':
print("This file is not meant to be called directly.")

18
scripts/wk/cfg/log.py Normal file
View file

@ -0,0 +1,18 @@
"""WizardKit: Config - Log"""
# vim: sts=2 sw=2 ts=2
DEBUG = {
'level': 'DEBUG',
'format': '[%(asctime)s %(levelname)s] [%(name)s.%(funcName)s] %(message)s',
'datefmt': '%Y-%m-%d %H%M%S%z',
}
DEFAULT = {
'level': 'INFO',
'format': '[%(asctime)s %(levelname)s] %(message)s',
'datefmt': '%Y-%m-%d %H%M%z',
}
if __name__ == '__main__':
print("This file is not meant to be called directly.")

36
scripts/wk/cfg/main.py Normal file
View file

@ -0,0 +1,36 @@
"""WizardKit: Config - Main
NOTE: Non-standard formating is used for BASH/BATCH/PYTHON compatibility
"""
# pylint: disable=bad-whitespace
# vim: sts=2 sw=2 ts=2
# Features
ENABLED_OPEN_LOGS=False
ENABLED_TICKET_NUMBERS=False
ENABLED_UPLOAD_DATA=False
# Main Kit
ARCHIVE_PASSWORD='Abracadabra'
KIT_NAME_FULL='WizardKit'
KIT_NAME_SHORT='WK'
SUPPORT_MESSAGE='Please let 2Shirt know by opening an issue on GitHub'
# Text Formatting
INDENT=4
WIDTH=32
# Live Linux
ROOT_PASSWORD='Abracadabra'
TECH_PASSWORD='Abracadabra'
# Time Zones
## See 'timedatectl list-timezones' for valid Linux values
## See 'tzutil /l' for valid Windows values
LINUX_TIME_ZONE='America/Denver'
WINDOWS_TIME_ZONE='Mountain Standard Time'
if __name__ == '__main__':
print("This file is not meant to be called directly.")

35
scripts/wk/cfg/net.py Normal file
View file

@ -0,0 +1,35 @@
"""WizardKit: Config - Net"""
# pylint: disable=bad-whitespace
# vim: sts=2 sw=2 ts=2
# Servers
BACKUP_SERVERS = {
#'Server One': {
# 'Address': '10.0.0.10',
# 'Share': 'Backups',
# 'RO-User': 'restore',
# 'RO-Pass': 'Abracadabra',
# 'RW-User': 'backup',
# 'RW-Pass': 'Abracadabra',
# },
#'Server Two': {
# 'Address': 'servertwo.example.com',
# 'Share': 'Backups',
# 'RO-User': 'restore',
# 'RO-Pass': 'Abracadabra',
# 'RW-User': 'backup',
# 'RW-Pass': 'Abracadabra',
# },
}
CRASH_SERVER = {
#'Name': 'CrashServer',
#'Url': '',
#'User': '',
#'Pass': '',
#'Headers': {'X-Requested-With': 'XMLHttpRequest'},
}
if __name__ == '__main__':
print("This file is not meant to be called directly.")

View file

@ -1,44 +1,15 @@
'''Wizard Kit: Settings - UFD'''
# pylint: disable=C0326,E0611
"""WizardKit: Config - UFD"""
# pylint: disable=bad-whitespace
# vim: sts=2 sw=2 ts=2
from collections import OrderedDict
from settings.main import KIT_NAME_FULL,KIT_NAME_SHORT
from wk.cfg.main import KIT_NAME_FULL
# General
DOCSTRING = '''WizardKit: Build UFD
Usage:
build-ufd [options] --ufd-device PATH --linux PATH
[--linux-minimal PATH]
[--main-kit PATH]
[--winpe PATH]
[--eset PATH]
[--hdclone PATH]
[--extra-dir PATH]
build-ufd (-h | --help)
Options:
-c PATH, --hdclone PATH
-d PATH, --linux-dgpu PATH
-e PATH, --extra-dir PATH
-k PATH, --main-kit PATH
-l PATH, --linux PATH
-m PATH, --linux-minimal PATH
-s PATH, --eset PATH
-u PATH, --ufd-device PATH
-w PATH, --winpe PATH
-h --help Show this page
-M --use-mbr Use real MBR instead of GPT w/ Protective MBR
-F --force Bypass all confirmation messages. USE WITH EXTREME CAUTION!
-U --update Don't format device, just update
'''
ISO_LABEL = '{}_LINUX'.format(KIT_NAME_SHORT)
UFD_LABEL = '{}_UFD'.format(KIT_NAME_SHORT)
UFD_SOURCES = OrderedDict({
SOURCES = OrderedDict({
'Linux': {'Arg': '--linux', 'Type': 'ISO'},
'Linux (dGPU)': {'Arg': '--linux-dgpu', 'Type': 'ISO'},
'Linux (Minimal)': {'Arg': '--linux-minimal', 'Type': 'ISO'},
'WinPE': {'Arg': '--winpe', 'Type': 'ISO'},
'ESET SysRescue': {'Arg': '--eset', 'Type': 'IMG'},
@ -52,7 +23,6 @@ BOOT_ENTRIES = {
# Path to check: Comment to remove
'/arch_minimal': 'UFD-MINIMAL',
'/casper': 'UFD-ESET',
'/dgpu': 'UFD-DGPU',
'/kernel.map': 'UFD-HDCLONE',
'/sources/boot.wim': 'UFD-WINPE',
}
@ -87,12 +57,6 @@ ITEMS = {
('/EFI/boot', '/EFI/'),
('/EFI/memtest86', '/EFI/'),
),
'Linux (dGPU)': (
('/arch/boot/x86_64/archiso.img', '/dgpu/'),
('/arch/boot/x86_64/vmlinuz', '/dgpu/'),
('/arch/pkglist.x86_64.txt', '/dgpu/'),
('/arch/x86_64', '/dgpu/'),
),
'Linux (Minimal)': (
('/arch/boot/x86_64/archiso.img', '/arch_minimal/'),
('/arch/boot/x86_64/vmlinuz', '/arch_minimal/'),
@ -100,7 +64,7 @@ ITEMS = {
('/arch/x86_64', '/arch_minimal/'),
),
'Main Kit': (
('/', '/{}/'.format(KIT_NAME_FULL)),
('/', f'/{KIT_NAME_FULL}/'),
),
'WinPE': (
('/bootmgr', '/'),
@ -124,12 +88,11 @@ ITEMS_HIDDEN = (
# Linux (all versions)
'arch',
'arch_minimal',
'dgpu',
'EFI',
'isolinux',
# Main Kit
'{}/.bin'.format(KIT_NAME_FULL),
'{}/.cbin'.format(KIT_NAME_FULL),
f'{KIT_NAME_FULL}/.bin',
f'{KIT_NAME_FULL}/.cbin',
# WinPE
'boot',
'bootmgr',
@ -139,5 +102,6 @@ ITEMS_HIDDEN = (
'sources',
)
if __name__ == '__main__':
print("This file is not meant to be called directly.")

45
scripts/wk/debug.py Normal file
View file

@ -0,0 +1,45 @@
"""WizardKit: Debug Functions"""
# pylint: disable=invalid-name
# vim: sts=2 sw=2 ts=2
# Classes
class Debug():
# pylint: disable=too-few-public-methods
"""Object used when dumping debug data."""
def method(self):
"""Dummy method used to identify functions vs data."""
# STATIC VARIABLES
DEBUG_CLASS = Debug()
METHOD_TYPE = type(DEBUG_CLASS.method)
# Functions
def generate_object_report(obj, indent=0):
"""Generate debug report for obj, returns list."""
report = []
# Dump object data
for name in dir(obj):
attr = getattr(obj, name)
# Skip methods and private attributes
if isinstance(attr, METHOD_TYPE) or name.startswith('_'):
continue
# Add attribute to report (expanded if necessary)
if isinstance(attr, dict):
report.append(f'{name}:')
for key, value in sorted(attr.items()):
report.append(f'{" "*(indent+1)}{key}: {str(value)}')
else:
report.append(f'{" "*indent}{name}: {str(attr)}')
# Done
return report
if __name__ == '__main__':
print("This file is not meant to be called directly.")

222
scripts/wk/exe.py Normal file
View file

@ -0,0 +1,222 @@
"""WizardKit: Execution functions"""
#vim: sts=2 sw=2 ts=2
import json
import logging
import re
import subprocess
from threading import Thread
from queue import Queue, Empty
import psutil
# STATIC VARIABLES
LOG = logging.getLogger(__name__)
# Classes
class NonBlockingStreamReader():
"""Class to allow non-blocking reads from a stream."""
# pylint: disable=too-few-public-methods
# Credits:
## https://gist.github.com/EyalAr/7915597
## https://stackoverflow.com/a/4896288
def __init__(self, stream):
self.stream = stream
self.queue = Queue()
def populate_queue(stream, queue):
"""Collect lines from stream and put them in queue."""
while True:
line = stream.read(1)
if line:
queue.put(line)
self.thread = start_thread(
populate_queue,
args=(self.stream, self.queue),
)
def read(self, timeout=None):
"""Read from queue if possible, returns item from queue."""
try:
return self.queue.get(block=timeout is not None, timeout=timeout)
except Empty:
return None
def save_to_file(self, proc, out_path):
"""Continuously save output to file while proc is running."""
while proc.poll() is None:
out = b''
out_bytes = b''
while out is not None:
out = self.read(0.1)
if out:
out_bytes += out
with open(out_path, 'a') as _f:
_f.write(out_bytes.decode('utf-8', errors='ignore'))
# Functions
def build_cmd_kwargs(cmd, minimized=False, pipe=True, shell=False, **kwargs):
"""Build kwargs for use by subprocess functions, returns dict.
Specifically subprocess.run() and subprocess.Popen().
NOTE: If no encoding specified then UTF-8 will be used.
"""
LOG.debug(
'cmd: %s, minimized: %s, pipe: %s, shell: %s',
cmd, minimized, pipe, shell,
)
LOG.debug('kwargs: %s', kwargs)
cmd_kwargs = {
'args': cmd,
'shell': shell,
}
# Add additional kwargs if applicable
for key in 'check cwd encoding errors stderr stdin stdout'.split():
if key in kwargs:
cmd_kwargs[key] = kwargs[key]
# Default to UTF-8 encoding
if not ('encoding' in cmd_kwargs or 'errors' in cmd_kwargs):
cmd_kwargs['encoding'] = 'utf-8'
cmd_kwargs['errors'] = 'ignore'
# Start minimized
if minimized:
startupinfo = subprocess.STARTUPINFO()
startupinfo.dwFlags = subprocess.STARTF_USESHOWWINDOW
startupinfo.wShowWindow = 6
cmd_kwargs['startupinfo'] = startupinfo
# Pipe output
if pipe:
cmd_kwargs['stderr'] = subprocess.PIPE
cmd_kwargs['stdout'] = subprocess.PIPE
# Done
LOG.debug('cmd_kwargs: %s', cmd_kwargs)
return cmd_kwargs
def get_json_from_command(cmd, check=True, encoding='utf-8', errors='ignore'):
"""Capture JSON content from cmd output, returns dict.
If the data can't be decoded then either an exception is raised
or an empty dict is returned depending on errors.
"""
json_data = {}
try:
proc = run_program(cmd, check=check, encoding=encoding, errors=errors)
json_data = json.loads(proc.stdout)
except (subprocess.CalledProcessError, json.decoder.JSONDecodeError):
if errors != 'ignore':
raise
return json_data
def get_procs(name, exact=True):
"""Get process object(s) based on name, returns list of proc objects."""
LOG.debug('name: %s, exact: %s', name, exact)
processes = []
regex = f'^{name}$' if exact else name
# Iterate over all processes
for proc in psutil.process_iter():
if re.search(regex, proc.name(), re.IGNORECASE):
processes.append(proc)
# Done
return processes
def kill_procs(name, exact=True, force=False, timeout=30):
"""Kill all processes matching name (case-insensitively).
NOTE: Under Posix systems this will send SIGINT to allow processes
to gracefully exit.
If force is True then it will wait until timeout specified and then
send SIGKILL to any processes still alive.
"""
LOG.debug(
'name: %s, exact: %s, force: %s, timeout: %s',
name, exact, force, timeout,
)
target_procs = get_procs(name, exact=exact)
for proc in target_procs:
proc.terminate()
# Force kill if necesary
if force:
results = psutil.wait_procs(target_procs, timeout=timeout)
for proc in results[1]: # Alive processes
proc.kill()
def popen_program(cmd, minimized=False, pipe=False, shell=False, **kwargs):
"""Run program and return a subprocess.Popen object."""
LOG.debug(
'cmd: %s, minimized: %s, pipe: %s, shell: %s',
cmd, minimized, pipe, shell,
)
LOG.debug('kwargs: %s', kwargs)
cmd_kwargs = build_cmd_kwargs(
cmd,
minimized=minimized,
pipe=pipe,
shell=shell,
**kwargs)
# Ready to run program
return subprocess.Popen(**cmd_kwargs)
def run_program(cmd, check=True, pipe=True, shell=False, **kwargs):
# pylint: disable=subprocess-run-check
"""Run program and return a subprocess.CompletedProcess object."""
LOG.debug(
'cmd: %s, check: %s, pipe: %s, shell: %s',
cmd, check, pipe, shell,
)
LOG.debug('kwargs: %s', kwargs)
cmd_kwargs = build_cmd_kwargs(
cmd,
check=check,
pipe=pipe,
shell=shell,
**kwargs)
# Ready to run program
return subprocess.run(**cmd_kwargs)
def start_thread(function, args=None, daemon=True):
"""Run function as thread in background, returns Thread object."""
args = args if args else []
thread = Thread(target=function, args=args, daemon=daemon)
thread.start()
return thread
def wait_for_procs(name, exact=True, timeout=None):
"""Wait for all process matching name."""
LOG.debug('name: %s, exact: %s, timeout: %s', name, exact, timeout)
target_procs = get_procs(name, exact=exact)
results = psutil.wait_procs(target_procs, timeout=timeout)
# Raise exception if necessary
if results[1]: # Alive processes
raise psutil.TimeoutExpired(name=name, seconds=timeout)
if __name__ == '__main__':
print("This file is not meant to be called directly.")

151
scripts/wk/graph.py Normal file
View file

@ -0,0 +1,151 @@
"""WizardKit: Graph Functions"""
# pylint: disable=bad-whitespace
# vim: sts=2 sw=2 ts=2
import logging
from wk.std import color_string
# STATIC VARIABLES
LOG = logging.getLogger(__name__)
GRAPH_HORIZONTAL = ('', '', '', '', '', '', '', '')
GRAPH_VERTICAL = (
'', '', '', '',
'', '', '', '',
'█▏', '█▎', '█▍', '█▌',
'█▋', '█▊', '█▉', '██',
'██▏', '██▎', '██▍', '██▌',
'██▋', '██▊', '██▉', '███',
'███▏', '███▎', '███▍', '███▌',
'███▋', '███▊', '███▉', '████',
)
# SCALE_STEPS: These scales allow showing differences between HDDs and SSDs
# on the same graph.
SCALE_STEPS = {
8: [2**(0.56*(x+1))+(16*(x+1)) for x in range(8)],
16: [2**(0.56*(x+1))+(16*(x+1)) for x in range(16)],
32: [2**(0.56*(x+1)/2)+(16*(x+1)/2) for x in range(32)],
}
# THRESHOLDS: These are the rate_list (in MB/s) used to color graphs
THRESH_FAIL = 65 * 1024**2
THRESH_WARN = 135 * 1024**2
THRESH_GREAT = 750 * 1024**2
# Functions
def generate_horizontal_graph(rate_list, graph_width=40, oneline=False):
"""Generate horizontal graph from rate_list, returns list."""
graph = ['', '', '', '']
scale = 8 if oneline else 32
# Build graph
for rate in merge_rates(rate_list, graph_width=graph_width):
step = get_graph_step(rate, scale=scale)
# Set color
rate_color = None
if rate < THRESH_FAIL:
rate_color = 'RED'
elif rate < THRESH_WARN:
rate_color = 'YELLOW'
elif rate > THRESH_GREAT:
rate_color = 'GREEN'
# Build graph
full_block = color_string((GRAPH_HORIZONTAL[-1],), (rate_color,))
if step >= 24:
graph[0] += 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[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[3] += full_block
else:
graph[0] += ' '
graph[1] += ' '
graph[2] += ' '
graph[3] += color_string((GRAPH_HORIZONTAL[step],), (rate_color,))
# Done
if oneline:
graph = graph[-1:]
return graph
def get_graph_step(rate, scale=16):
"""Get graph step based on rate and scale, returns int."""
rate_in_mb = rate / (1024**2)
step = 0
# Iterate over scale_steps backwards
for _r in range(scale-1, -1, -1):
if rate_in_mb >= SCALE_STEPS[scale][_r]:
step = _r
break
# Done
return step
def merge_rates(rates, graph_width=40):
"""Merge rates to have entries equal to the width, returns list."""
merged_rates = []
offset = 0
slice_width = int(len(rates) / graph_width)
# Merge rates
for _i in range(graph_width):
merged_rates.append(sum(rates[offset:offset+slice_width])/slice_width)
offset += slice_width
# Done
return merged_rates
def vertical_graph_line(percent, rate, scale=32):
"""Build colored graph string using thresholds, returns str."""
color_bar = None
color_rate = None
step = get_graph_step(rate, scale=scale)
# Set colors
if rate < THRESH_FAIL:
color_bar = 'RED'
color_rate = 'YELLOW'
elif rate < THRESH_WARN:
color_bar = 'YELLOW'
color_rate = 'YELLOW'
elif rate > THRESH_GREAT:
color_bar = 'GREEN'
color_rate = 'GREEN'
# Build string
line = color_string(
strings=(
f'{percent:5.1f}%',
f'{GRAPH_VERTICAL[step]:<4}',
f'{rate/(1000**2):6.1f} MB/s',
),
colors=(
None,
color_bar,
color_rate,
),
sep=' ',
)
# Done
return line
if __name__ == '__main__':
print("This file is not meant to be called directly.")

View file

@ -0,0 +1,6 @@
"""WizardKit: hw module init"""
from wk.hw import ddrescue
from wk.hw import diags
from wk.hw import obj
from wk.hw import sensors

2061
scripts/wk/hw/ddrescue.py Normal file

File diff suppressed because it is too large Load diff

1369
scripts/wk/hw/diags.py Normal file

File diff suppressed because it is too large Load diff

844
scripts/wk/hw/obj.py Normal file
View file

@ -0,0 +1,844 @@
"""WizardKit: Hardware objects (mostly)"""
# vim: sts=2 sw=2 ts=2
import logging
import pathlib
import plistlib
import re
from collections import OrderedDict
from wk.cfg.hw import (
ATTRIBUTE_COLORS,
KEY_NVME,
KEY_SMART,
KNOWN_DISK_ATTRIBUTES,
KNOWN_DISK_MODELS,
KNOWN_RAM_VENDOR_IDS,
REGEX_POWER_ON_TIME,
)
from wk.cfg.main import KIT_NAME_SHORT
from wk.exe import get_json_from_command, run_program
from wk.std import (
PLATFORM,
bytes_to_string,
color_string,
sleep,
string_to_bytes,
)
# STATIC VARIABLES
LOG = logging.getLogger(__name__)
NVME_WARNING_KEYS = (
'spare_below_threshold',
'reliability_degraded',
'volatile_memory_backup_failed',
)
WK_LABEL_REGEX = re.compile(
fr'{KIT_NAME_SHORT}_(LINUX|UFD)',
re.IGNORECASE,
)
# Exception Classes
class CriticalHardwareError(RuntimeError):
"""Exception used for critical hardware failures."""
class SMARTNotSupportedError(TypeError):
"""Exception used for disks lacking SMART support."""
class SMARTSelfTestInProgressError(RuntimeError):
"""Exception used when a SMART self-test is in progress."""
# Classes
class BaseObj():
"""Base object for tracking device data."""
def __init__(self):
self.tests = OrderedDict()
def all_tests_passed(self):
"""Check if all tests passed, returns bool."""
return all([results.passed for results in self.tests.values()])
def any_test_failed(self):
"""Check if any test failed, returns bool."""
return any([results.failed for results in self.tests.values()])
class CpuRam(BaseObj):
"""Object for tracking CPU & RAM specific data."""
def __init__(self):
super().__init__()
self.description = 'Unknown'
self.details = {}
self.ram_total = 'Unknown'
self.ram_dimms = []
self.tests = OrderedDict()
# Update details
self.get_cpu_details()
self.get_ram_details()
def generate_report(self):
"""Generate CPU & RAM report, returns list."""
report = []
report.append(color_string('Device', 'BLUE'))
report.append(f' {self.description}')
# Include RAM details
report.append(color_string('RAM', 'BLUE'))
report.append(f' {self.ram_total} ({", ".join(self.ram_dimms)})')
# Tests
for test in self.tests.values():
report.extend(test.report)
return report
def get_cpu_details(self):
"""Get CPU details using OS specific methods."""
if PLATFORM == 'Darwin':
cmd = 'sysctl -n machdep.cpu.brand_string'.split()
proc = run_program(cmd, check=False)
self.description = re.sub(r'\s+', ' ', proc.stdout.strip())
elif PLATFORM == 'Linux':
cmd = ['lscpu', '--json']
json_data = get_json_from_command(cmd)
for line in json_data.get('lscpu', [{}]):
_field = line.get('field', '').replace(':', '')
_data = line.get('data', '')
if not (_field or _data):
# Skip
continue
self.details[_field] = _data
self.description = self.details.get('Model name', '')
# Replace empty description
if not self.description:
self.description = 'Unknown CPU'
def get_ram_details(self):
"""Get RAM details using OS specific methods."""
if PLATFORM == 'Darwin':
dimm_list = get_ram_list_macos()
elif PLATFORM == 'Linux':
dimm_list = get_ram_list_linux()
details = {'Total': 0}
for dimm_details in dimm_list:
size, manufacturer = dimm_details
if size <= 0:
# Skip empty DIMMs
continue
description = f'{bytes_to_string(size)} {manufacturer}'
details['Total'] += size
if description in details:
details[description] += 1
else:
details[description] = 1
# Save details
self.ram_total = bytes_to_string(details.pop('Total', 0))
self.ram_dimms = [
f'{count}x {desc}' for desc, count in sorted(details.items())
]
class Disk(BaseObj):
"""Object for tracking disk specific data."""
def __init__(self, path):
super().__init__()
self.attributes = {}
self.description = 'Unknown'
self.details = {}
self.notes = []
self.path = pathlib.Path(path).resolve()
self.smartctl = {}
self.tests = OrderedDict()
# Update details
self.get_details()
self.enable_smart()
self.update_smart_details()
if not self.is_4k_aligned():
self.add_note('One or more partitions are not 4K aligned', 'YELLOW')
def abort_self_test(self):
"""Abort currently running non-captive self-test."""
cmd = ['sudo', 'smartctl', '--abort', self.path]
run_program(cmd, check=False)
def add_note(self, note, color=None):
"""Add note that will be included in the disk report."""
if color:
note = color_string(note, color)
if note not in self.notes:
self.notes.append(note)
self.notes.sort()
def check_attributes(self, only_blocking=False):
"""Check if any known attributes are failing, returns bool."""
attributes_ok = True
known_attributes = get_known_disk_attributes(self.details['model'])
for attr, value in self.attributes.items():
# Skip unknown attributes
if attr not in known_attributes:
continue
# Get thresholds
blocking_attribute = known_attributes[attr].get('Blocking', False)
err_thresh = known_attributes[attr].get('Error', None)
max_thresh = known_attributes[attr].get('Maximum', None)
if not max_thresh:
max_thresh = float('inf')
# Skip non-blocking attributes if necessary
if only_blocking and not blocking_attribute:
continue
# Skip informational attributes
if not err_thresh:
continue
# Check attribute
if err_thresh <= value['raw'] < max_thresh:
attributes_ok = False
# Done
return attributes_ok
def disable_disk_tests(self):
"""Disable all tests."""
LOG.warning('Disabling all tests for: %s', self.path)
for test in self.tests.values():
if test.status in ('Pending', 'Working'):
test.set_status('Denied')
test.disabled = True
def enable_smart(self):
"""Try enabling SMART for this disk."""
cmd = [
'sudo',
'smartctl',
'--tolerance=permissive',
'--smart=on',
self.path,
]
run_program(cmd, check=False)
def generate_attribute_report(self):
"""Generate attribute report, returns list."""
known_attributes = get_known_disk_attributes(self.details['model'])
report = []
for attr, value in sorted(self.attributes.items()):
note = ''
value_color = 'GREEN'
# Skip attributes not in our list
if attr not in known_attributes:
continue
# Check for attribute note
note = known_attributes[attr].get('Note', '')
# ID / Name
label = f'{attr:>3}'
if isinstance(attr, int):
# Assuming SMART, include hex ID and name
label += f' / {str(hex(attr))[2:].upper():0>2}: {value["name"]}'
label = f' {label.replace("_", " "):38}'
# Value color
for threshold, color in ATTRIBUTE_COLORS:
threshold_val = known_attributes[attr].get(threshold, None)
if threshold_val and value['raw'] >= threshold_val:
value_color = color
if threshold == 'Error':
note = '(failed)'
elif threshold == 'Maximum':
note = '(invalid?)'
# 199/C7 warning
if str(attr) == '199' and value['raw'] > 0:
note = '(bad cable?)'
# Build colored string and append to report
line = color_string(
[label, value['raw_str'], note],
[None, value_color, 'YELLOW'],
)
report.append(line)
# Done
return report
def generate_report(self, header=True):
"""Generate Disk report, returns list."""
report = []
if header:
report.append(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.extend(self.generate_attribute_report())
# Notes
if self.notes:
report.append(color_string('Notes', 'BLUE'))
for note in self.notes:
report.append(f' {note}')
# Tests
for test in self.tests.values():
report.extend(test.report)
return report
def get_details(self):
"""Get disk details using OS specific methods.
Required details default to generic descriptions
and are converted to the correct type.
"""
if PLATFORM == 'Darwin':
self.details = get_disk_details_macos(self.path)
elif PLATFORM == 'Linux':
self.details = get_disk_details_linux(self.path)
# Set necessary details
self.details['bus'] = str(self.details.get('bus', '???')).upper()
self.details['bus'] = self.details['bus'].replace('IMAGE', 'Image')
self.details['bus'] = self.details['bus'].replace('NVME', 'NVMe')
self.details['log-sec'] = self.details.get('log-sec', 512)
self.details['model'] = self.details.get('model', 'Unknown Model')
self.details['name'] = self.details.get('name', self.path)
self.details['phy-sec'] = self.details.get('phy-sec', 512)
self.details['serial'] = self.details.get('serial', 'Unknown Serial')
self.details['size'] = self.details.get('size', -1)
self.details['ssd'] = self.details.get('ssd', False)
# Ensure certain attributes types
for attr in ['bus', 'model', 'name', 'serial']:
if not isinstance(self.details[attr], str):
self.details[attr] = str(self.details[attr])
for attr in ['phy-sec', 'size']:
if not isinstance(self.details[attr], int):
try:
self.details[attr] = int(self.details[attr])
except (TypeError, ValueError):
LOG.error('Invalid disk %s: %s', attr, self.details[attr])
self.details[attr] = -1
# Set description
self.description = '{size_str} ({bus}) {model} {serial}'.format(
size_str=bytes_to_string(self.details['size'], use_binary=False),
**self.details,
)
def get_labels(self):
"""Build list of labels for this disk, returns list."""
labels = []
# Add all labels from lsblk
for disk in [self.details, *self.details.get('children', [])]:
labels.append(disk.get('label', ''))
labels.append(disk.get('partlabel', ''))
# Remove empty labels
labels = [str(label) for label in labels if label]
# Done
return labels
def get_smart_self_test_details(self):
"""Shorthand to get deeply nested self-test details, returns dict."""
details = {}
try:
details = self.smartctl['ata_smart_data']['self_test']
except (KeyError, TypeError):
# Assuming disk lacks SMART support, ignore and return empty dict.
pass
# Done
return details
def is_4k_aligned(self):
"""Check that all disk partitions are aligned, returns bool."""
aligned = True
if PLATFORM == 'Darwin':
aligned = is_4k_aligned_macos(self.details)
elif PLATFORM == 'Linux':
aligned = is_4k_aligned_linux(self.path, self.details['phy-sec'])
#TODO: Add checks for other OS
return aligned
def safety_checks(self):
"""Run safety checks and raise an exception if necessary."""
blocking_event_encountered = False
self.update_smart_details()
# Attributes
if not self.check_attributes(only_blocking=True):
blocking_event_encountered = True
LOG.error('%s: Blocked for failing attribute(s)', self.path)
# NVMe status
nvme_status = self.smartctl.get('smart_status', {}).get('nvme', {})
if nvme_status.get('media_read_only', False):
blocking_event_encountered = True
msg = 'Media has been placed in read-only mode'
self.add_note(msg, 'RED')
LOG.error('%s %s', self.path, msg)
for key in NVME_WARNING_KEYS:
if nvme_status.get(key, False):
msg = key.replace('_', ' ')
self.add_note(msg, 'YELLOW')
LOG.warning('%s %s', self.path, msg)
# SMART overall assessment
smart_passed = True
try:
smart_passed = self.smartctl['smart_status']['passed']
except (KeyError, TypeError):
# Assuming disk doesn't support SMART overall assessment
pass
if not smart_passed:
blocking_event_encountered = True
msg = 'SMART overall self-assessment: Failed'
self.add_note(msg, 'RED')
LOG.error('%s %s', self.path, msg)
# Raise blocking exception if necessary
if blocking_event_encountered:
raise CriticalHardwareError(f'Critical error(s) for: {self.path}')
# SMART self-test status
test_details = self.get_smart_self_test_details()
if 'remaining_percent' in test_details.get('status', ''):
msg = f'SMART self-test in progress for: {self.path}'
LOG.error(msg)
raise SMARTSelfTestInProgressError(msg)
def run_self_test(self, log_path):
"""Run disk self-test and check if it passed, returns bool.
NOTE: This function is here to reserve a place for future
NVMe self-tests announced in NVMe spec v1.3.
"""
result = self.run_smart_self_test(log_path)
return result
def run_smart_self_test(self, log_path):
"""Run SMART self-test and check if it passed, returns bool.
NOTE: An exception will be raised if the disk lacks SMART support.
"""
finished = False
result = None
started = False
status_str = 'Starting self-test...'
test_details = self.get_smart_self_test_details()
test_minutes = 15
# Check if disk supports self-tests
if not test_details:
raise SMARTNotSupportedError(
f'SMART self-test not supported for {self.path}')
# Get real test length
test_minutes = test_details.get('polling_minutes', {}).get('short', 5)
test_minutes = int(test_minutes) + 10
# Start test
cmd = [
'sudo',
'smartctl',
'--tolerance=normal',
'--test=short',
self.path,
]
run_program(cmd, check=False)
# Monitor progress (in five second intervals)
for _i in range(int(test_minutes*60/5)):
sleep(5)
# Update status
self.update_smart_details()
test_details = self.get_smart_self_test_details()
# Check test progress
if started:
status_str = test_details.get('status', {}).get('string', 'Unknown')
status_str = status_str.capitalize()
# Update log
with open(log_path, 'w') as _f:
_f.write(f'SMART self-test status for {self.path}:\n {status_str}')
# Check if finished
if 'remaining_percent' not in test_details['status']:
finished = True
break
elif 'remaining_percent' in test_details['status']:
started = True
# Check result
if finished:
result = test_details.get('status', {}).get('passed', False)
elif started:
raise TimeoutError(f'SMART self-test timed out for {self.path}')
# Done
return result
def update_smart_details(self):
"""Update SMART details via smartctl."""
self.attributes = {}
cmd = [
'sudo',
'smartctl',
'--tolerance=verypermissive',
'--all',
'--json',
self.path,
]
self.smartctl = get_json_from_command(cmd, check=False)
# Check for attributes
if KEY_NVME in self.smartctl:
for name, value in self.smartctl[KEY_NVME].items():
try:
self.attributes[name] = {
'name': name,
'raw': int(value),
'raw_str': str(value),
}
except ValueError:
# Ignoring invalid attribute
LOG.error('Invalid NVMe attribute: %s %s', name, value)
elif KEY_SMART in self.smartctl:
for attribute in self.smartctl[KEY_SMART].get('table', {}):
try:
_id = int(attribute['id'])
except (KeyError, ValueError):
# Ignoring invalid attribute
LOG.error('Invalid SMART attribute: %s', attribute)
continue
name = str(attribute.get('name', 'Unknown')).replace('_', ' ').title()
raw = int(attribute.get('raw', {}).get('value', -1))
raw_str = attribute.get('raw', {}).get('string', 'Unknown')
# Fix power-on time
match = REGEX_POWER_ON_TIME.match(raw_str)
if _id == 9 and match:
raw = int(match.group(1))
# Add to dict
self.attributes[_id] = {
'name': name, 'raw': raw, 'raw_str': raw_str}
# Add note if necessary
if not self.attributes:
self.add_note('No NVMe or SMART data available', 'YELLOW')
class Test():
# pylint: disable=too-few-public-methods
"""Object for tracking test specific data."""
def __init__(self, dev, label):
self.dev = dev
self.disabled = False
self.failed = False
self.label = label
self.passed = False
self.report = []
self.status = 'Pending'
def set_status(self, status):
"""Update status string."""
if self.disabled:
# Don't change status if disabled
return
self.status = status
# Functions
def get_disk_details_linux(path):
"""Get disk details using lsblk, returns dict."""
cmd = ['lsblk', '--bytes', '--json', '--output-all', '--paths', path]
json_data = get_json_from_command(cmd, check=False)
details = json_data.get('blockdevices', [{}])[0]
# Fix details
for dev in [details, *details.get('children', [])]:
dev['bus'] = dev.pop('tran', '???')
dev['parent'] = dev.pop('pkname', None)
dev['ssd'] = not dev.pop('rota', True)
if 'loop' in str(path) and dev['bus'] is None:
dev['bus'] = 'Image'
dev['model'] = ''
dev['serial'] = ''
# Done
return details
def get_disk_details_macos(path):
"""Get disk details using diskutil, returns dict."""
details = {}
# Get "list" details
cmd = ['diskutil', 'list', '-plist', path]
proc = run_program(cmd, check=False, encoding=None, errors=None)
try:
plist_data = plistlib.loads(proc.stdout)
except (TypeError, ValueError):
# Invalid / corrupt plist data? return empty dict to avoid crash
LOG.error('Failed to get diskutil list for %s', path)
return details
# Parse "list" details
details = plist_data.get('AllDisksAndPartitions', [{}])[0]
details['children'] = details.pop('Partitions', [])
details['path'] = path
for child in details['children']:
child['path'] = path.with_name(child.get('DeviceIdentifier', 'null'))
# Get "info" details
for dev in [details, *details['children']]:
cmd = ['diskutil', 'info', '-plist', dev['path']]
proc = run_program(cmd, check=False, encoding=None, errors=None)
try:
plist_data = plistlib.loads(proc.stdout)
except (TypeError, ValueError):
LOG.error('Failed to get diskutil info for %s', path)
continue #Skip
# Parse "info" details
dev.update(plist_data)
dev['bus'] = dev.pop('BusProtocol', '???')
dev['fstype'] = dev.pop('FilesystemType', '')
dev['label'] = dev.pop('VolumeName', '')
dev['model'] = dev.pop('MediaName', 'Unknown')
dev['mountpoint'] = dev.pop('MountPoint', '')
dev['phy-sec'] = dev.pop('DeviceBlockSize', 512)
dev['serial'] = get_disk_serial_macos(dev['path'])
dev['size'] = dev.pop('Size', -1)
dev['ssd'] = dev.pop('SolidState', False)
dev['vendor'] = ''
if not dev.get('WholeDisk', True):
dev['parent'] = dev.pop('ParentWholeDisk', None)
# Done
return details
def get_disk_serial_macos(path):
"""Get disk serial using system_profiler, returns str."""
cmd = ['sudo', 'smartctl', '--info', '--json', path]
smart_info = get_json_from_command(cmd)
return smart_info.get('serial_number', 'Unknown Serial')
def get_disks(skip_kits=False):
"""Get disks using OS-specific methods, returns list."""
disks = []
if PLATFORM == 'Darwin':
disks = get_disks_macos()
elif PLATFORM == 'Linux':
disks = get_disks_linux()
# Skip WK disks
if skip_kits:
disks = [
disk_obj for disk_obj in disks
if not any(
[WK_LABEL_REGEX.search(label) for label in disk_obj.get_labels()]
)
]
# Done
return disks
def get_disks_linux():
"""Get disks via lsblk, returns list."""
cmd = ['lsblk', '--json', '--nodeps', '--paths']
disks = []
# Add valid disks
json_data = get_json_from_command(cmd)
for disk in json_data.get('blockdevices', []):
disk_obj = Disk(disk['name'])
# Skip loopback devices, optical devices, etc
if disk_obj.details['type'] != 'disk':
continue
# Add disk
disks.append(disk_obj)
# Done
return disks
def get_disks_macos():
"""Get disks via diskutil, returns list."""
cmd = ['diskutil', 'list', '-plist', 'physical']
disks = []
# Get info from diskutil
proc = run_program(cmd, encoding=None, errors=None)
try:
plist_data = plistlib.loads(proc.stdout)
except (TypeError, ValueError):
# Invalid / corrupt plist data? return empty list to avoid crash
LOG.error('Failed to get diskutil list')
return disks
# Add valid disks
for disk in plist_data['WholeDisks']:
disks.append(Disk(f'/dev/{disk}'))
# Done
return disks
def get_known_disk_attributes(model):
"""Get known NVMe/SMART attributes (model specific), returns str."""
known_attributes = KNOWN_DISK_ATTRIBUTES.copy()
# Apply model-specific data
for regex, data in KNOWN_DISK_MODELS.items():
if re.search(regex, model):
for attr, thresholds in data.items():
if attr in known_attributes:
known_attributes[attr].update(thresholds)
else:
known_attributes[attr] = thresholds
# Done
return known_attributes
def get_ram_list_linux():
"""Get RAM list using dmidecode."""
cmd = ['sudo', 'dmidecode', '--type', 'memory']
dimm_list = []
manufacturer = 'Unknown'
size = 0
# Get DMI data
proc = run_program(cmd)
dmi_data = proc.stdout.splitlines()
# Parse data
for line in dmi_data:
line = line.strip()
if line == 'Memory Device':
# Reset vars
manufacturer = 'Unknown'
size = 0
elif line.startswith('Size:'):
size = line.replace('Size: ', '')
try:
size = string_to_bytes(size, assume_binary=True)
except ValueError:
# Assuming empty module
size = 0
elif line.startswith('Manufacturer:'):
manufacturer = line.replace('Manufacturer: ', '')
dimm_list.append([size, manufacturer])
# Save details
return dimm_list
def get_ram_list_macos():
"""Get RAM list using system_profiler."""
dimm_list = []
# Get and parse plist data
cmd = [
'system_profiler',
'-xml',
'SPMemoryDataType',
]
proc = run_program(cmd, check=False, encoding=None, errors=None)
try:
plist_data = plistlib.loads(proc.stdout)
except (TypeError, ValueError):
# Ignore and return an empty list
return dimm_list
# Check DIMM data
dimm_details = plist_data[0].get('_items', [{}])[0].get('_items', [])
for dimm in dimm_details:
manufacturer = dimm.get('dimm_manufacturer', None)
manufacturer = KNOWN_RAM_VENDOR_IDS.get(
manufacturer,
f'Unknown ({manufacturer})')
size = dimm.get('dimm_size', '0 GB')
try:
size = string_to_bytes(size, assume_binary=True)
except ValueError:
# Empty DIMM?
LOG.error('Invalid DIMM size: %s', size)
continue
dimm_list.append([size, manufacturer])
# Save details
return dimm_list
def is_4k_aligned_macos(disk_details):
"""Check partition alignment using diskutil info, returns bool."""
aligned = True
# Check partitions
for part in disk_details.get('children', []):
offset = part.get('PartitionMapPartitionOffset', 0)
if not offset:
# Assuming offset couldn't be found and it defaulted to 0
# NOTE: Just logging the error, not bailing
LOG.error('Failed to get partition offset for %s', part['path'])
aligned = aligned and offset >= 0 and offset % 4096 == 0
# Done
return aligned
def is_4k_aligned_linux(dev_path, physical_sector_size):
"""Check partition alignment using lsblk, returns bool."""
aligned = True
cmd = [
'sudo',
'sfdisk',
'--json',
dev_path,
]
# Get partition details
json_data = get_json_from_command(cmd)
# Check partitions
for part in json_data.get('partitiontable', {}).get('partitions', []):
offset = physical_sector_size * part.get('start', -1)
aligned = aligned and offset >= 0 and offset % 4096 == 0
# Done
return aligned
if __name__ == '__main__':
print("This file is not meant to be called directly.")

412
scripts/wk/hw/sensors.py Normal file
View file

@ -0,0 +1,412 @@
"""WizardKit: Hardware sensors"""
# vim: sts=2 sw=2 ts=2
import json
import logging
import pathlib
import re
from subprocess import CalledProcessError
from wk.cfg.hw import CPU_CRITICAL_TEMP, SMC_IDS, TEMP_COLORS
from wk.exe import run_program, start_thread
from wk.std import PLATFORM, color_string, sleep
# STATIC VARIABLES
LOG = logging.getLogger(__name__)
LM_SENSORS_CPU_REGEX = re.compile(r'(core|k\d+)temp', re.IGNORECASE)
SMC_REGEX = re.compile(
r'^\s*(?P<ID>\w{4})'
r'\s+\[(?P<Type>.*)\]'
r'\s+(?P<Value>.*?)'
r'\s*\(bytes (?P<Bytes>.*)\)$'
)
SENSOR_SOURCE_WIDTH = 25 if PLATFORM == 'Darwin' else 20
# Error Classes
class ThermalLimitReachedError(RuntimeError):
"""Raised when the thermal threshold is reached."""
# Classes
class Sensors():
"""Class for holding sensor specific data."""
def __init__(self):
self.background_thread = None
self.data = get_sensor_data()
self.out_path = None
def clear_temps(self):
"""Clear saved temps but keep structure"""
for adapters in self.data.values():
for sources in adapters.values():
for source_data in sources.values():
source_data['Temps'] = []
def cpu_max_temp(self):
"""Get max temp from any CPU source, returns float.
NOTE: If no temps are found this returns zero.
"""
max_temp = 0.0
# Check all CPU Temps
for section, adapters in self.data.items():
if not section.startswith('CPU'):
continue
for sources in adapters.values():
for source_data in sources.values():
max_temp = max(max_temp, source_data.get('Max', 0))
# Done
return max_temp
def cpu_reached_critical_temp(self):
"""Check if CPU reached CPU_CRITICAL_TEMP, returns bool."""
for section, adapters in self.data.items():
if not section.startswith('CPU'):
# Limit to CPU temps
continue
# Ugly section
for sources in adapters.values():
for source_data in sources.values():
if source_data.get('Max', -1) >= CPU_CRITICAL_TEMP:
return True
# Didn't return above so temps are within the threshold
return False
def generate_report(self, *temp_labels, colored=True, only_cpu=False):
"""Generate report based on given temp_labels, returns list."""
report = []
for section, adapters in sorted(self.data.items()):
if only_cpu and not section.startswith('CPU'):
continue
# Ugly section
for adapter, sources in sorted(adapters.items()):
report.append(fix_sensor_name(adapter))
for source, source_data in sorted(sources.items()):
line = f'{fix_sensor_name(source):{SENSOR_SOURCE_WIDTH}} '
for label in temp_labels:
if label != 'Current':
line += f' {label.lower()}: '
line += get_temp_str(
source_data.get(label, '???'),
colored=colored,
)
report.append(line)
if not only_cpu:
report.append('')
# Handle empty reports
if not report:
report = [
color_string('WARNING: No sensors found', 'YELLOW'),
'',
'Please monitor temps manually',
]
# Done
return report
def monitor_to_file(
self, out_path,
exit_on_thermal_limit=True, temp_labels=None, thermal_action=None):
"""Write report to path every second until stopped.
thermal_action is a cmd to run if ThermalLimitReachedError is caught.
"""
stop_path = pathlib.Path(out_path).resolve().with_suffix('.stop')
if not temp_labels:
temp_labels = ('Current', 'Max')
# Start loop
while True:
try:
self.update_sensor_data(exit_on_thermal_limit)
except ThermalLimitReachedError:
if thermal_action:
run_program(thermal_action, check=False)
report = self.generate_report(*temp_labels)
with open(out_path, 'w') as _f:
_f.write('\n'.join(report))
# Check if we should stop
if stop_path.exists():
break
# Sleep before next loop
sleep(0.5)
def save_average_temps(self, temp_label, seconds=10):
# pylint: disable=unused-variable
"""Save average temps under temp_label over provided seconds.."""
self.clear_temps()
# Get temps
for i in range(seconds):
self.update_sensor_data()
sleep(1)
# Calculate averages
for adapters in self.data.values():
for sources in adapters.values():
for source_data in sources.values():
temps = source_data['Temps']
source_data[temp_label] = sum(temps) / len(temps)
def start_background_monitor(
self, out_path,
exit_on_thermal_limit=True, temp_labels=None, thermal_action=None):
"""Start background thread to save report to file.
thermal_action is a cmd to run if ThermalLimitReachedError is caught.
"""
if self.background_thread:
raise RuntimeError('Background thread already running')
self.out_path = pathlib.Path(out_path)
self.background_thread = start_thread(
self.monitor_to_file,
args=(out_path, exit_on_thermal_limit, temp_labels, thermal_action),
)
def stop_background_monitor(self):
"""Stop background thread."""
self.out_path.with_suffix('.stop').touch()
self.background_thread.join()
# Reset vars to None
self.background_thread = None
self.out_path = None
def update_sensor_data(self, exit_on_thermal_limit=True):
"""Update sensor data via OS-specific means."""
if PLATFORM == 'Darwin':
self.update_sensor_data_macos(exit_on_thermal_limit)
elif PLATFORM == 'Linux':
self.update_sensor_data_linux(exit_on_thermal_limit)
def update_sensor_data_linux(self, exit_on_thermal_limit=True):
"""Update sensor data via lm_sensors."""
lm_sensor_data = get_sensor_data_lm()
for section, adapters in self.data.items():
for adapter, sources in adapters.items():
for source, source_data in sources.items():
try:
label = source_data['Label']
temp = lm_sensor_data[adapter][source][label]
source_data['Current'] = temp
source_data['Max'] = max(temp, source_data['Max'])
source_data['Temps'].append(temp)
except KeyError:
# Dumb workaround for Dell sensors with changing source names
pass
# Raise exception if thermal limit reached
if exit_on_thermal_limit and section == 'CPUTemps':
if source_data['Current'] >= CPU_CRITICAL_TEMP:
raise ThermalLimitReachedError('CPU temps reached limit')
def update_sensor_data_macos(self, exit_on_thermal_limit=True):
"""Update sensor data via SMC."""
for section, adapters in self.data.items():
for sources in adapters.values():
for source_data in sources.values():
cmd = ['smc', '-k', source_data['Label'], '-r']
proc = run_program(cmd)
match = SMC_REGEX.match(proc.stdout.strip())
try:
temp = float(match.group('Value'))
except (TypeError, ValueError):
LOG.error('Failed to update temp %s', source_data['Label'])
continue
# Update source
source_data['Current'] = temp
source_data['Max'] = max(temp, source_data['Max'])
source_data['Temps'].append(temp)
# Raise exception if thermal limit reached
if exit_on_thermal_limit and section == 'CPUTemps':
if source_data['Current'] >= CPU_CRITICAL_TEMP:
raise ThermalLimitReachedError('CPU temps reached limit')
# Functions
def fix_sensor_name(name):
"""Cleanup sensor name, returns str."""
name = re.sub(r'^(\w+)-(\w+)-(\w+)', r'\1 (\2 \3)', name, re.IGNORECASE)
name = name.title()
name = name.replace('Acpi', 'ACPI')
name = name.replace('ACPItz', 'ACPI TZ')
name = name.replace('Coretemp', 'CoreTemp')
name = name.replace('Cpu', 'CPU')
name = name.replace('Id ', 'ID ')
name = name.replace('Isa ', 'ISA ')
name = name.replace('Pci ', 'PCI ')
name = name.replace('Smc', 'SMC')
name = re.sub(r'(\D+)(\d+)', r'\1 \2', name, re.IGNORECASE)
name = re.sub(r'^K (\d+)Temp', r'AMD K\1 Temps', name, re.IGNORECASE)
name = re.sub(r'T(ctl|die)', r'CPU (T\1)', name, re.IGNORECASE)
name = re.sub(r'\s+', ' ', name)
return name
def get_sensor_data():
"""Get sensor data via OS-specific means, returns dict."""
sensor_data = {}
if PLATFORM == 'Darwin':
sensor_data = get_sensor_data_macos()
elif PLATFORM == 'Linux':
sensor_data = get_sensor_data_linux()
return sensor_data
def get_sensor_data_linux():
"""Get sensor data via lm_sensors, returns dict."""
raw_lm_sensor_data = get_sensor_data_lm()
sensor_data = {'CPUTemps': {}, 'Others': {}}
# Parse lm_sensor data
for adapter, sources in raw_lm_sensor_data.items():
section = 'Others'
if LM_SENSORS_CPU_REGEX.search(adapter):
section = 'CPUTemps'
sensor_data[section][adapter] = {}
sources.pop('Adapter', None)
# Find current temp and add to dict
## current temp is labeled xxxx_input
for source, labels in sources.items():
for label, temp in labels.items():
if label.startswith('fan') or label.startswith('in'):
# Skip fan RPMs and voltages
continue
if 'input' in label:
sensor_data[section][adapter][source] = {
'Current': temp,
'Label': label,
'Max': temp,
'Temps': [temp],
}
# Remove empty adapters
if not sensor_data[section][adapter]:
sensor_data[section].pop(adapter)
# Remove empty sections
for adapters in sensor_data.values():
adapters = {source: source_data for source, source_data in adapters.items()
if source_data}
# Done
return sensor_data
def get_sensor_data_lm():
"""Get raw sensor data via lm_sensors, returns dict."""
raw_lm_sensor_data = {}
cmd = ['sensors', '-j']
# Get raw data
try:
proc = run_program(cmd)
except CalledProcessError:
# Assuming no sensors available, return empty dict
return {}
# Workaround for bad sensors
raw_data = []
for line in proc.stdout.splitlines():
if line.strip() == ',':
# Assuming malformatted line caused by missing data
continue
raw_data.append(line)
# Parse JSON data
try:
raw_lm_sensor_data = json.loads('\n'.join(raw_data))
except json.JSONDecodeError:
# Still broken, just return the empty dict
pass
# Done
return raw_lm_sensor_data
def get_sensor_data_macos():
"""Get sensor data via SMC, returns dict.
NOTE: The data is structured like the lm_sensor data.
"""
cmd = ['smc', '-l']
sensor_data = {'CPUTemps': {'SMC (CPU)': {}}, 'Others': {'SMC (Other)': {}}}
# Parse SMC data
proc = run_program(cmd)
for line in proc.stdout.splitlines():
tmp = SMC_REGEX.match(line.strip())
if tmp:
value = tmp.group('Value')
try:
LOG.debug('Invalid sensor: %s', tmp.group('ID'))
value = float(value)
except (TypeError, ValueError):
# Skip this sensor
continue
# Only add known sensor IDs
sensor_id = tmp.group('ID')
if sensor_id not in SMC_IDS:
continue
# Add to dict
section = 'Others'
adapter = 'SMC (Other)'
if SMC_IDS[sensor_id].get('CPU Temp', False):
section = 'CPUTemps'
adapter = 'SMC (CPU)'
source = SMC_IDS[sensor_id]['Source']
sensor_data[section][adapter][source] = {
'Current': value,
'Label': sensor_id,
'Max': value,
'Temps': [value],
}
# Done
return sensor_data
def get_temp_str(temp, colored=True):
"""Get colored string based on temp, returns str."""
temp_color = None
# Safety check
try:
temp = float(temp)
except (TypeError, ValueError):
# Invalid temp?
return color_string(temp, 'PURPLE')
# Determine color
if colored:
for threshold, color in sorted(TEMP_COLORS.items(), reverse=True):
if temp >= threshold:
temp_color = color
break
# Done
return color_string(f'{"-" if temp < 0 else ""}{temp:2.0f}°C', temp_color)
if __name__ == '__main__':
print("This file is not meant to be called directly.")

View file

@ -0,0 +1,196 @@
"""WizardKit: I/O Functions"""
# vim: sts=2 sw=2 ts=2
import logging
import os
import pathlib
import re
import shutil
# STATIC VARIABLES
LOG = logging.getLogger(__name__)
# Functions
def case_insensitive_path(path):
"""Find path case-insensitively, returns pathlib.Path obj."""
given_path = pathlib.Path(path).resolve()
real_path = None
# Quick check
if given_path.exists():
return given_path
# Search for real path
parts = list(given_path.parts)
real_path = parts.pop(0)
for part in parts:
try:
real_path = case_insensitive_search(real_path, part)
except NotADirectoryError:
# Reclassify error
raise FileNotFoundError(given_path)
real_path = pathlib.Path(real_path)
# Done
return real_path
def case_insensitive_search(path, item):
"""Search path for item case insensitively, returns pathlib.Path obj."""
path = pathlib.Path(path).resolve()
given_path = path.joinpath(item)
real_path = None
regex = fr'^{item}'
# Quick check
if given_path.exists():
return given_path
# Check all items in path
for entry in os.scandir(path):
if re.match(regex, entry.name, re.IGNORECASE):
real_path = path.joinpath(entry.name)
# Raise exception if necessary
if not real_path:
raise FileNotFoundError(given_path)
# Done
return real_path
def delete_empty_folders(path):
"""Recursively delete all empty folders in path."""
LOG.debug('path: %s', path)
# Delete empty subfolders first
for item in os.scandir(path):
if item.is_dir():
delete_empty_folders(item.path)
# Attempt to remove (top) path
try:
delete_folder(path, force=False)
except OSError:
# Assuming it's not empty
pass
def delete_folder(path, force=False, ignore_errors=False):
"""Delete folder if empty or if forced.
NOTE: Exceptions are not caught by this function,
ignore_errors is passed to shutil.rmtree to allow partial deletions.
"""
LOG.debug(
'path: %s, force: %s, ignore_errors: %s',
path, force, ignore_errors,
)
if force:
shutil.rmtree(path, ignore_errors=ignore_errors)
else:
os.rmdir(path)
def delete_item(path, force=False, ignore_errors=False):
"""Delete file or folder, optionally recursively.
NOTE: Exceptions are not caught by this function,
ignore_errors is passed to delete_folder to allow partial deletions.
"""
LOG.debug(
'path: %s, force: %s, ignore_errors: %s',
path, force, ignore_errors,
)
path = pathlib.Path(path)
if path.is_dir():
delete_folder(path, force=force, ignore_errors=ignore_errors)
else:
os.remove(path)
def non_clobber_path(path):
"""Update path as needed to non-existing path, returns pathlib.Path."""
LOG.debug('path: %s', path)
path = pathlib.Path(path)
name = path.name
new_path = None
suffix = ''.join(path.suffixes)
name = name.replace(suffix, '')
# Bail early
if not path.exists():
return path
# Find non-existant path
for _i in range(1000):
test_path = path.with_name(f'{name}_{_i}').with_suffix(suffix)
if not test_path.exists():
new_path = test_path
break
# Raise error if viable path not found
if not new_path:
raise FileExistsError(new_path)
# Done
LOG.debug('new path: %s', new_path)
return new_path
def recursive_copy(source, dest, overwrite=False):
"""Copy source to dest recursively.
NOTE: This uses rsync style source/dest syntax.
If the source has a trailing slash then it's contents are copied,
otherwise the source itself is copied.
Examples assuming "ExDir/ExFile.txt" exists:
recursive_copy("ExDir", "Dest/") results in "Dest/ExDir/ExFile.txt"
recursive_copy("ExDir/", "Dest/") results in "Dest/ExFile.txt"
NOTE 2: dest does not use find_path because it might not exist.
"""
copy_contents = str(source).endswith(('/', '\\'))
source = case_insensitive_path(source)
dest = pathlib.Path(dest).resolve().joinpath(source.name)
os.makedirs(dest.parent, exist_ok=True)
# Recursively copy source to dest
if source.is_dir():
if copy_contents:
# Trailing slash syntax
for item in os.scandir(source):
recursive_copy(item.path, dest.parent, overwrite=overwrite)
elif not dest.exists():
# No conflict, copying whole tree (no merging needed)
shutil.copytree(source, dest)
elif not dest.is_dir():
# Refusing to replace file with dir
raise FileExistsError(f'Refusing to replace file: {dest}')
else:
# Dest exists and is a dir, merge dirs
for item in os.scandir(source):
recursive_copy(item.path, dest, overwrite=overwrite)
elif source.is_file():
if not dest.exists():
# No conflict, copying file
shutil.copy2(source, dest)
elif not dest.is_file():
# Refusing to replace dir with file
raise FileExistsError(f'Refusing to replace dir: {dest}')
elif overwrite:
# Dest file exists, deleting and replacing file
os.remove(dest)
shutil.copy2(source, dest)
else:
# Refusing to delete file when overwrite=False
raise FileExistsError(f'Refusing to delete file: {dest}')
if __name__ == '__main__':
print("This file is not meant to be called directly.")

View file

@ -0,0 +1,7 @@
"""WizardKit: kit module init"""
# vim: sts=2 sw=2 ts=2
import platform
if platform.system() == 'Linux':
from wk.kit import ufd

472
scripts/wk/kit/ufd.py Normal file
View file

@ -0,0 +1,472 @@
"""WizardKit: UFD Functions"""
# vim: sts=2 sw=2 ts=2
# TODO: Replace some lsblk usage with hw_obj?
# TODO: Reduce imports if possible
# TODO: Needs testing
import logging
import os
import shutil
from collections import OrderedDict
from docopt import docopt
from wk import io, log, std
from wk.cfg.main import KIT_NAME_FULL, KIT_NAME_SHORT
from wk.cfg.ufd import BOOT_ENTRIES, BOOT_FILES, ITEMS, ITEMS_HIDDEN, SOURCES
from wk.exe import run_program
from wk.os import linux
# STATIC VARIABLES
DOCSTRING = '''WizardKit: Build UFD
Usage:
build-ufd [options] --ufd-device PATH --linux PATH
[--linux-minimal PATH]
[--main-kit PATH]
[--winpe PATH]
[--eset PATH]
[--hdclone PATH]
[--extra-dir PATH]
build-ufd (-h | --help)
Options:
-c PATH, --hdclone PATH
-e PATH, --extra-dir PATH
-k PATH, --main-kit PATH
-l PATH, --linux PATH
-m PATH, --linux-minimal PATH
-s PATH, --eset PATH
-u PATH, --ufd-device PATH
-w PATH, --winpe PATH
-h --help Show this page
-M --use-mbr Use real MBR instead of GPT w/ Protective MBR
-F --force Bypass all confirmation messages. USE WITH EXTREME CAUTION!
-U --update Don't format device, just update
'''
LOG = logging.getLogger(__name__)
ISO_LABEL = f'{KIT_NAME_SHORT}_LINUX'
UFD_LABEL = f'{KIT_NAME_SHORT}_UFD'
# Functions
def build_ufd():
"""Build UFD using selected sources."""
args = docopt(DOCSTRING)
log.update_log_path(dest_name='build-ufd', timestamp=True)
try_print = std.TryAndPrint()
try_print.indent = 2
# Check if running with root permissions
if not linux.running_as_root():
std.print_error('This script is meant to be run as root')
std.abort()
# Show header
std.print_success(KIT_NAME_FULL)
std.print_warning('UFD Build Tool')
std.print_warning(' ')
# Verify selections
ufd_dev = verify_ufd(args['--ufd-device'])
sources = verify_sources(args, SOURCES)
show_selections(args, sources, ufd_dev, SOURCES)
if not args['--force']:
confirm_selections(update=args['--update'])
# Prep UFD
if not args['--update']:
std.print_info('Prep UFD')
prep_device(ufd_dev, UFD_LABEL, use_mbr=args['--use-mbr'])
# Mount UFD
try_print.run(
message='Mounting UFD...',
function=linux.mount,
mount_source=find_first_partition(ufd_dev),
mount_point='/mnt/UFD',
read_write=True,
)
# Remove Arch folder
if args['--update']:
try_print.run(
message='Removing Linux...',
function=remove_arch,
)
# Copy sources
std.print_standard(' ')
std.print_info('Copy Sources')
for s_label, s_path in sources.items():
try_print.run(
message='Copying {}...'.format(s_label),
function=copy_source,
source=s_path,
items=ITEMS[s_label],
overwrite=True,
)
# Update boot entries
std.print_standard(' ')
std.print_info('Boot Setup')
try_print.run(
message='Updating boot entries...',
function=update_boot_entries,
)
# Install syslinux (to partition)
try_print.run(
message='Syslinux (partition)...',
function=install_syslinux_to_partition,
partition=find_first_partition(ufd_dev),
)
# Unmount UFD
try_print.run(
message='Unmounting UFD...',
function=linux.unmount,
mount_point='/mnt/UFD',
)
# Install syslinux (to device)
try_print.run(
message='Syslinux (device)...',
function=install_syslinux_to_dev,
ufd_dev=ufd_dev,
use_mbr=args['--use-mbr'],
)
# Hide items
std.print_standard(' ')
std.print_info('Final Touches')
try_print.run(
message='Hiding items...',
function=hide_items,
ufd_dev=ufd_dev,
items=ITEMS_HIDDEN,
)
# Done
std.print_standard('\nDone.')
if not args['--force']:
std.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()
# Safety check
if not update:
std.print_standard(' ')
std.print_warning('SAFETY CHECK')
std.print_standard(
'All data will be DELETED from the disk and partition(s) listed above.')
std.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()
std.print_standard(' ')
def copy_source(source, items, overwrite=False):
"""Copy source items to /mnt/UFD."""
is_image = source.is_file()
# Mount source if necessary
if is_image:
linux.mount(source, '/mnt/Source')
# Copy items
for i_source, i_dest in items:
i_source = f'{"/mnt/Source" if is_image else source}{i_source}'
i_dest = f'/mnt/UFD{i_dest}'
try:
io.recursive_copy(i_source, i_dest, overwrite=overwrite)
except FileNotFoundError:
# Going to assume (hope) that this is fine
pass
# Unmount source if necessary
if is_image:
linux.unmount('/mnt/Source')
def find_first_partition(dev_path):
"""Find path to first partition of dev, returns str.
NOTE: This assumes the dev was just partitioned with
a single partition.
"""
cmd = [
'lsblk',
'--list',
'--noheadings',
'--output', 'name',
'--paths',
dev_path,
]
# Run cmd
proc = run_program(cmd)
part_path = proc.stdout.splitlines()[-1].strip()
# Done
return part_path
def hide_items(ufd_dev, items):
"""Set FAT32 hidden flag for items."""
first_partition = find_first_partition(ufd_dev)
with open('/root/.mtoolsrc', 'w') as _f:
_f.write(f'drive U: file="{first_partition}"\n')
_f.write('mtools_skip_check=1\n')
# Hide items
for item in items:
cmd = [f'yes | mattrib +h "U:/{item}"']
run_program(cmd, check=False, shell=True)
def install_syslinux_to_dev(ufd_dev, use_mbr):
"""Install Syslinux to UFD (dev)."""
cmd = [
'dd',
'bs=440',
'count=1',
f'if=/usr/lib/syslinux/bios/{"mbr" if use_mbr else "gptmbr"}.bin',
f'of={ufd_dev}',
]
run_program(cmd)
def install_syslinux_to_partition(partition):
"""Install Syslinux to UFD (partition)."""
cmd = [
'syslinux',
'--install',
'--directory',
'/arch/boot/syslinux/',
partition,
]
run_program(cmd)
def is_valid_path(path_obj, path_type):
"""Verify path_obj is valid by type, returns bool."""
valid_path = False
if path_type == 'DIR':
valid_path = path_obj.is_dir()
elif path_type == 'KIT':
valid_path = path_obj.is_dir() and path_obj.joinpath('.bin').exists()
elif path_type == 'IMG':
valid_path = path_obj.is_file() and path_obj.suffix.lower() == '.img'
elif path_type == 'ISO':
valid_path = path_obj.is_file() and path_obj.suffix.lower() == '.iso'
elif path_type == 'UFD':
valid_path = path_obj.is_block_device()
return valid_path
def prep_device(dev_path, label, use_mbr=False):
"""Format device in preparation for applying the WizardKit components
This is done is four steps:
1. Zero-out first 64MB (this deletes the partition table and/or bootloader)
2. Create a new partition table (GPT by default, optionally MBR)
3. Set boot flag
4. Format partition (FAT32, 4K aligned)
"""
try_print = std.TryAndPrint()
try_print.indent = 2
# Zero-out first 64MB
cmd = [
'dd',
'bs=4M',
'count=16',
'if=/dev/zero',
f'of={dev_path}',
]
try_print.run(
message='Zeroing first 64MiB...',
function=run_program,
cmd=cmd,
)
# Create partition table
cmd = [
'parted', dev_path,
'--script',
'--',
'mklabel', 'msdos' if use_mbr else 'gpt',
'-1s' if use_mbr else '-4MiB',
]
try_print.run(
message='Creating partition table...',
function=run_program,
cmd=cmd,
)
# Set boot flag
cmd = [
'parted', dev_path,
'set', '1',
'boot' if use_mbr else 'legacy_boot',
'on',
]
try_print.run(
message='Setting boot flag...',
function=run_program,
cmd=cmd,
)
# Format partition
cmd = [
'mkfs.vfat',
'-F', '32',
'-n', label,
find_first_partition(dev_path),
]
try_print.run(
message='Formatting partition...',
function=run_program,
cmd=cmd,
)
def remove_arch():
"""Remove arch dir from UFD.
This ensures a clean installation to the UFD and resets the boot files
"""
shutil.rmtree(io.case_insensitive_path('/mnt/UFD/arch'))
def show_selections(args, sources, ufd_dev, ufd_sources):
"""Show selections including non-specified options."""
# Sources
std.print_info('Sources')
for label in ufd_sources.keys():
if label in sources:
std.print_standard(f' {label+":":<18} {sources["label"]}')
else:
std.print_colored(
[f' {label+":":<18}', 'Not Specified'],
[None, 'YELLOW'],
)
std.print_standard(' ')
# Destination
std.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())
cmd = [
'lsblk', '--noheadings', '--paths',
'--output', 'NAME,SIZE,FSTYPE,LABEL,MOUNTPOINT',
ufd_dev,
]
proc = run_program(cmd, check=False)
for line in proc.stdout.splitlines()[1:]:
std.print_standard(line)
# Notes
if args['--update']:
std.print_warning('Updating kit in-place')
elif args['--use-mbr']:
std.print_warning('Formatting using legacy MBR')
std.print_standard(' ')
def update_boot_entries():
"""Update boot files for UFD usage"""
configs = []
# Find config files
for c_path, c_ext in BOOT_FILES.items():
c_path = io.case_insensitive_path('/mnt/UFD{c_path}')
for item in os.scandir(c_path):
if item.name.lower().endswith(c_ext.lower()):
configs.append(item.path)
# Update Linux labels
cmd = [
'sed',
'--in-place',
'--regexp-extended',
f's/(eSysRescueLiveCD|{ISO_LABEL})/{UFD_LABEL}/',
*configs,
]
run_program(cmd)
# Uncomment extra entries if present
for b_path, b_comment in BOOT_ENTRIES.items():
try:
io.case_insensitive_path(f'/mnt/UFD{b_path}')
except (FileNotFoundError, NotADirectoryError):
# Entry not found, continue to next entry
continue
# Entry found, update config files
cmd = [
'sed',
'--in-place',
f's/#{b_comment}#//',
*configs,
]
run_program(cmd, check=False)
def verify_sources(args, ufd_sources):
"""Check all sources and abort if necessary, returns dict."""
sources = OrderedDict()
for label, data in ufd_sources.items():
s_path = args[data['Arg']]
if s_path:
try:
s_path_obj = io.case_insensitive_path(s_path)
except FileNotFoundError:
std.print_error(f'ERROR: {label} not found: {s_path}')
std.abort()
if not is_valid_path(s_path_obj, data['Type']):
std.print_error(f'ERROR: Invalid {label} source: {s_path}')
std.abort()
sources[label] = s_path_obj
return sources
def verify_ufd(dev_path):
"""Check that dev_path is a valid UFD, returns pathlib.Path obj."""
ufd_dev = None
try:
ufd_dev = io.case_insensitive_path(dev_path)
except FileNotFoundError:
std.print_error(f'ERROR: UFD device not found: {dev_path}')
std.abort()
if not is_valid_path(ufd_dev, 'UFD'):
std.print_error(f'ERROR: Invalid UFD device: {ufd_dev}')
std.abort()
return ufd_dev
if __name__ == '__main__':
print("This file is not meant to be called directly.")

154
scripts/wk/log.py Normal file
View file

@ -0,0 +1,154 @@
"""WizardKit: Log Functions"""
# vim: sts=2 sw=2 ts=2
import atexit
import logging
import os
import pathlib
import shutil
import time
from wk import cfg
from wk.io import non_clobber_path
# STATIC VARIABLES
if os.name == 'nt':
# Example: "C:\WK\1955-11-05\WizardKit"
DEFAULT_LOG_DIR = (
f'{os.environ.get("SYSTEMDRIVE", "C:")}/'
f'{cfg.main.KIT_NAME_SHORT}/'
f'{time.strftime("%Y-%m-%d")}'
)
else:
# Example: "/home/tech/Logs"
DEFAULT_LOG_DIR = f'{os.path.expanduser("~")}/Logs'
DEFAULT_LOG_NAME = cfg.main.KIT_NAME_FULL
# Functions
def enable_debug_mode():
"""Configures logging for better debugging."""
root_logger = logging.getLogger()
for handler in root_logger.handlers:
formatter = logging.Formatter(
datefmt=cfg.log.DEBUG['datefmt'],
fmt=cfg.log.DEBUG['format'],
)
handler.setFormatter(formatter)
root_logger.setLevel('DEBUG')
def format_log_path(
log_dir=None, log_name=None, timestamp=False,
kit=False, tool=False):
"""Format path based on args passed, returns pathlib.Path obj."""
log_path = pathlib.Path(
f'{log_dir if log_dir else DEFAULT_LOG_DIR}/'
f'{cfg.main.KIT_NAME_FULL+"/" if kit else ""}'
f'{"Tools/" if tool else ""}'
f'{log_name if log_name else DEFAULT_LOG_NAME}'
f'{"_" if timestamp else ""}'
f'{time.strftime("%Y-%m-%d_%H%M%S%z") if timestamp else ""}'
'.log'
)
log_path = log_path.resolve()
# Avoid clobbering
log_path = non_clobber_path(log_path)
# Done
return log_path
def get_root_logger_path():
"""Get path to log file from root logger, returns pathlib.Path obj."""
log_path = None
root_logger = logging.getLogger()
# Check all handlers and use the first fileHandler found
for handler in root_logger.handlers:
if isinstance(handler, logging.FileHandler):
log_path = pathlib.Path(handler.baseFilename).resolve()
break
# Done
return log_path
def remove_empty_log():
"""Remove log if empty."""
is_empty = False
# Check if log is empty
log_path = get_root_logger_path()
try:
is_empty = log_path and log_path.exists() and log_path.stat().st_size == 0
except (FileNotFoundError, AttributeError):
# File doesn't exist or couldn't verify it's empty
pass
# Delete log
if is_empty:
log_path.unlink()
def start(config=None):
"""Configure and start logging using safe defaults."""
log_path = format_log_path(timestamp=os.name != 'nt')
root_logger = logging.getLogger()
# Safety checks
if not config:
config = cfg.log.DEFAULT
if root_logger.hasHandlers():
raise UserWarning('Logging already started, results may be unpredictable.')
# Create log_dir
os.makedirs(log_path.parent, exist_ok=True)
# Config logger
logging.basicConfig(filename=log_path, **config)
# Register shutdown to run atexit
atexit.register(remove_empty_log)
atexit.register(logging.shutdown)
def update_log_path(
dest_dir=None, dest_name=None, keep_history=True, timestamp=True):
"""Moves current log file to new path and updates the root logger."""
root_logger = logging.getLogger()
cur_handler = None
cur_path = get_root_logger_path()
new_path = format_log_path(dest_dir, dest_name, timestamp=timestamp)
os.makedirs(new_path.parent, exist_ok=True)
# Get current logging file handler
for handler in root_logger.handlers:
if isinstance(handler, logging.FileHandler):
cur_handler = handler
break
if not cur_handler:
raise RuntimeError('Logging FileHandler not found')
# Copy original log to new location
if keep_history:
if new_path.exists():
raise FileExistsError(f'Refusing to clobber: {new_path}')
shutil.move(cur_path, new_path)
# Remove old log if empty
remove_empty_log()
# Create new cur_handler (preserving formatter settings)
new_handler = logging.FileHandler(new_path, mode='a')
new_handler.setFormatter(cur_handler.formatter)
# Replace current handler
root_logger.removeHandler(cur_handler)
root_logger.addHandler(new_handler)
if __name__ == '__main__':
print("This file is not meant to be called directly.")

View file

@ -0,0 +1,260 @@
"""WizardKit: Net Functions"""
# vim: sts=2 sw=2 ts=2
import os
import pathlib
import re
import psutil
from wk.exe import get_json_from_command, run_program
from wk.std import PLATFORM, GenericError, show_data
from wk.cfg.net import BACKUP_SERVERS
# REGEX
REGEX_VALID_IP = re.compile(
r'(10.\d+.\d+.\d+'
r'|172.(1[6-9]|2\d|3[0-1])'
r'|192.168.\d+.\d+)',
re.IGNORECASE)
# Functions
def connected_to_private_network(raise_on_error=False):
"""Check if connected to a private network, returns bool.
This checks for a valid private IP assigned to this system.
NOTE: If one isn't found and raise_on_error=True then an exception is raised.
NOTE 2: If one is found and raise_on_error=True then None is returned.
"""
connected = False
# Check IPs
devs = psutil.net_if_addrs()
for dev in devs.values():
for family in dev:
if REGEX_VALID_IP.search(family.address):
# Valid IP found
connected = True
break
if connected:
break
# No valid IP found
if not connected and raise_on_error:
raise GenericError('Not connected to a network')
# Done
if raise_on_error:
connected = None
return connected
def mount_backup_shares(read_write=False):
"""Mount backup shares using OS specific methods."""
report = []
for name, details in BACKUP_SERVERS.items():
mount_point = None
mount_str = f'{name} (//{details["Address"]}/{details["Share"]})'
# Prep mount point
if PLATFORM in ('Darwin', 'Linux'):
mount_point = pathlib.Path(f'/Backups/{name}')
try:
if not mount_point.exists():
# Script should be run as user so sudo is required
run_program(['sudo', 'mkdir', '-p', mount_point])
except OSError:
# Assuming permission denied under macOS
pass
if mount_point:
mount_str += f' to {mount_point}'
# Check if already mounted
if share_is_mounted(details):
report.append(f'(Already) Mounted {mount_str}')
# Skip to next share
continue
# Mount share
proc = mount_network_share(details, mount_point, read_write=read_write)
if proc.returncode:
report.append(f'Failed to Mount {mount_str}')
else:
report.append(f'Mounted {mount_str}')
# Done
return report
def mount_network_share(details, mount_point=None, read_write=False):
"""Mount network share using OS specific methods."""
cmd = None
address = details['Address']
share = details['Share']
username = details['RO-User']
password = details['RO-Pass']
if read_write:
username = details['RW-User']
password = details['RW-Pass']
# Network check
if not connected_to_private_network():
raise RuntimeError('Not connected to a network')
# Build OS-specific command
if PLATFORM == 'Darwin':
cmd = [
'sudo',
'mount',
'-t', 'smbfs',
'-o', f'{"rw" if read_write else "ro"}',
f'//{username}:{password}@{address}/{share}',
mount_point,
]
elif PLATFORM == 'Linux':
cmd = [
'sudo',
'mount',
'-t', 'cifs',
'-o', (
f'{"rw" if read_write else "ro"}'
f',uid={os.getuid()}'
f',gid={os.getgid()}'
f',username={username}'
f',{"password=" if password else "guest"}{password}'
),
f'//{address}/{share}',
mount_point
]
elif PLATFORM == 'Windows':
cmd = ['net', 'use']
if mount_point:
cmd.append(f'{mount_point}:')
cmd.append(f'/user:{username}')
cmd.append(fr'\\{address}\{share}')
cmd.append(password)
# Mount share
return run_program(cmd, check=False)
def ping(addr='google.com'):
"""Attempt to ping addr."""
cmd = (
'ping',
'-n' if psutil.WINDOWS else '-c',
'2',
addr,
)
run_program(cmd)
def share_is_mounted(details):
"""Check if dev/share/etc is mounted, returns bool."""
mounted = False
if PLATFORM == 'Darwin':
# Weak and naive text search
proc = run_program(['mount'], check=False)
for line in proc.stdout.splitlines():
if f'{details["Address"]}/{details["Share"]}' in line:
mounted = True
break
elif PLATFORM == 'Linux':
cmd = [
'findmnt',
'--list',
'--json',
'--invert',
'--types', (
'autofs,binfmt_misc,bpf,cgroup,cgroup2,configfs,debugfs,devpts,'
'devtmpfs,hugetlbfs,mqueue,proc,pstore,securityfs,sysfs,tmpfs'
),
'--output', 'SOURCE',
]
mount_data = get_json_from_command(cmd)
for row in mount_data.get('filesystems', []):
if row['source'] == f'//{details["Address"]}/{details["Share"]}':
mounted = True
break
#TODO: Check mount status under Windows
#elif PLATFORM == 'Windows':
# Done
return mounted
def show_valid_addresses():
"""Show all valid private IP addresses assigned to the system."""
devs = psutil.net_if_addrs()
for dev, families in sorted(devs.items()):
for family in families:
if REGEX_VALID_IP.search(family.address):
# Valid IP found
show_data(message=dev, data=family.address)
def speedtest():
"""Run a network speedtest using speedtest-cli."""
cmd = ['speedtest-cli', '--simple']
proc = run_program(cmd, check=False)
output = [line.strip() for line in proc.stdout.splitlines() if line.strip()]
output = [line.split() for line in output]
output = [(a, float(b), c) for a, b, c in output]
return [f'{a:<10}{b:6.2f} {c}' for a, b, c in output]
def unmount_backup_shares():
"""Unmount backup shares."""
report = []
for name, details in BACKUP_SERVERS.items():
kwargs = {}
source_str = f'{name} (//{details["Address"]}/{details["Share"]})'
# Check if mounted
if not share_is_mounted(details):
report.append(f'Not mounted {source_str}')
continue
# Build OS specific kwargs
if PLATFORM in ('Darwin', 'Linux'):
kwargs['mount_point'] = f'/Backups/{name}'
elif PLATFORM == 'Windows':
kwargs['details'] = details
# Unmount and add to report
proc = unmount_network_share(**kwargs)
if proc.returncode:
report.append(f'Failed to unmount {source_str}')
else:
report.append(f'Unmounted {source_str}')
# Done
return report
def unmount_network_share(details=None, mount_point=None):
"""Unmount network share"""
cmd = []
# Build OS specific command
if PLATFORM in ('Darwin', 'Linux'):
cmd = ['sudo', 'umount', mount_point]
elif PLATFORM == 'Windows':
cmd = ['net', 'use']
if mount_point:
cmd.append(f'{mount_point}:')
elif details:
cmd.append(fr'\\{details["Address"]}\{details["Share"]}')
cmd.append('/delete')
# Unmount share
return run_program(cmd, check=False)
if __name__ == '__main__':
print("This file is not meant to be called directly.")

10
scripts/wk/os/__init__.py Normal file
View file

@ -0,0 +1,10 @@
"""WizardKit: os module init"""
# vim: sts=2 sw=2 ts=2
import platform
#if platform.system() == 'Darwin':
if platform.system() == 'Linux':
from wk.os import linux
if platform.system() == 'Windows':
from wk.os import win

240
scripts/wk/os/linux.py Normal file
View file

@ -0,0 +1,240 @@
"""WizardKit: Linux Functions"""
# vim: sts=2 sw=2 ts=2
import logging
import os
import pathlib
import re
import subprocess
from wk import std
from wk.exe import popen_program, run_program
from wk.hw.obj import Disk
# STATIC VARIABLES
LOG = logging.getLogger(__name__)
UUID_CORESTORAGE = '53746f72-6167-11aa-aa11-00306543ecac'
# Functions
def get_user_home(user):
"""Get path to user's home dir, returns pathlib.Path obj."""
home = None
# Get path from user details
cmd = ['getent', 'passwd', user]
proc = run_program(cmd, check=False)
try:
home = proc.stdout.split(':')[5]
except IndexError:
# Try using environment variable
home = os.environ.get('HOME')
# Raise exception if necessary
if not home:
raise RuntimeError(f'Failed to find home for: {user}')
# Done
return pathlib.Path(home)
def get_user_name():
"""Get real user name, returns str."""
user = None
# Query environment
user = os.environ.get('SUDO_USER')
if not user:
user = os.environ.get('USER')
# Raise exception if necessary
if not user:
raise RuntimeError('Failed to determine user')
# Done
return user
def make_temp_file():
"""Make temporary file, returns pathlib.Path() obj."""
proc = run_program(['mktemp'], check=False)
return pathlib.Path(proc.stdout.strip())
def mount(source, mount_point=None, read_write=False):
"""Mount source (on mount_point if provided).
NOTE: If not running_as_root() then udevil will be used.
"""
cmd = [
'mount',
'-o', 'rw' if read_write else 'ro',
source,
]
if not running_as_root():
cmd.insert(0, 'udevil')
if mount_point:
cmd.append(mount_point)
# Run mount command
proc = run_program(cmd, check=False)
if not proc.returncode == 0:
raise RuntimeError(f'Failed to mount: {source} on {mount_point}')
def mount_volumes(device_path=None, read_write=False, scan_corestorage=False):
"""Mount all detected volumes, returns list.
NOTE: If device_path is specified then only volumes
under that path will be mounted.
"""
report = []
volumes = []
containers = []
# Get list of volumes
cmd = [
'lsblk',
'--list',
'--noheadings',
'--output=name',
'--paths',
]
if device_path:
cmd.append(device_path)
proc = run_program(cmd, check=False)
for line in sorted(proc.stdout.splitlines()):
volumes.append(Disk(line.strip()))
# Get list of CoreStorage containers
containers = [
vol for vol in volumes if vol.details.get('parttype', '') == UUID_CORESTORAGE
]
# Scan CoreStorage containers
if scan_corestorage:
if containers:
std.print_warning(
f'Detected CoreStorage container{"s" if len(containers) > 1 else ""}',
)
std.print_standard('Scanning for inner volume(s)...')
for container in containers:
volumes.extend(scan_corestorage_container(container))
# Mount volumes
for vol in volumes:
already_mounted = vol.details.get('mountpoint', '')
result = f'{vol.details["name"].replace("/dev/mapper/", ""):<20}'
# Parent devices
if vol.details.get('children', False):
if vol.details.get('fstype', ''):
result += vol.details['fstype']
if vol.details.get('label', ''):
result += f' "{vol.details["label"]}"'
report.append(std.color_string(result, 'BLUE'))
continue
# Attempt to mount volume
if not already_mounted:
mount(vol.path, read_write=read_write)
proc = run_program(cmd, check=False)
if proc.returncode:
result += 'Failed to mount'
report.append(std.color_string(result, 'RED'))
continue
# Add size to result
vol.get_details()
vol.details['fsused'] = vol.details.get('fsused', -1)
vol.details['fsavail'] = vol.details.get('fsavail', -1)
result += f'{"Mounted on "+vol.details.get("mountpoint", "?"):<40}'
result = (
f'{result} ({vol.details.get("fstype", "Unknown FS")+",":<5} '
f'{std.bytes_to_string(vol.details["fsused"], decimals=1):>9} used, '
f'{std.bytes_to_string(vol.details["fsavail"], decimals=1):>9} free)'
)
report.append(
std.color_string(
result,
'YELLOW' if already_mounted else None,
),
)
# Done
return report
def running_as_root():
"""Check if running with effective UID of 0, returns bool."""
return os.geteuid() == 0
def scan_corestorage_container(container, timeout=300):
"""Scan CoreStorage container for inner volumes, returns list."""
# TODO: Test Scanning CoreStorage containers
detected_volumes = {}
inner_volumes = []
log_path = make_temp_file()
# Run scan via TestDisk
cmd = [
'sudo', 'testdisk',
'/logname', log_path,
'/debug',
'/log',
'/cmd', container.path, 'partition_none,analyze',
]
proc = popen_program(cmd)
try:
proc.wait(timeout=timeout)
except subprocess.TimeoutExpired:
# Failed to find any volumes, stop scan
run_program(['sudo', 'kill', proc.pid], check=False)
# Check results
if proc.returncode == 0 and log_path.exists():
results = log_path.read_text(encoding='utf-8', errors='ignore')
for line in results.splitlines():
line = line.lower().strip()
match = re.match(r'^.*echo "([^"]+)" . dmsetup create test(\d)$', line)
if match:
cs_name = f'CoreStorage_{container.path.name}_{match.group(2)}'
detected_volumes[cs_name] = match.group(1)
# Create mapper device(s) if necessary
for name, cmd in detected_volumes.items():
cmd_file = make_temp_file()
cmd_file.write_text(cmd)
proc = run_program(
cmd=['sudo', 'dmsetup', 'create', name, cmd_file],
check=False,
)
if proc.returncode == 0:
inner_volumes.append(Disk(f'/dev/mapper/{name}'))
# Done
return inner_volumes
def unmount(source_or_mountpoint):
"""Unmount source_or_mountpoint.
NOTE: If not running_as_root() then udevil will be used.
"""
cmd = [
'umount',
source_or_mountpoint,
]
if not running_as_root():
cmd.insert(0, 'udevil')
# Run unmount command
proc = run_program(cmd, check=False)
if not proc.returncode == 0:
raise RuntimeError(f'Failed to unmount: {source_or_mountpoint}')
if __name__ == '__main__':
print("This file is not meant to be called directly.")

181
scripts/wk/os/win.py Normal file
View file

@ -0,0 +1,181 @@
"""WizardKit: Windows Functions"""
# vim: sts=2 sw=2 ts=2
import logging
import os
import pathlib
import platform
from wk.borrowed import acpi
from wk.exe import run_program
from wk.io import non_clobber_path
from wk.log import format_log_path
from wk.std import GenericError, GenericWarning, sleep
# STATIC VARIABLES
LOG = logging.getLogger(__name__)
OS_VERSION = float(platform.win32_ver()[0]) # TODO: Check if Win8.1 returns '8'
REG_MSISERVER = r'HKLM\SYSTEM\CurrentControlSet\Control\SafeBoot\Network\MSIServer'
SLMGR = pathlib.Path(f'{os.environ.get("SYSTEMROOT")}/System32/slmgr.vbs')
# Functions
def activate_with_bios():
"""Attempt to activate Windows with a key stored in the BIOS."""
# Code borrowed from https://github.com/aeruder/get_win8key
#####################################################
#script to query windows 8.x OEM key from PC firmware
#ACPI -> table MSDM -> raw content -> byte offset 56 to end
#ck, 03-Jan-2014 (christian@korneck.de)
#####################################################
bios_key = None
table = b"MSDM"
if acpi.FindAcpiTable(table) is True:
rawtable = acpi.GetAcpiTable(table)
#http://msdn.microsoft.com/library/windows/hardware/hh673514
#byte offset 36 from beginning
# = Microsoft 'software licensing data structure'
# / 36 + 20 bytes offset from beginning = Win Key
bios_key = rawtable[56:len(rawtable)].decode("utf-8")
if not bios_key:
raise GenericError('BIOS key not found.')
# Check if activation is needed
if is_activated():
raise GenericWarning('System already activated')
# Install Key
cmd = ['cscript', '//nologo', SLMGR, '/ipk', bios_key]
run_program(cmd, check=False)
sleep(5)
# Attempt activation
cmd = ['cscript', '//nologo', SLMGR, '/ato']
run_program(cmd, check=False)
sleep(5)
# Check status
if not is_activated():
raise GenericError('Activation Failed')
def disable_safemode():
"""Edit BCD to remove safeboot value."""
cmd = ['bcdedit', '/deletevalue', '{default}', 'safeboot']
run_program(cmd)
def disable_safemode_msi():
"""Disable MSI access under safemode."""
cmd = ['reg', 'delete', REG_MSISERVER, '/f']
run_program(cmd)
def enable_safemode():
"""Edit BCD to set safeboot as default."""
cmd = ['bcdedit', '/set', '{default}', 'safeboot', 'network']
run_program(cmd)
def enable_safemode_msi():
"""Enable MSI access under safemode."""
cmd = ['reg', 'add', REG_MSISERVER, '/f']
run_program(cmd)
cmd = [
'reg', 'add', REG_MSISERVER, '/ve',
'/t', 'REG_SZ',
'/d', 'Service', '/f',
]
run_program(cmd)
def get_activation_string():
"""Get activation status, returns str."""
cmd = ['cscript', '//nologo', SLMGR, '/xpr']
proc = run_program(cmd, check=False)
act_str = proc.stdout
act_str = act_str.splitlines()[1]
act_str = act_str.strip()
return act_str
def is_activated():
"""Check if Windows is activated via slmgr.vbs and return bool."""
act_str = get_activation_string()
# Check result.
return act_str and 'permanent' in act_str
def run_chkdsk_offline():
"""Set filesystem 'dirty bit' to force a CHKDSK during startup."""
cmd = f'fsutil dirty set {os.environ.get("SYSTEMDRIVE")}'
proc = run_program(cmd.split(), check=False)
# Check result
if proc.returncode > 0:
raise GenericError('Failed to set dirty bit.')
def run_chkdsk_online():
"""Run CHKDSK in a split window.
NOTE: If run on Windows 8+ online repairs are attempted.
"""
cmd = ['CHKDSK', os.environ.get('SYSTEMDRIVE', 'C:')]
if OS_VERSION >= 8:
cmd.extend(['/scan', '/perf'])
log_path = format_log_path(log_name='CHKDSK', tool=True)
err_path = log_path.with_suffix('.err')
# Run scan
proc = run_program(cmd, check=False)
# Check result
if proc.returncode == 1:
raise GenericWarning('Repaired (or manually aborted)')
if proc.returncode > 1:
raise GenericError('Issue(s) detected')
# Save output
os.makedirs(log_path.parent, exist_ok=True)
with open(log_path, 'w') as _f:
_f.write(proc.stdout)
with open(err_path, 'w') as _f:
_f.write(proc.stderr)
def run_sfc_scan():
"""Run SFC and save results."""
cmd = ['sfc', '/scannow']
log_path = format_log_path(log_name='SFC', tool=True)
err_path = log_path.with_suffix('.err')
# Run SFC
proc = run_program(cmd, check=False, encoding='utf-16')
# Fix paths
log_path = non_clobber_path(log_path)
err_path = non_clobber_path(err_path)
# Save output
os.makedirs(log_path.parent, exist_ok=True)
with open(log_path, 'w') as _f:
_f.write(proc.stdout)
with open(err_path, 'w') as _f:
_f.write(proc.stderr)
# Check result
if 'did not find any integrity violations' in proc.stdout:
pass
elif 'successfully repaired' in proc.stdout:
raise GenericWarning('Repaired')
elif 'found corrupt files' in proc.stdout:
raise GenericError('Corruption detected')
else:
raise OSError
if __name__ == '__main__':
print("This file is not meant to be called directly.")

File diff suppressed because it is too large Load diff

View file

@ -0,0 +1 @@
"""WizardKit: sw module init"""

287
scripts/wk/tmux.py Normal file
View file

@ -0,0 +1,287 @@
"""WizardKit: tmux Functions"""
# vim: sts=2 sw=2 ts=2
import logging
import pathlib
from wk.exe import run_program
from wk.std import PLATFORM
# STATIC_VARIABLES
LOG = logging.getLogger(__name__)
# Functions
def capture_pane(pane_id=None):
"""Capture text from current or target pane, returns str."""
cmd = ['tmux', 'capture-pane', '-p']
if pane_id:
cmd.extend(['-t', pane_id])
# Capture and return
proc = run_program(cmd, check=False)
return proc.stdout.strip()
def clear_pane(pane_id=None):
"""Clear pane buffer for current or target pane."""
cmd = ['tmux', 'send-keys', '-R']
if pane_id:
cmd.extend(['-t', pane_id])
# Clear pane
run_program(cmd, check=False)
def fix_layout(panes, layout, forced=False):
"""Fix pane sizes based on layout."""
if not (forced or layout_needs_fixed(panes, layout)):
# Layout should be fine
return
# Update panes
for name, data in layout.items():
# Skip missing panes
if name not in panes:
continue
# 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
def get_pane_size(pane_id=None):
"""Get current or target pane size, returns tuple."""
cmd = ['tmux', 'display', '-p']
if pane_id:
cmd.extend(['-t', pane_id])
cmd.append('#{pane_width} #{pane_height}')
# Get resolution
proc = run_program(cmd, check=False)
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']
if pane_id:
cmd.extend(['-t', pane_id])
# Kill
run_program(cmd, check=False)
def kill_pane(*pane_ids):
"""Kill pane(s) by id."""
cmd = ['tmux', 'kill-pane', '-t']
# Iterate over all passed pane IDs
for pane_id in pane_ids:
run_program(cmd+[pane_id], check=False)
def layout_needs_fixed(panes, 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
# 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
# Done
return needs_fixed
def poll_pane(pane_id):
"""Check if pane exists, returns bool."""
cmd = ['tmux', 'list-panes', '-F', '#D']
# Get list of panes
proc = run_program(cmd, check=False)
existant_panes = proc.stdout.splitlines()
# Check if pane exists
return pane_id in existant_panes
def prep_action(
cmd=None, working_dir=None, text=None, watch_file=None, watch_cmd='cat'):
"""Prep action to perform during a tmux call, returns list.
This will prep for running a basic command, displaying text on screen,
or monitoring a file. The last option uses cat by default but can be
overridden by using the watch_cmd.
"""
action_cmd = []
if working_dir:
action_cmd.extend(['-c', working_dir])
if cmd:
# Basic command
action_cmd.append(cmd)
elif text:
# Display text
echo_cmd = ['echo']
if PLATFORM == 'Linux':
echo_cmd.append('-e')
action_cmd.extend([
'watch',
'--color',
'--exec',
'--no-title',
'--interval', '1',
])
action_cmd.extend(echo_cmd)
action_cmd.append(text)
elif watch_file:
# Monitor file
prep_file(watch_file)
if watch_cmd == 'cat':
action_cmd.extend([
'watch',
'--color',
'--no-title',
'--interval', '1',
'cat',
])
elif watch_cmd == 'tail':
action_cmd.extend(['tail', '-f'])
action_cmd.append(watch_file)
else:
LOG.error('No action specified')
raise RuntimeError('No action specified')
# Done
return action_cmd
def prep_file(path):
"""Check if file exists and create empty file if not."""
path = pathlib.Path(path).resolve()
try:
path.touch(exist_ok=False)
except FileExistsError:
# Leave existing files alone
pass
def resize_pane(pane_id=None, width=None, height=None, **kwargs):
# pylint: disable=unused-argument
"""Resize current or target pane.
NOTE: kwargs is only here to make calling this function easier
by dropping any extra kwargs passed.
"""
cmd = ['tmux', 'resize-pane']
# Safety checks
if not (width or height):
LOG.error('Neither width nor height specified')
raise RuntimeError('Neither width nor height specified')
# Finish building cmd
if pane_id:
cmd.extend(['-t', pane_id])
if width:
cmd.extend(['-x', str(width)])
if height:
cmd.extend(['-y', str(height)])
# Resize
run_program(cmd, check=False)
def split_window(
lines=None, percent=None,
behind=False, vertical=False,
target_id=None, **action):
"""Split tmux window, run action, and return pane_id as str."""
cmd = ['tmux', 'split-window', '-d', '-PF', '#D']
# Safety checks
if not (lines or percent):
LOG.error('Neither lines nor percent specified')
raise RuntimeError('Neither lines nor percent specified')
# New pane placement
if behind:
cmd.append('-b')
if vertical:
cmd.append('-v')
else:
cmd.append('-h')
if target_id:
cmd.extend(['-t', target_id])
# New pane size
if lines:
cmd.extend(['-l', str(lines)])
elif percent:
cmd.extend(['-p', str(percent)])
# New pane action
cmd.extend(prep_action(**action))
# Run and return pane_id
proc = run_program(cmd, check=False)
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']
if pane_id:
cmd.extend(['-t', pane_id])
# Toggle
run_program(cmd, check=False)
if __name__ == '__main__':
print("This file is not meant to be called directly.")

View file

@ -11,10 +11,10 @@ set -o pipefail
DATE="$(date +%F)"
DATETIME="$(date +%F_%H%M)"
ROOT_DIR="$(realpath $(dirname "$0"))"
BUILD_DIR="$ROOT_DIR/BUILD_LINUX"
BUILD_DIR="$ROOT_DIR/setup/BUILD_LINUX"
LIVE_DIR="$BUILD_DIR/live"
LOG_DIR="$BUILD_DIR/logs"
OUT_DIR="$ROOT_DIR/OUT_LINUX"
OUT_DIR="$ROOT_DIR/setup/OUT_LINUX"
REPO_DIR="$BUILD_DIR/repo"
SKEL_DIR="$LIVE_DIR/airootfs/etc/skel"
TEMP_DIR="$BUILD_DIR/temp"
@ -57,7 +57,7 @@ function cleanup() {
function fix_kit_permissions() {
# GitHub zip archives don't preserve the correct permissions
for d in .bin .cbin .kit_items .linux_items .pe_items Images; do
for d in docs images scripts setup; do
find "$ROOT_DIR/$d" -type d -exec chmod 755 "{}" \;
done
}
@ -80,7 +80,7 @@ function load_settings() {
if [[ "${1:-}" == "--edit" ]]; then
# Copy settings
if [[ ! -e "$BUILD_DIR/main.py" ]] || ask "Overwrite main.py?"; then
cp -bv "$ROOT_DIR/.bin/Scripts/settings/main.py" "$BUILD_DIR/main.py"
cp -bv "$ROOT_DIR/scripts/wk/cfg/main.py" "$BUILD_DIR/main.py"
dos2unix "$BUILD_DIR/main.py"
fi
@ -89,7 +89,7 @@ function load_settings() {
"$EDITOR" "$BUILD_DIR/main.py"
else
# Load settings from $LIVE_DIR
_main_path="$LIVE_DIR/airootfs/usr/local/bin/settings/main.py"
_main_path="$LIVE_DIR/airootfs/usr/local/bin/wk/cfg/main.py"
fi
# Load settings
@ -118,13 +118,13 @@ function copy_live_env() {
rm "$LIVE_DIR/syslinux"/*.cfg "$LIVE_DIR/syslinux"/*.png
# Add items
rsync -aI "$ROOT_DIR/.linux_items/include/" "$LIVE_DIR/"
rsync -aI "$ROOT_DIR/setup/linux/include/" "$LIVE_DIR/"
if [[ "${1:-}" != "--minimal" ]]; then
rsync -aI "$ROOT_DIR/.linux_items/include_x/" "$LIVE_DIR/"
rsync -aI "$ROOT_DIR/setup/linux/include_x/" "$LIVE_DIR/"
fi
mkdir -p "$LIVE_DIR/airootfs/usr/local/bin"
rsync -aI "$ROOT_DIR/.bin/Scripts/" "$LIVE_DIR/airootfs/usr/local/bin/"
cp -a "$BUILD_DIR/main.py" "$LIVE_DIR/airootfs/usr/local/bin/settings/"
rsync -aI "$ROOT_DIR/scripts/" "$LIVE_DIR/airootfs/usr/local/bin/"
cp -a "$BUILD_DIR/main.py" "$LIVE_DIR/airootfs/usr/local/bin/wk/cfg/"
}
function run_elevated() {
@ -155,8 +155,8 @@ function update_live_env() {
# Boot config (legacy)
mkdir -p "$LIVE_DIR/arch"
cp "$ROOT_DIR/Images/Pxelinux.jpg" "$LIVE_DIR/arch/pxelinux.jpg"
cp "$ROOT_DIR/Images/Syslinux.jpg" "$LIVE_DIR/arch/syslinux.jpg"
cp "$ROOT_DIR/images/Pxelinux.png" "$LIVE_DIR/arch/pxelinux.png"
cp "$ROOT_DIR/images/Syslinux.png" "$LIVE_DIR/arch/syslinux.png"
sed -i -r "s/_+/$KIT_NAME_FULL/" "$LIVE_DIR/syslinux/wk_head.cfg"
mkdir -p "$TEMP_DIR" 2>/dev/null
curl -Lo "$TEMP_DIR/wimboot.zip" "http://git.ipxe.org/releases/wimboot/wimboot-latest.zip"
@ -165,7 +165,7 @@ function update_live_env() {
# Boot config (UEFI)
mkdir -p "$LIVE_DIR/EFI/boot"
cp "/usr/share/refind/refind_x64.efi" "$LIVE_DIR/EFI/boot/bootx64.efi"
cp "$ROOT_DIR/Images/rEFInd.png" "$LIVE_DIR/EFI/boot/rEFInd.png"
cp "$ROOT_DIR/images/rEFInd.png" "$LIVE_DIR/EFI/boot/rEFInd.png"
rsync -aI "/usr/share/refind/drivers_x64/" "$LIVE_DIR/EFI/boot/drivers_x64/"
rsync -aI "/usr/share/refind/icons/" "$LIVE_DIR/EFI/boot/icons/" --exclude "/usr/share/refind/icons/svg"
sed -i "s/%ARCHISO_LABEL%/${label}/" "$LIVE_DIR/EFI/boot/refind.conf"
@ -199,12 +199,12 @@ function update_live_env() {
# Live packages
while read -r p; do
sed -i "/$p/d" "$LIVE_DIR/packages.x86_64"
done < "$ROOT_DIR/.linux_items/packages/live_remove"
cat "$ROOT_DIR/.linux_items/packages/live_add" >> "$LIVE_DIR/packages.x86_64"
done < "$ROOT_DIR/setup/linux/packages/live_remove"
cat "$ROOT_DIR/setup/linux/packages/live_add" >> "$LIVE_DIR/packages.x86_64"
if [[ "${1:-}" == "--minimal" ]]; then
cat "$ROOT_DIR/.linux_items/packages/live_add_min" >> "$LIVE_DIR/packages.x86_64"
cat "$ROOT_DIR/setup/linux/packages/live_add_min" >> "$LIVE_DIR/packages.x86_64"
else
cat "$ROOT_DIR/.linux_items/packages/live_add_x" >> "$LIVE_DIR/packages.x86_64"
cat "$ROOT_DIR/setup/linux/packages/live_add_x" >> "$LIVE_DIR/packages.x86_64"
fi
echo "[custom]" >> "$LIVE_DIR/pacman.conf"
echo "SigLevel = Optional TrustAll" >> "$LIVE_DIR/pacman.conf"
@ -246,7 +246,7 @@ function update_live_env() {
echo 'rm /root/.zlogin' >> "$LIVE_DIR/airootfs/root/customize_airootfs.sh"
sed -i -r '/.*PermitRootLogin.*/d' "$LIVE_DIR/airootfs/root/customize_airootfs.sh"
echo "sed -i -r '/.*PermitRootLogin.*/d' /etc/ssh/sshd_config" >> "$LIVE_DIR/airootfs/root/customize_airootfs.sh"
cp "$ROOT_DIR/.linux_items/authorized_keys" "$SKEL_DIR/.ssh/authorized_keys"
cp "$ROOT_DIR/setup/linux/authorized_keys" "$SKEL_DIR/.ssh/authorized_keys"
# Root user
echo "echo 'root:$ROOT_PASSWORD' | chpasswd" >> "$LIVE_DIR/airootfs/root/customize_airootfs.sh"
@ -279,11 +279,11 @@ function update_live_env() {
# Wallpaper
mkdir -p "$LIVE_DIR/airootfs/usr/share/wallpaper"
cp "$ROOT_DIR/Images/Linux.jpg" "$LIVE_DIR/airootfs/usr/share/wallpaper/burned.in"
cp "$ROOT_DIR/images/Linux.png" "$LIVE_DIR/airootfs/usr/share/wallpaper/burned.in"
fi
# WiFi
cp "$ROOT_DIR/.linux_items/known_networks" "$LIVE_DIR/airootfs/root/known_networks"
cp "$ROOT_DIR/setup/linux/known_networks" "$LIVE_DIR/airootfs/root/known_networks"
echo "add-known-networks --user=$username" >> "$LIVE_DIR/airootfs/root/customize_airootfs.sh"
}
@ -315,7 +315,7 @@ function update_repo() {
makepkg -d
popd >/dev/null
mv -n $p/*xz "$REPO_DIR"/
done < "$ROOT_DIR/.linux_items/packages/aur"
done < "$ROOT_DIR/setup/linux/packages/aur"
popd >/dev/null
# Build custom repo database
@ -329,7 +329,7 @@ function install_deps() {
packages=
while read -r line; do
packages="$packages $line"
done < "$ROOT_DIR/.linux_items/packages/dependencies"
done < "$ROOT_DIR/setup/linux/packages/dependencies"
run_elevated pacman -Syu --needed --noconfirm $packages
}
@ -347,7 +347,7 @@ function build_iso() {
chmod 600 "$LIVE_DIR/airootfs/etc/skel/.ssh/id_rsa"
# Removing cached (and possibly outdated) custom repo packages
for package in $(cat "$ROOT_DIR/.linux_items/packages/aur"); do
for package in $(cat "$ROOT_DIR/setup/linux/packages/aur"); do
for p in /var/cache/pacman/pkg/*${package}*; do
if [[ -f "${p}" ]]; then
rm "${p}"

View file

@ -12,3 +12,6 @@ PS1='[\u@\h \W]\$ '
# Update LS_COLORS
eval $(dircolors ~/.dircolors)
# WizardKit
export PYTHONPATH='/usr/local/bin'

View file

@ -9,3 +9,4 @@ source $ZSH/oh-my-zsh.sh
# Wizard Kit
. $HOME/.aliases
eval $(dircolors ~/.dircolors)
export PYTHONPATH="/usr/local/bin"