From 7c16d13f6592f81e614fbb6d1846448d342b7b26 Mon Sep 17 00:00:00 2001
From: 2Shirt <2xShirt@gmail.com>
Date: Sun, 9 Jul 2023 13:58:59 -0700
Subject: [PATCH] Generate test maps at runtime in ddrescue-tui

---
 scripts/wk/cfg/ddrescue.py   |  2 +-
 scripts/wk/clone/ddrescue.py | 36 +++++++++++++++++++++++++++++++++++-
 2 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/scripts/wk/cfg/ddrescue.py b/scripts/wk/cfg/ddrescue.py
index 5a463ea9..d8e9b8ac 100644
--- a/scripts/wk/cfg/ddrescue.py
+++ b/scripts/wk/cfg/ddrescue.py
@@ -37,7 +37,7 @@ DDRESCUE_SETTINGS = {
     '--retry-passes':     {'Selected': True,  'Value': '0',                 },
     '--reverse':          {'Selected': False,                               },
     '--skip-size':        {'Selected': True,  'Value': '0.001,0.02',        },  # Percentages of source size
-    '--test-mode':        {'Selected': False, 'Value': 'test.map',          },
+    '--test-mode':        {'Selected': False,                               },
     '--timeout':          {'Selected': True,  'Value': '30m',               },
     '-vvvv':              {'Selected': True,                'Hidden': True, },
     },
diff --git a/scripts/wk/clone/ddrescue.py b/scripts/wk/clone/ddrescue.py
index bb0d4f5b..7fe78873 100644
--- a/scripts/wk/clone/ddrescue.py
+++ b/scripts/wk/clone/ddrescue.py
@@ -14,6 +14,7 @@ import shutil
 import subprocess
 import time
 
+from random import randint
 from typing import Any
 
 import psutil
@@ -125,6 +126,7 @@ class BlockPair():
       'trim':       'Pending',
       'scrape':     'Pending',
       }
+    self.test_map: pathlib.Path | None = None
     self.view_map: bool = 'DISPLAY' in os.environ or 'WAYLAND_DISPLAY' in os.environ
     self.view_proc: subprocess.Popen | None = None
 
@@ -1196,6 +1198,16 @@ def build_ddrescue_cmd(block_pair, pass_name, settings_menu) -> list[str]:
   # Add source physical sector size (if possible)
   cmd.append(f'--sector-size={block_pair.sector_size}')
 
+  # Generate test map if needed
+  if '--test-mode' in cmd:
+    cmd.remove('--test-mode')
+    if not block_pair.test_map:
+      block_pair.test_map = block_pair.map_path.with_stem(
+        f'{block_pair.map_path.stem}-testing'
+      )
+      generate_test_map(map_path=block_pair.test_map, size=domain_size)
+    cmd.append(f'--test-mode={block_pair.test_map}')
+
   # Add block pair and map file
   if PLATFORM == 'Darwin':
     # Use Raw disks if possible
@@ -1484,6 +1496,28 @@ def fstype_is_ok(path, map_dir=False) -> bool:
   return is_ok
 
 
+def generate_test_map(map_path: pathlib.Path, size: int) -> None:
+  """Generate test map with roughly 20% of the space marked as bad."""
+  chunk = 2*1024**2
+  output = [
+    '# Mapfile. Created by WizardKit',
+    '0x0 ? 1',
+    ]
+  position = 0
+
+  # Generate "holes"
+  steps, remainder = divmod(size, chunk)
+  for _ in range(steps):
+    bad = randint(1, 5) % 5 == 0
+    output.append(f'{hex(position)}  {hex(chunk)}  {"-" if bad else "+"}')
+    position += chunk
+  if remainder:
+    output.append(f'{hex(position)}  {hex(remainder)}  +')
+
+  # Save map
+  map_path.write_text('\n'.join(output), encoding='utf-8')
+
+
 def get_ddrescue_settings(settings_menu) -> list:
   """Get ddrescue settings from menu selections, returns list."""
   settings = []
@@ -2238,7 +2272,7 @@ def zero_fill_destination(state: State, dry_run: bool = True) -> None:
 
     # Prep zero-fill map file
     zero_map_path = block_pair.map_path.with_stem(
-      f'{block_pair.map_path.name}_zero-fill',
+      f'{block_pair.map_path.stem}_zero-fill',
       )
     io.copy_file(block_pair.map_path, zero_map_path, overwrite=True)
     if larger_destination: