diff --git a/scripts/wk/clone/ddrescue.py b/scripts/wk/clone/ddrescue.py index f019e0de..fa9c570a 100644 --- a/scripts/wk/clone/ddrescue.py +++ b/scripts/wk/clone/ddrescue.py @@ -17,7 +17,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, @@ -177,10 +177,68 @@ def get_ddrescue_settings(settings_menu) -> list: 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 = ( + 'Zero-fill Extra Space', + ) + toggles = ( + 'Relocate Backup GPT', + 'Zero-fill Gaps', + ) + + # 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_toggle = 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: + menu.add_option(name, details={'Selected': True}) + for name in toggles: + details = {'Selected': True} + if 'GPT' in name and disable_gpt_toggle: + details['Disabled'] = True + details['Selected'] = False + menu.add_toggle(name, details) + + # Show menu + selection = menu.advanced_select() + if 'Quit' in selection: + return + + # Run functions + if menu.toggles['Relocate Backup GPT']['Selected']: relocate_backup_gpt(state, dry_run=dry_run) + if menu.toggles['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: @@ -587,25 +645,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 @@ -613,11 +674,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 @@ -625,7 +688,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 f0dd9c66..f9badfce 100644 --- a/scripts/wk/clone/state.py +++ b/scripts/wk/clone/state.py @@ -497,29 +497,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() @@ -1009,6 +996,24 @@ def get_working_dir(mode, destination, force_local=False) -> pathlib.Path: 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: