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

scripts/wk/clone/block_pair.py

@@ -264,7 +264,7 @@ class BlockPair():
 
 
 # Functions
-def add_clone_block_pairs(state) -> None:
+def add_clone_block_pairs(state) -> list[hw_disk.Disk]:
   """Add device to device block pairs and set settings if necessary."""
   source_sep = get_partition_separator(state.source.path.name)
   dest_sep = get_partition_separator(state.destination.path.name)
@@ -276,6 +276,7 @@ def add_clone_block_pairs(state) -> None:
 
   # Add pairs from previous run
   if settings['Partition Mapping']:
+    source_parts = []
     for part_map in settings['Partition Mapping']:
       bp_source = hw_disk.Disk(
         f'{state.source.path}{source_sep}{part_map[0]}',
@@ -283,8 +284,9 @@ def add_clone_block_pairs(state) -> None:
       bp_dest = pathlib.Path(
         f'{state.destination.path}{dest_sep}{part_map[1]}',
         )
+      source_parts.append(bp_source)
       state.add_block_pair(bp_source, bp_dest)
-      return
+    return source_parts
 
   # Add pairs from selection
   source_parts = menus.select_disk_parts('Clone', state.source)
@@ -292,7 +294,7 @@ def add_clone_block_pairs(state) -> None:
     # Whole disk (or single partition via args), skip settings
     bp_dest = state.destination.path
     state.add_block_pair(state.source, bp_dest)
-    return
+    return source_parts
 
   # New run, use new settings file
   settings['Needs Format'] = True
@@ -327,13 +329,19 @@ def add_clone_block_pairs(state) -> None:
   # Save settings
   state.save_settings(settings)
 
+  # Done
+  return source_parts
 
-def add_image_block_pairs(state) -> None:
+
+def add_image_block_pairs(state) -> list[hw_disk.Disk]:
   """Add device to image file block pairs."""
   source_parts = menus.select_disk_parts(state.mode, state.source)
   for part in source_parts:
     state.add_block_pair(part, state.destination)
 
+  # Done
+  return source_parts
+
 
 def build_block_pair_report(block_pairs, settings) -> list:
   """Build block pair report, returns list."""
@@ -399,7 +407,7 @@ def build_sfdisk_partition_line(table_type, dev_path, size, details) -> str:
 
   # Set dest type
   if re.match(r'^0x\w+$', source_type):
-    # Both source and dest are MBR
+    # Source is a MBR type
     source_table_type = 'MBR'
     if table_type == 'MBR':
       dest_type = source_type.replace('0x', '').lower()

scripts/wk/clone/ddrescue.py

@@ -18,7 +18,7 @@ from docopt import docopt
 from wk import cfg, exe, io, log, std
 from wk.cfg.ddrescue import DDRESCUE_SPECIFIC_PASS_SETTINGS
 from wk.clone import menus
-from wk.clone.state import State
+from wk.clone.state import State, mark_non_recovered_as_non_tried
 from wk.hw import disk as hw_disk
 from wk.hw.smart import (
   check_attributes,
@@ -208,10 +208,64 @@ def get_stats(output: str | None = None) -> dict[str, Any]:
 
 
 def finalize_recovery(state: State, dry_run: bool = True) -> None:
-  """Show recovery finalization options."""
+  """Show recovery finalization options and run selected functions."""
-  zero_fill_destination(state, dry_run=dry_run)
+  options = (
-  if state.mode == 'Clone':
+    'Relocate Backup GPT',
+    'Zero-fill Gaps',
+    'Zero-fill Extra Space',
+    )
+
+  # Get destination size
+  dest_size = -1
+  if hasattr(state.destination, 'size'):
+    # hw_disk.Disk
+    dest_size = state.destination.size
+  if hasattr(state.destination, 'stat'):
+    # pathlib.Path
+    dest_size = state.destination.stat().st_size
+
+  # Run checks to disable item(s)
+  whole_disk = bool(
+    len(state.block_pairs) == 1
+    and not state.source.parent
+    )
+  disable_gpt_option = not bool(
+    ## Breakdown of below tests:
+    ##  Only offer this option when cloning a whole, non-child device
+    ##  where the source is using a GUID_Partition_Table
+    ##  and the source is smaller than the destination
+      whole_disk
+      and str(state.source.raw_details.get('pttype', 'Unknown')).lower() == 'gpt'
+      and state.source.size < dest_size
+      )
+
+  # Build menu
+  menu = cli.Menu(title=ansi.color_string('ddrescue TUI: Finalization', 'GREEN'))
+  menu.separator = ' '
+  menu.add_action('Start')
+  menu.add_action('Quit')
+  for name in options:
+    details = {'Selected': True}
+    if 'GPT' in name and disable_gpt_option:
+      details['Disabled'] = True
+      details['Selected'] = False
+    menu.add_option(name, details)
+
+  # Show menu
+  selection = menu.advanced_select()
+  if 'Quit' in selection:
+    return
+
+  # Run functions
+  if menu.options['Relocate Backup GPT']['Selected']:
     relocate_backup_gpt(state, dry_run=dry_run)
+  if menu.options['Zero-fill Gaps']['Selected']:
+    zero_fill_gaps(
+      state,
+      dest_size=dest_size,
+      dry_run=dry_run,
+      extend_to_end=menu.options['Zero-fill Extra Space']['Selected'],
+      )
 
 
 def is_missing_source_or_destination(state) -> bool:
@@ -636,25 +690,28 @@ def source_or_destination_changed(state) -> bool:
   return changed
 
 
-def zero_fill_destination(state: State, dry_run: bool = True) -> None:
+def zero_fill_gaps(
-  """Zero-fill any gaps and space on destination beyond the source size."""
+    state: State,
+    dest_size: int,
+    dry_run: bool = True,
+    extend_to_end: bool = False,
+    ) -> None:
+  """Zero-fill any gaps on the destination."""
+  #fake_settings_menu = menus.settings(state.mode)
   full_disk_clone = False
-  larger_destination = False
+  larger_destination = state.source.size < dest_size
   percent_recovered = state.get_percent_recovered()
   if state.mode == 'Clone' and len(state.block_pairs) == 1:
     full_disk_clone = True
 
   # Bail early
-  if not (
+  if percent_recovered == 100 and not (larger_destination and extend_to_end):
-      (percent_recovered < 100
-        or (full_disk_clone and state.source.size < state.destination.size))
-      and cli.ask('Fill gaps with zeros?')):
     return
 
   for block_pair in state.block_pairs:
-    destination_size = block_pair.size
+    domain_size = block_pair.size
     if (full_disk_clone and state.source.size < state.destination.size):
-      destination_size = state.destination.size
+      domain_size = dest_size
       larger_destination = True
 
     # Prep zero-fill map file
@@ -662,11 +719,13 @@ def zero_fill_destination(state: State, dry_run: bool = True) -> None:
       f'{block_pair.map_path.stem}_zero-fill',
       )
     io.copy_file(block_pair.map_path, zero_map_path, overwrite=True)
-    if larger_destination:
+    mark_non_recovered_as_non_tried(zero_map_path)
+    if full_disk_clone and larger_destination and extend_to_end:
+      # Extend domain size
       with open(zero_map_path, 'a', encoding='utf-8') as f:
         f.write(
-          f'{hex(block_pair.size)}  '
+          f'\n{hex(block_pair.size)}  '
-          f'{hex(destination_size - block_pair.size)}  ?'
+          f'{hex(domain_size - block_pair.size)}  ?'
           )
 
     # Build cmd
@@ -674,7 +733,7 @@ def zero_fill_destination(state: State, dry_run: bool = True) -> None:
       'sudo',
       'ddrescue',
       '--force',
-      f'--size={destination_size}',
+      f'--size={domain_size}',
       '--binary-prefixes',
       '--complete-only',
       '--data-preview=5',

scripts/wk/clone/state.py

@@ -441,11 +441,12 @@ class State():
     # Add block pairs
     if advanced_selection:
       if self.mode == 'Clone':
-        add_clone_block_pairs(self)
+        source_parts = add_clone_block_pairs(self)
       else:
-        add_image_block_pairs(self)
+        source_parts = add_image_block_pairs(self)
     else:
       if self.mode == 'Clone':
+        source_parts.append(hw_disk.Disk(self.source.path))
         self.add_block_pair(
           hw_disk.Disk(self.source.path),
           pathlib.Path(self.destination.path),
@@ -594,29 +595,16 @@ class State():
 
   def retry_all_passes(self) -> None:
     """Prep block_pairs for a retry recovery attempt."""
-    bad_statuses = ('*', '/', '-')
     LOG.warning('Updating block_pairs for retry')
 
     # Update all block_pairs
     for pair in self.block_pairs:
-      map_data = []
+      mark_non_recovered_as_non_tried(pair.map_path)
 
       # Reset status strings
       for name in pair.status.keys():
         pair.status[name] = 'Pending'
 
-      # Mark all non-trimmed, non-scraped, and bad areas as non-tried
-      with open(pair.map_path, 'r', encoding='utf-8') as _f:
-        for line in _f.readlines():
-          line = line.strip()
-          if line.startswith('0x') and line.endswith(bad_statuses):
-            line = f'{line[:-1]}?'
-          map_data.append(line)
-
-      # Save updated map
-      with open(pair.map_path, 'w', encoding='utf-8') as _f:
-        _f.write('\n'.join(map_data))
-
       # Reinitialize status
       pair.set_initial_status()
 
@@ -1151,6 +1139,24 @@ def get_working_dir(
   return working_dir
 
 
+def mark_non_recovered_as_non_tried(map_path: pathlib.Path) -> None:
+  """Read map file and mark all non-recovered sections as non-tried."""
+  bad_statuses = ('*', '/', '-')
+  map_data = []
+
+  # Mark all non-trimmed, non-scraped, and bad areas as non-tried
+  with open(map_path, 'r', encoding='utf-8') as _f:
+    for line in _f.readlines():
+      line = line.strip()
+      if line.startswith('0x') and line.endswith(bad_statuses):
+        line = f'{line[:-1]}?'
+      map_data.append(line)
+
+  # Save updated map
+  with open(map_path, 'w', encoding='utf-8') as _f:
+    _f.write('\n'.join(map_data))
+
+
 def select_disk_obj(label:str, disk_menu: cli.Menu, disk_path: str) -> hw_disk.Disk:
   """Get disk based on path or menu selection, returns Disk."""
   if not disk_path:

scripts/wk/hw/diags.py

@@ -405,6 +405,7 @@ def cpu_tests_end(state) -> None:
 
 def cpu_test_cooling(state, test_object, test_mode=False) -> None:
   """CPU cooling test via sensor data assessment."""
+  _ = test_mode
   LOG.info('CPU Test (Cooling)')
 
   # Bail early
@@ -569,6 +570,7 @@ def cpu_test_sysbench(state, test_object, test_mode=False) -> None:
 
 def disk_attribute_check(state, test_objects, test_mode=False) -> None:
   """Disk attribute check."""
+  _ = test_mode
   LOG.info('Disk Attribute Check')
   for test in test_objects:
     disk_smart_status_check(test.dev, mid_run=False)
@@ -651,6 +653,7 @@ def disk_io_benchmark(
 
 def disk_self_test(state, test_objects, test_mode=False) -> None:
   """Disk self-test if available."""
+  _ = test_mode
   LOG.info('Disk Self-Test(s)')
   aborted = False
   threads = []