diff --git a/scripts/wk/clone/block_pair.py b/scripts/wk/clone/block_pair.py index 0c98e33d..bd667105 100644 --- a/scripts/wk/clone/block_pair.py +++ b/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() diff --git a/scripts/wk/clone/ddrescue.py b/scripts/wk/clone/ddrescue.py index 0ebb2fcc..70557625 100644 --- a/scripts/wk/clone/ddrescue.py +++ b/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.""" - zero_fill_destination(state, dry_run=dry_run) - if state.mode == 'Clone': + """Show recovery finalization options and run selected functions.""" + options = ( + '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: - """Zero-fill any gaps and space on destination beyond the source size.""" +def zero_fill_gaps( + 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 ( - (percent_recovered < 100 - or (full_disk_clone and state.source.size < state.destination.size)) - and cli.ask('Fill gaps with zeros?')): + if percent_recovered == 100 and not (larger_destination and extend_to_end): 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'{hex(destination_size - block_pair.size)} ?' + f'\n{hex(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', diff --git a/scripts/wk/clone/state.py b/scripts/wk/clone/state.py index 152a5350..81192f3a 100644 --- a/scripts/wk/clone/state.py +++ b/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: diff --git a/scripts/wk/hw/diags.py b/scripts/wk/hw/diags.py index d01ae06f..e14a299b 100644 --- a/scripts/wk/hw/diags.py +++ b/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 = []