Merge branch 'project-overhaul' into dev

scripts/wk.prev/settings/windows_builds.py

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

scripts/wk/exe.py

@@ -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:
+  if cmd[0] == 'sudo':
-    cmd.pop(0)
+    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)

scripts/wk/hw/ddrescue.py

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

scripts/wk/hw/sensors.py

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

scripts/wk/kit/ufd.py

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

scripts/wk/log.py

@@ -79,14 +79,20 @@ def get_root_logger_path():
   return log_path
 
 
-def remove_empty_log():
+def remove_empty_log(log_path=None):
-  """Remove log if empty."""
+  """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
+  # Create new handler (preserving formatter settings)
-  remove_empty_log()
-
-  # Create new cur_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
+  # Remove old_handler and log if empty
-  root_logger.removeHandler(cur_handler)
+  root_logger.removeHandler(old_handler)
+  old_handler.close()
+  remove_empty_log(old_path)
+
+  # Add new handler
   root_logger.addHandler(new_handler)
 
 

scripts/wk/net.py

@@ -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',uid={os.getuid()}' # pylint: disable=no-member
-        f',gid={os.getgid()}'
+        f',gid={os.getgid()}' # pylint: disable=no-member
         f',username={username}'
         f',{"password=" if password else "guest"}{password}'
         ),

scripts/wk/os/win.py

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

setup/linux/include/airootfs/etc/skel/.aliases

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