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

Made working_dir a State() variable

Browse source
This commit is contained in:
2Shirt 2020-01-03 16:08:38 -07:00
parent 4f2b31c705
commit a4b5e81ef1
Signed by: 2Shirt
GPG key ID: 152FAC923B0E132C
1 changed files with 37 additions and 37 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -256,27 +256,27 @@ class State():
def __init__(self): def __init__(self):
self.block_pairs = [] self.block_pairs = []
self.destination = None self.destination = None
self.disks = []
self.layout = cfg.ddrescue.TMUX_LAYOUT.copy() self.layout = cfg.ddrescue.TMUX_LAYOUT.copy()
self.log_dir = None self.log_dir = None
self.panes = {} self.panes = {}
self.source = None self.source = None
self.working_dir = None
# Start a background process to maintain layout # Start a background process to maintain layout
self.init_tmux() self.init_tmux()
exe.start_thread(self.fix_tmux_layout_loop) exe.start_thread(self.fix_tmux_layout_loop)
def add_block_pair(self, source, destination, working_dir): def add_block_pair(self, source, destination):
"""Add BlockPair object and run safety checks.""" """Add BlockPair object and run safety checks."""
self.block_pairs.append( self.block_pairs.append(
BlockPair( BlockPair(
source=source, source=source,
destination=destination, destination=destination,
model=self.source.details['model'], model=self.source.details['model'],
working_dir=working_dir, working_dir=self.working_dir,
)) ))
def add_clone_block_pairs(self, working_dir): def add_clone_block_pairs(self):
"""Add device to device block pairs and set settings if necessary.""" """Add device to device block pairs and set settings if necessary."""
source_sep = get_partition_separator(self.source.path.name) source_sep = get_partition_separator(self.source.path.name)
dest_sep = get_partition_separator(self.destination.path.name) dest_sep = get_partition_separator(self.destination.path.name)
@ -284,7 +284,7 @@ class State():
source_parts = [] source_parts = []
# Clone settings # Clone settings
settings = self.load_settings(working_dir, discard_unused_settings=True) settings = self.load_settings(discard_unused_settings=True)
# Add pairs # Add pairs
if settings['Partition Mapping']: if settings['Partition Mapping']:
@ -296,13 +296,13 @@ class State():
bp_dest = pathlib.Path( bp_dest = pathlib.Path(
f'{self.destination.path}{dest_sep}{part_map[1]}', f'{self.destination.path}{dest_sep}{part_map[1]}',
) )
self.add_block_pair(bp_source, bp_dest, working_dir) self.add_block_pair(bp_source, bp_dest)
else: else:
source_parts = select_disk_parts('Clone', self.source) source_parts = select_disk_parts('Clone', self.source)
if self.source.path.samefile(source_parts[0].path): if self.source.path.samefile(source_parts[0].path):
# Whole disk (or single partition via args), skip settings # Whole disk (or single partition via args), skip settings
bp_dest = self.destination.path bp_dest = self.destination.path
self.add_block_pair(self.source, bp_dest, working_dir) self.add_block_pair(self.source, bp_dest)
else: else:
# New run, use new settings file # New run, use new settings file
settings['Needs Format'] = True settings['Needs Format'] = True
@ -328,26 +328,25 @@ class State():
bp_dest = pathlib.Path( bp_dest = pathlib.Path(
f'{self.destination.path}{dest_sep}{dest_num}', f'{self.destination.path}{dest_sep}{dest_num}',
) )
self.add_block_pair(part, bp_dest, working_dir) self.add_block_pair(part, bp_dest)
# Add to settings file # Add to settings file
source_num = re.sub(r'^.*?(\d+)$', r'\1', part.path.name) source_num = re.sub(r'^.*?(\d+)$', r'\1', part.path.name)
settings['Partition Mapping'].append([source_num, dest_num]) settings['Partition Mapping'].append([source_num, dest_num])
# Save settings # Save settings
self.save_settings(settings, working_dir) self.save_settings(settings)
# Done # Done
return source_parts return source_parts
def add_image_block_pairs(self, source_parts, working_dir): def add_image_block_pairs(self, source_parts):
"""Add device to image file block pairs.""" """Add device to image file block pairs."""
for part in source_parts: for part in source_parts:
bp_dest = self.destination bp_dest = self.destination
self.add_block_pair(part, bp_dest, working_dir) self.add_block_pair(part, bp_dest)
def confirm_selections( def confirm_selections(self, mode, prompt, source_parts=None):
self, mode, prompt, working_dir=None, source_parts=None):
"""Show selection details and prompt for confirmation.""" """Show selection details and prompt for confirmation."""
report = [] report = []
@ -384,17 +383,17 @@ class State():
report.extend( report.extend(
build_block_pair_report( build_block_pair_report(
self.block_pairs, self.block_pairs,
self.load_settings(working_dir) if mode == 'Clone' else {}, self.load_settings() if mode == 'Clone' else {},
), ),
) )
report.append(' ') report.append(' ')
# Map dir # Map dir
if working_dir: if self.working_dir:
report.append(std.color_string('Map Save Directory', 'GREEN')) report.append(std.color_string('Map Save Directory', 'GREEN'))
report.append(f'{working_dir}/') report.append(f'{self.working_dir}/')
report.append(' ') report.append(' ')
if not fstype_is_ok(working_dir, map_dir=True): if not fstype_is_ok(self.working_dir, map_dir=True):
report.append( report.append(
std.color_string( std.color_string(
'Map file(s) are being saved to a non-recommended filesystem.', 'Map file(s) are being saved to a non-recommended filesystem.',
@ -518,35 +517,33 @@ class State():
self.update_progress_pane('Idle') self.update_progress_pane('Idle')
# Set working dir # Set working dir
working_dir = get_working_dir( self.working_dir = get_working_dir(
mode, self.destination, force_local=docopt_args['--force-local-map'], mode, self.destination, force_local=docopt_args['--force-local-map'],
) )
# Start fresh if requested # Start fresh if requested
if docopt_args['--start-fresh']: if docopt_args['--start-fresh']:
clean_working_dir(working_dir) clean_working_dir(self.working_dir)
# Add block pairs # Add block pairs
if mode == 'Clone': if mode == 'Clone':
source_parts = self.add_clone_block_pairs(working_dir) source_parts = self.add_clone_block_pairs()
else: else:
source_parts = select_disk_parts(mode, self.source) source_parts = select_disk_parts(mode, self.source)
self.add_image_block_pairs(source_parts, working_dir) self.add_image_block_pairs(source_parts)
# Safety Checks #1 # Safety Checks #1
if mode == 'Clone': if mode == 'Clone':
self.safety_check_destination() self.safety_check_destination()
self.safety_check_size(mode, working_dir) self.safety_check_size(mode)
# Confirmation #2 # Confirmation #2
self.update_progress_pane('Idle') self.update_progress_pane('Idle')
self.confirm_selections(mode, 'Start recovery?', working_dir=working_dir) self.confirm_selections(mode, 'Start recovery?')
# Prep destination # Prep destination
if mode == 'Clone': if mode == 'Clone':
self.prep_destination( self.prep_destination(source_parts, dry_run=docopt_args['--dry-run'])
source_parts, working_dir, dry_run=docopt_args['--dry-run'],
)
# Safety Checks #2 # Safety Checks #2
if not docopt_args['--dry-run']: if not docopt_args['--dry-run']:
@ -579,11 +576,11 @@ class State():
# Source / Dest # Source / Dest
self.update_top_panes() self.update_top_panes()
def load_settings(self, working_dir, discard_unused_settings=False): def load_settings(self, discard_unused_settings=False):
"""Load settings from previous run, returns dict.""" """Load settings from previous run, returns dict."""
settings = {} settings = {}
settings_file = pathlib.Path( settings_file = pathlib.Path(
f'{working_dir}/Clone_{self.source.details["model"]}.json', f'{self.working_dir}/Clone_{self.source.details["model"]}.json',
) )
# Try loading JSON data # Try loading JSON data
@ -641,7 +638,7 @@ class State():
"""Check if all block_pairs completed pass_name, returns bool.""" """Check if all block_pairs completed pass_name, returns bool."""
return all([p.pass_complete(pass_name) for p in self.block_pairs]) return all([p.pass_complete(pass_name) for p in self.block_pairs])
def prep_destination(self, source_parts, working_dir, dry_run=True): def prep_destination(self, source_parts, dry_run=True):
"""Prep destination as necessary.""" """Prep destination as necessary."""
dest_prefix = str(self.destination.path) dest_prefix = str(self.destination.path)
dest_prefix += get_partition_separator(self.destination.path.name) dest_prefix += get_partition_separator(self.destination.path.name)
@ -649,7 +646,7 @@ class State():
msr_type = 'E3C9E316-0B5C-4DB8-817D-F92DF00215AE' msr_type = 'E3C9E316-0B5C-4DB8-817D-F92DF00215AE'
part_num = 0 part_num = 0
sfdisk_script = [] sfdisk_script = []
settings = self.load_settings(working_dir) settings = self.load_settings()
# Bail early # Bail early
if not settings['Needs Format']: if not settings['Needs Format']:
@ -710,7 +707,10 @@ class State():
) )
# Save sfdisk script # Save sfdisk script
script_path = f'{working_dir}/sfdisk_{self.destination.path.name}.script' script_path = (
f'{self.working_dir}/'
f'sfdisk_{self.destination.path.name}.script'
)
with open(script_path, 'w') as _f: with open(script_path, 'w') as _f:
_f.write('\n'.join(sfdisk_script)) _f.write('\n'.join(sfdisk_script))
@ -731,7 +731,7 @@ class State():
# Update settings # Update settings
settings['Needs Format'] = False settings['Needs Format'] = False
self.save_settings(settings, working_dir) self.save_settings(settings)
def retry_all_passes(self): def retry_all_passes(self):
"""Set all statuses to Pending.""" """Set all statuses to Pending."""
@ -750,10 +750,10 @@ class State():
raise std.GenericAbort() raise std.GenericAbort()
def safety_check_size(self, mode, working_dir): def safety_check_size(self, mode):
"""Run size safety check and abort if necessary.""" """Run size safety check and abort if necessary."""
required_size = sum([pair.size for pair in self.block_pairs]) required_size = sum([pair.size for pair in self.block_pairs])
settings = self.load_settings(working_dir) if mode == 'Clone' else {} settings = self.load_settings() if mode == 'Clone' else {}
# Increase required_size if necessary # Increase required_size if necessary
if mode == 'Clone' and settings.get('Needs Format', False): if mode == 'Clone' and settings.get('Needs Format', False):
@ -812,11 +812,11 @@ class State():
with open(f'{debug_dir}/bp_part#.report', 'a') as _f: with open(f'{debug_dir}/bp_part#.report', 'a') as _f:
_f.write('\n'.join(debug.generate_object_report(_bp))) _f.write('\n'.join(debug.generate_object_report(_bp)))
def save_settings(self, settings, working_dir): def save_settings(self, settings):
# pylint: disable=no-self-use # pylint: disable=no-self-use
"""Save settings for future runs.""" """Save settings for future runs."""
settings_file = pathlib.Path( settings_file = pathlib.Path(
f'{working_dir}/Clone_{self.source.details["model"]}.json', f'{self.working_dir}/Clone_{self.source.details["model"]}.json',
) )
# Try saving JSON data # Try saving JSON data
@ -1676,7 +1676,7 @@ def run_ddrescue(state, block_pair, pass_name, settings, dry_run=True):
# Wait a bit to let ddrescue exit safely # Wait a bit to let ddrescue exit safely
warning_message = 'Aborted' warning_message = 'Aborted'
std.sleep(2) std.sleep(2)
exe.run_program(['sudo', 'killall', 'ddrescue'], check=False) exe.run_program(['sudo', 'kill', str(proc.pid)], check=False)
break break
except subprocess.TimeoutExpired: except subprocess.TimeoutExpired:
# Continue to next loop to update panes # Continue to next loop to update panes