From c3ec690318a52adf04d37078f426fbeb4182a60e Mon Sep 17 00:00:00 2001
From: 2Shirt <2xShirt@gmail.com>
Date: Mon, 7 Mar 2022 23:49:15 -0700
Subject: [PATCH] Add new ddrescue argument options

Addresses #184
---
 scripts/wk/cfg/ddrescue.py |  4 ++++
 scripts/wk/hw/ddrescue.py  | 34 +++++++++++++++++++++++++++++-----
 2 files changed, 33 insertions(+), 5 deletions(-)

diff --git a/scripts/wk/cfg/ddrescue.py b/scripts/wk/cfg/ddrescue.py
index 738a1c8c..6e3459b8 100644
--- a/scripts/wk/cfg/ddrescue.py
+++ b/scripts/wk/cfg/ddrescue.py
@@ -23,15 +23,18 @@ AUTO_PASS_THRESHOLDS = {
 DDRESCUE_SETTINGS = {
   'Default': {
     '--binary-prefixes':  {'Selected': True,                'Hidden': True, },
+    '--complete-only':    {'Selected': True,                'Hidden': True, },
     '--data-preview':     {'Selected': True,  'Value': '5', 'Hidden': True, },
     '--idirect':          {'Selected': True,                                },
     '--odirect':          {'Selected': True,                                },
+    '--input-position':   {'Selected': False, 'Value': '0',                 },
     '--max-error-rate':   {'Selected': True,  'Value': '100MiB',            },
     '--max-read-rate':    {'Selected': False, 'Value': '1MiB',              },
     '--min-read-rate':    {'Selected': True,  'Value': '64KiB',             },
     '--reopen-on-error':  {'Selected': True,                                },
     '--retry-passes':     {'Selected': True,  'Value': '0',                 },
     '--reverse':          {'Selected': False,                               },
+    '--skip-size':        {'Selected': True,  'Value': '0.0001,0.01',       },  # Percentages of source size
     '--test-mode':        {'Selected': False, 'Value': 'test.map',          },
     '--timeout':          {'Selected': True,  'Value': '30m',               },
     '-vvvv':              {'Selected': True,                'Hidden': True, },
@@ -46,6 +49,7 @@ DDRESCUE_SETTINGS = {
     '--max-read-rate':    {'Selected': True,  'Value': '64MiB',             },
     '--min-read-rate':    {'Selected': True,  'Value': '1KiB',              },
     '--reopen-on-error':  {'Selected': True,                                },
+    '--skip-size':        {'Selected': True,  'Value': '0.001,0.05',        },  # Percentages of source size
     '--timeout':          {'Selected': False, 'Value': '30m',               },
     },
   }
diff --git a/scripts/wk/hw/ddrescue.py b/scripts/wk/hw/ddrescue.py
index 3e441127..391a4dcc 100644
--- a/scripts/wk/hw/ddrescue.py
+++ b/scripts/wk/hw/ddrescue.py
@@ -67,6 +67,7 @@ DDRESCUE_LOG_REGEX = re.compile(
   r'.*\(\s*(?P<percent>\d+\.?\d*)%\)$',
   re.IGNORECASE,
   )
+INITIAL_SKIP_MIN = 64 * 1024 # This is ddrescue's minimum accepted value
 REGEX_REMAINING_TIME = re.compile(
   r'remaining time:'
   r'\s*((?P<days>\d+)d)?'
@@ -1202,7 +1203,7 @@ def build_block_pair_report(block_pairs, settings):
   return report
 
 
-def build_ddrescue_cmd(block_pair, pass_name, settings):
+def build_ddrescue_cmd(block_pair, pass_name, settings_menu):
   """Build ddrescue cmd using passed details, returns list."""
   cmd = ['sudo', 'ddrescue']
   if (block_pair.destination.is_block_device()
@@ -1216,8 +1217,30 @@ def build_ddrescue_cmd(block_pair, pass_name, settings):
   elif pass_name == 'scrape':
     # Allow trimming and scraping
     pass
-  cmd.extend(settings)
-  cmd.append(f'--size={block_pair.size}')
+
+  # Fix domain size based on starting position
+  domain_size = block_pair.size
+  if settings_menu.options['--input-position']['Selected']:
+    settings_menu.options['--reverse']['Selected'] = False
+    input_position = std.string_to_bytes(
+      settings_menu.options['--input-position']['Value'],
+      )
+    domain_size -= input_position
+  cmd.append(f'--size={domain_size}')
+
+  # Determine skip sizes
+  if settings_menu.options['--skip-size']['Selected']:
+    skip_sizes = settings_menu.options['--skip-size']['Value'].split(',')
+    skip_sizes = [float(s) for s in skip_sizes]
+    initial_skip = min(INITIAL_SKIP_MIN, int(block_pair.size * skip_sizes[0]))
+    max_skip = min(int(block_pair.size * skip_sizes[1]), domain_size)
+    cmd.append(f'--skip-size={initial_skip},{max_skip}')
+  cmd.extend(get_ddrescue_settings(settings_menu))
+
+  # Add source physical sector size (if possible)
+  cmd.append(f'--sector-size={block_pair.source.details.get("phy-sec", 512)}')
+
+  # Add block pair and map file
   if PLATFORM == 'Darwin':
     # Use Raw disks if possible
     for dev in (block_pair.source, block_pair.destination):
@@ -1569,6 +1592,8 @@ def get_ddrescue_settings(settings_menu):
 
   # Check menu selections
   for name, details in settings_menu.options.items():
+    if name == '--skip-size':
+      continue
     if details['Selected']:
       if 'Value' in details:
         settings.append(f'{name}={details["Value"]}')
@@ -2130,7 +2155,6 @@ def run_recovery(state, main_menu, settings_menu, dry_run=True):
     if 'Retry' in name and details['Selected']:
       details['Selected'] = False
       state.retry_all_passes()
-  settings = get_ddrescue_settings(settings_menu)
 
   # Start SMART/Journal
   state.panes['SMART'] = tmux.split_window(
@@ -2158,7 +2182,7 @@ def run_recovery(state, main_menu, settings_menu, dry_run=True):
         attempted_recovery = True
         state.mark_started()
         try:
-          run_ddrescue(state, pair, pass_name, settings, dry_run=dry_run)
+          run_ddrescue(state, pair, pass_name, settings_menu, dry_run=dry_run)
         except (FileNotFoundError, KeyboardInterrupt, std.GenericAbort):
           is_missing_source_or_destination(state)
           abort = True