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

Merge branch 'project-overhaul' into dev

Browse source
This commit is contained in:
2Shirt 2021-01-10 17:24:54 -07:00
commit 2497fec389
Signed by: 2Shirt
GPG key ID: 152FAC923B0E132C
9 changed files with 383 additions and 28 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

33
scripts/wk.prev/settings/windows_builds.py
View file

@ -205,6 +205,7 @@ WINDOWS_BUILDS = {
'18358': ('10', None, '19H1', None, 'preview build'),
'18361': ('10', None, '19H1', None, 'preview build'),
'18362': ('10', 'v1903', '19H1', 'May 2019 Update', None),
'18363': ('10', 'v1909', '19H2', 'November 2019 Update', None),
'18836': ('10', None, '20H1', None, 'preview build'),
'18841': ('10', None, '20H1', None, 'preview build'),
'18845': ('10', None, '20H1', None, 'preview build'),
@ -218,6 +219,38 @@ WINDOWS_BUILDS = {
'18894': ('10', None, '20H1', None, 'preview build'),
'18895': ('10', None, '20H1', None, 'preview build'),
'18898': ('10', None, '20H1', None, 'preview build'),
'18908': ('10', None, '20H1', None, 'preview build'),
'18912': ('10', None, '20H1', None, 'preview build'),
'18917': ('10', None, '20H1', None, 'preview build'),
'18922': ('10', None, '20H1', None, 'preview build'),
'18932': ('10', None, '20H1', None, 'preview build'),
'18936': ('10', None, '20H1', None, 'preview build'),
'18941': ('10', None, '20H1', None, 'preview build'),
'18945': ('10', None, '20H1', None, 'preview build'),
'18950': ('10', None, '20H1', None, 'preview build'),
'18956': ('10', None, '20H1', None, 'preview build'),
'18963': ('10', None, '20H1', None, 'preview build'),
'18965': ('10', None, '20H1', None, 'preview build'),
'18970': ('10', None, '20H1', None, 'preview build'),
'18975': ('10', None, '20H1', None, 'preview build'),
'18980': ('10', None, '20H1', None, 'preview build'),
'18985': ('10', None, '20H1', None, 'preview build'),
'18990': ('10', None, '20H1', None, 'preview build'),
'18995': ('10', None, '20H1', None, 'preview build'),
'18999': ('10', None, '20H1', None, 'preview build'),
'19002': ('10', None, '20H1', None, 'preview build'),
'19008': ('10', None, '20H1', None, 'preview build'),
'19013': ('10', None, '20H1', None, 'preview build'),
'19018': ('10', None, '20H1', None, 'preview build'),
'19023': ('10', None, '20H1', None, 'preview build'),
'19025': ('10', None, '20H1', None, 'preview build'),
'19028': ('10', None, '20H1', None, 'preview build'),
'19030': ('10', None, '20H1', None, 'preview build'),
'19033': ('10', None, '20H1', None, 'preview build'),
'19035': ('10', None, '20H1', None, 'preview build'),
'19037': ('10', None, '20H1', None, 'preview build'),
'19041': ('10', 'v2004', '20H1', 'May 2020 Update', None),
'19042': ('10', 'v20H2', '20H2', 'October 2020 Update', None),
}

28
scripts/wk/exe.py
View file

@ -6,6 +6,7 @@ import logging
import os
import re
import subprocess
import time
from threading import Thread
from queue import Queue, Empty
@ -79,8 +80,9 @@ def build_cmd_kwargs(cmd, minimized=False, pipe=True, shell=False, **kwargs):
}
# Strip sudo if appropriate
if cmd[0] == 'sudo' and os.name == 'posix' and os.geteuid() == 0:
cmd.pop(0)
if cmd[0] == 'sudo':
if os.name == 'posix' and os.geteuid() == 0: # pylint: disable=no-member
cmd.pop(0)
# Add additional kwargs if applicable
for key in 'check cwd encoding errors stderr stdin stdout'.split():
@ -214,6 +216,28 @@ def start_thread(function, args=None, daemon=True):
return thread
def stop_process(proc, graceful=True):
"""Stop process.
NOTES: proc should be a subprocess.Popen obj.
If graceful is True then a SIGTERM is sent before SIGKILL.
"""
# Graceful exit
if graceful:
if os.name == 'posix' and os.geteuid() != 0: # pylint: disable=no-member
run_program(['sudo', 'kill', str(proc.pid)], check=False)
else:
proc.terminate()
time.sleep(2)
# Force exit
if os.name == 'posix' and os.geteuid() != 0: # pylint: disable=no-member
run_program(['sudo', 'kill', '-9', str(proc.pid)], check=False)
else:
proc.kill()
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)

58
scripts/wk/hw/ddrescue.py
View file

@ -1390,6 +1390,48 @@ def build_sfdisk_partition_line(table_type, dev_path, size, details):
return line
def check_destination_health(destination):
"""Check destination health, returns str."""
result = ''
# Bail early
if not isinstance(destination, hw_obj.Disk):
# Return empty string
return result
# Run safety checks
try:
destination.safety_checks()
except hw_obj.CriticalHardwareError:
result = 'Critical hardware error detected on destination'
except hw_obj.SMARTSelfTestInProgressError:
result = 'SMART self-test in progress on destination'
except hw_obj.SMARTNotSupportedError:
pass
# Done
return result
def check_for_missing_items(state):
"""Check if source or destination dissapeared."""
items = {
'Source': state.source,
'Destination': state.destination,
}
for name, item in items.items():
if not item:
continue
if hasattr(item, 'path'):
if not item.path.exists():
std.print_error(f'{name} disappeared')
elif hasattr(item, 'exists'):
if not item.exists():
std.print_error(f'{name} disappeared')
else:
LOG.error('Unknown %s type: %s', name, item)
def clean_working_dir(working_dir):
"""Clean working directory to ensure a fresh recovery session.
@ -1686,7 +1728,8 @@ def main():
state = State()
try:
state.init_recovery(args)
except std.GenericAbort:
except (FileNotFoundError, std.GenericAbort):
check_for_missing_items(state)
std.abort()
# Show menu
@ -1800,6 +1843,7 @@ def mount_raw_image_macos(path):
def run_ddrescue(state, block_pair, pass_name, settings, dry_run=True):
# pylint: disable=too-many-statements
"""Run ddrescue using passed settings."""
cmd = build_ddrescue_cmd(block_pair, pass_name, settings)
state.update_progress_pane('Active')
@ -1834,6 +1878,13 @@ def run_ddrescue(state, block_pair, pass_name, settings, dry_run=True):
if _i % 30 == 0:
# Update SMART pane
_update_smart_pane()
# Check destination
warning_message = check_destination_health(state.destination)
if warning_message:
# Error detected on destination, stop recovery
exe.stop_process(proc)
break
if _i % 60 == 0:
# Clear ddrescue pane
tmux.clear_pane()
@ -1852,7 +1903,7 @@ def run_ddrescue(state, block_pair, pass_name, settings, dry_run=True):
LOG.warning('ddrescue stopped by user')
warning_message = 'Aborted'
std.sleep(2)
exe.run_program(['sudo', 'kill', str(proc.pid)], check=False)
exe.stop_process(proc, graceful=False)
break
except subprocess.TimeoutExpired:
# Continue to next loop to update panes
@ -1926,7 +1977,8 @@ def run_recovery(state, main_menu, settings_menu, dry_run=True):
state.mark_started()
try:
run_ddrescue(state, pair, pass_name, settings, dry_run=dry_run)
except (KeyboardInterrupt, std.GenericAbort):
except (FileNotFoundError, KeyboardInterrupt, std.GenericAbort):
check_for_missing_items(state)
abort = True
break

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

@ -150,7 +150,7 @@ class Sensors():
# Get temps
for i in range(seconds):
self.update_sensor_data()
self.update_sensor_data(exit_on_thermal_limit=False)
sleep(1)
# Calculate averages
@ -158,7 +158,15 @@ class Sensors():
for sources in adapters.values():
for source_data in sources.values():
temps = source_data['Temps']
source_data[temp_label] = sum(temps) / len(temps)
try:
source_data[temp_label] = sum(temps) / len(temps)
except ZeroDivisionError:
# Going to use unrealistic 0°C instead
LOG.error(
'No temps saved for %s',
source_data.get('label', 'UNKNOWN'),
)
source_data[temp_label] = 0
def start_background_monitor(
self, out_path,
@ -253,7 +261,7 @@ def fix_sensor_name(name):
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'T(ccd\s+\d+|ctl|die)', r'CPU (T\1)', name, re.IGNORECASE)
name = re.sub(r'\s+', ' ', name)
return name
@ -286,7 +294,7 @@ def get_sensor_data_linux():
## 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'):
if label.startswith('fan') or label.startswith('in') or label.startswith('curr'):
# Skip fan RPMs and voltages
continue
if 'input' in label:

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

@ -212,7 +212,8 @@ def copy_source(source, items, overwrite=False):
linux.unmount('/mnt/Source')
# Raise exception if item(s) were not found
raise FileNotFoundError('One or more items not found')
if items_not_found:
raise FileNotFoundError('One or more items not found')
def create_table(dev_path, use_mbr=False):

39
scripts/wk/log.py
View file

@ -79,14 +79,20 @@ def get_root_logger_path():
return log_path
def remove_empty_log():
"""Remove log if empty."""
def remove_empty_log(log_path=None):
"""Remove log if empty.
NOTE: Under Windows an empty log is 2 bytes long.
"""
is_empty = False
# Get log path
if not log_path:
log_path = get_root_logger_path()
# 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
is_empty = log_path and log_path.exists() and log_path.stat().st_size <= 2
except (FileNotFoundError, AttributeError):
# File doesn't exist or couldn't verify it's empty
pass
@ -122,34 +128,35 @@ 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)
old_handler = None
old_path = get_root_logger_path()
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
old_handler = handler
break
if not cur_handler:
if not old_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)
shutil.copy(old_path, new_path)
# Remove old log if empty
remove_empty_log()
# Create new cur_handler (preserving formatter settings)
# Create new handler (preserving formatter settings)
new_handler = logging.FileHandler(new_path, mode='a')
new_handler.setFormatter(cur_handler.formatter)
new_handler.setFormatter(old_handler.formatter)
# Replace current handler
root_logger.removeHandler(cur_handler)
# Remove old_handler and log if empty
root_logger.removeHandler(old_handler)
old_handler.close()
remove_empty_log(old_path)
# Add new handler
root_logger.addHandler(new_handler)

4
scripts/wk/net.py
View file

@ -122,8 +122,8 @@ def mount_network_share(details, mount_point=None, read_write=False):
'-t', 'cifs',
'-o', (
f'{"rw" if read_write else "ro"}'
f',uid={os.getuid()}'
f',gid={os.getgid()}'
f',uid={os.getuid()}' # pylint: disable=no-member
f',gid={os.getgid()}' # pylint: disable=no-member
f',username={username}'
f',{"password=" if password else "guest"}{password}'
),

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

@ -5,6 +5,9 @@ import logging
import os
import pathlib
import platform
import winreg
from contextlib import suppress
from wk.borrowed import acpi
from wk.exe import run_program
@ -15,6 +18,39 @@ from wk.std import GenericError, GenericWarning, sleep
# STATIC VARIABLES
LOG = logging.getLogger(__name__)
KNOWN_DATA_TYPES = {
'BINARY': winreg.REG_BINARY,
'DWORD': winreg.REG_DWORD,
'DWORD_LITTLE_ENDIAN': winreg.REG_DWORD_LITTLE_ENDIAN,
'DWORD_BIG_ENDIAN': winreg.REG_DWORD_BIG_ENDIAN,
'EXPAND_SZ': winreg.REG_EXPAND_SZ,
'LINK': winreg.REG_LINK,
'MULTI_SZ': winreg.REG_MULTI_SZ,
'NONE': winreg.REG_NONE,
'QWORD': winreg.REG_QWORD,
'QWORD_LITTLE_ENDIAN': winreg.REG_QWORD_LITTLE_ENDIAN,
'SZ': winreg.REG_SZ,
}
KNOWN_HIVES = {
'HKCR': winreg.HKEY_CLASSES_ROOT,
'HKCU': winreg.HKEY_CURRENT_USER,
'HKLM': winreg.HKEY_LOCAL_MACHINE,
'HKU': winreg.HKEY_USERS,
'HKEY_CLASSES_ROOT': winreg.HKEY_CLASSES_ROOT,
'HKEY_CURRENT_USER': winreg.HKEY_CURRENT_USER,
'HKEY_LOCAL_MACHINE': winreg.HKEY_LOCAL_MACHINE,
'HKEY_USERS': winreg.HKEY_USERS,
}
KNOWN_HIVE_NAMES = {
winreg.HKEY_CLASSES_ROOT: 'HKCR',
winreg.HKEY_CURRENT_USER: 'HKCU',
winreg.HKEY_LOCAL_MACHINE: 'HKLM',
winreg.HKEY_USERS: 'HKU',
winreg.HKEY_CLASSES_ROOT: 'HKEY_CLASSES_ROOT',
winreg.HKEY_CURRENT_USER: 'HKEY_CURRENT_USER',
winreg.HKEY_LOCAL_MACHINE: 'HKEY_LOCAL_MACHINE',
winreg.HKEY_USERS: 'HKEY_USERS',
}
OS_VERSION = float(platform.win32_ver()[0])
REG_MSISERVER = r'HKLM\SYSTEM\CurrentControlSet\Control\SafeBoot\Network\MSIServer'
SLMGR = pathlib.Path(f'{os.environ.get("SYSTEMROOT")}/System32/slmgr.vbs')
@ -177,5 +213,198 @@ def run_sfc_scan():
raise OSError
# Registry Functions
def reg_delete_key(hive, key, recurse=False):
"""Delete a key from the registry.
NOTE: If recurse is False then it will only work on empty keys.
"""
hive = reg_get_hive(hive)
hive_name = KNOWN_HIVE_NAMES.get(hive, '???')
# Delete subkeys first
if recurse:
with suppress(WindowsError), winreg.OpenKey(hive, key) as open_key:
while True:
subkey = fr'{key}\{winreg.EnumKey(open_key, 0)}'
reg_delete_key(hive, subkey, recurse=recurse)
# Delete key
try:
winreg.DeleteKey(hive, key)
LOG.warning(r'Deleting registry key: %s\%s', hive_name, key)
except FileNotFoundError:
# Ignore
pass
except PermissionError:
LOG.error(r'Failed to delete registry key: %s\%s', hive_name, key)
if recurse:
# Re-raise exception
raise
# recurse is not True so assuming we tried to remove a non-empty key
msg = fr'Refusing to remove non-empty key: {hive_name}\{key}'
raise FileExistsError(msg)
def reg_delete_value(hive, key, value):
"""Delete a value from the registry."""
access = winreg.KEY_ALL_ACCESS
hive = reg_get_hive(hive)
hive_name = KNOWN_HIVE_NAMES.get(hive, '???')
# Delete value
with winreg.OpenKey(hive, key, access=access) as open_key:
try:
winreg.DeleteValue(open_key, value)
LOG.warning(
r'Deleting registry value: %s\%s "%s"', hive_name, key, value,
)
except FileNotFoundError:
# Ignore
pass
except PermissionError:
LOG.error(
r'Failed to delete registry value: %s\%s "%s"', hive_name, key, value,
)
# Re-raise exception
raise
def reg_get_hive(hive):
"""Get winreg HKEY constant from string, returns HKEY constant."""
if isinstance(hive, int):
# Assuming we're already a winreg HKEY constant
pass
else:
hive = KNOWN_HIVES[hive.upper()]
# Done
return hive
def reg_get_data_type(data_type):
"""Get registry data type from string, returns winreg constant."""
if isinstance(data_type, int):
# Assuming we're already a winreg value type constant
pass
else:
data_type = KNOWN_DATA_TYPES[data_type.upper()]
# Done
return data_type
def reg_key_exists(hive, key):
"""Test if the specified hive/key exists, returns bool."""
exists = False
hive = reg_get_hive(hive)
# Query key
try:
winreg.QueryValue(hive, key)
except FileNotFoundError:
# Leave set to False
pass
else:
exists = True
# Done
return exists
def reg_read_value(hive, key, value, force_32=False, force_64=False):
"""Query value from hive/hey, returns multiple types."""
access = winreg.KEY_READ
data = None
hive = reg_get_hive(hive)
# Set access
if force_32:
access = access | winreg.KEY_WOW64_32KEY
elif force_64:
access = access | winreg.KEY_WOW64_64KEY
# Query value
with winreg.OpenKey(hive, key, access=access) as open_key:
# Returning first part of tuple and ignoreing type
data = winreg.QueryValueEx(open_key, value)[0]
# Done
return data
def reg_write_settings(settings):
"""Set registry values in bulk from a custom data structure.
Data structure should be as follows:
EXAMPLE_SETTINGS = {
# See KNOWN_HIVES for valid hives
'HKLM': {
r'Software\\2Shirt\\WizardKit': (
# Value tuples should be in the form:
# (name, data, data-type, option),
# See KNOWN_DATA_TYPES for valid types
# The option item is optional
('Sample Value #1', 'Sample Data', 'SZ'),
('Sample Value #2', 14, 'DWORD'),
),
# An empty key will be created if no values are specified
r'Software\\2Shirt\\WizardKit\\Empty': (),
r'Software\\2Shirt\\WizardKit\\Test': (
('Sample Value #3', 14000000000000, 'QWORD'),
),
},
'HKCU': {
r'Software\\2Shirt\\WizardKit': (
# The 4th item forces using the 32-bit registry
# See reg_set_value() for valid options
('Sample Value #4', 'Sample Data', 'SZ', '32'),
),
},
}
"""
for hive, keys in settings.items():
hive = reg_get_hive(hive)
for key, values in keys.items():
if not values:
# Create an empty key
winreg.CreateKey(hive, key)
for value in values:
reg_set_value(hive, key, *value)
def reg_set_value(hive, key, name, data, data_type, option=None):
# pylint: disable=too-many-arguments
"""Set value for hive/key."""
access = winreg.KEY_WRITE
data_type = reg_get_data_type(data_type)
hive = reg_get_hive(hive)
option = str(option)
# Safety check
if not name and option in ('32', '64'):
raise NotImplementedError(
'Unable to set default values using alternate registry views',
)
# Set access
if option == '32':
access = access | winreg.KEY_WOW64_32KEY
elif option == '64':
access = access | winreg.KEY_WOW64_64KEY
# Create key
winreg.CreateKeyEx(hive, key, access=access)
# Set value
if name:
with winreg.OpenKey(hive, key, access=access) as open_key:
winreg.SetValueEx(open_key, name, 0, data_type, data)
else:
# Set default value instead
winreg.SetValue(hive, key, data_type, data)
if __name__ == '__main__':
print("This file is not meant to be called directly.")

1
setup/linux/include/airootfs/etc/skel/.aliases
View file

@ -11,6 +11,7 @@ alias du='du -sch --apparent-size'
alias fix-perms='find -type d -exec chmod 755 "{}" \; && find -type f -exec chmod 644 "{}" \;'
alias hexedit='hexedit --color'
alias hw-info='sudo hw-info | less -S'
alias ip='ip -br -color'
alias less='less -S'
alias ls='ls --color=auto'
alias mkdir='mkdir -p'