From f2adf6f5c003cb4f134d80fec5f8f404a4fb7372 Mon Sep 17 00:00:00 2001
From: 2Shirt <2xShirt@gmail.com>
Date: Thu, 16 May 2019 18:23:57 -0600
Subject: [PATCH] Improved ddrescue map parsing

* Get data recovered instead of percentage, should be more acurate
---
 .bin/Scripts/functions/ddrescue.py | 42 ++++++++++++++++--------------
 .bin/Scripts/settings/ddrescue.py  |  6 +++++
 2 files changed, 28 insertions(+), 20 deletions(-)

diff --git a/.bin/Scripts/functions/ddrescue.py b/.bin/Scripts/functions/ddrescue.py
index 79a58960..31122edc 100644
--- a/.bin/Scripts/functions/ddrescue.py
+++ b/.bin/Scripts/functions/ddrescue.py
@@ -107,19 +107,19 @@ class BlockPair():
   def load_map_data(self):
     """Load data from map file and set progress."""
     map_data = read_map_file(self.map_path)
-    self.rescued_percent = map_data['rescued']
-    self.rescued = (self.rescued_percent * self.size) / 100
+    self.rescued = map_data.get('rescued', 0)
+    self.rescued_percent = (self.rescued / self.size) * 100
     if map_data['full recovery']:
       self.pass_done = [True, True, True]
       self.rescued = self.size
       self.status = ['Skipped', 'Skipped', 'Skipped']
-    elif map_data['non-tried'] > 0:
+    elif map_data.get('non-tried', 0) > 0:
       # Initial pass incomplete
       pass
-    elif map_data['non-trimmed'] > 0:
+    elif map_data.get('non-trimmed', 0) > 0:
       self.pass_done = [True, False, False]
       self.status = ['Skipped', 'Pending', 'Pending']
-    elif map_data['non-scraped'] > 0:
+    elif map_data.get('non-scraped', 0) > 0:
       self.pass_done = [True, True, False]
       self.status = ['Skipped', 'Skipped', 'Pending']
     else:
@@ -147,8 +147,8 @@ class BlockPair():
     """Update progress using map file."""
     if os.path.exists(self.map_path):
       map_data = read_map_file(self.map_path)
-      self.rescued_percent = map_data.get('rescued', 0)
-      self.rescued = (self.rescued_percent * self.size) / 100
+      self.rescued = map_data.get('rescued', 0)
+      self.rescued_percent = (self.rescued / self.size) * 100
       self.status[pass_num] = get_formatted_status(
         label='Pass {}'.format(pass_num+1),
         data=(self.rescued/self.size)*100)
@@ -1022,25 +1022,27 @@ def menu_settings(state):
 
 def read_map_file(map_path):
   """Read map file with ddrescuelog and return data as dict."""
-  map_data = {'full recovery': False}
+  cmd = [
+    'ddrescuelog',
+    '--binary-prefixes',
+    '--show-status',
+    map_path,
+    ]
+  map_data = {'full recovery': False, 'pass completed': False}
   try:
-    result = run_program(['ddrescuelog', '-t', map_path])
+    result = run_program(cmd, encoding='utf-8', errors='ignore')
   except CalledProcessError:
     # (Grossly) assuming map_data hasn't been saved yet, return empty dict
     return map_data
 
   # Parse output
-  for line in result.stdout.decode().splitlines():
-    data = re.match(
-      r'^\s*(?P<key>\S+):.*\(\s*(?P<value>\d+\.?\d*)%.*', line.strip())
-    if data:
-      try:
-        map_data[data.group('key')] = float(data.group('value'))
-      except ValueError:
-        raise GenericError('Failed to read map data')
-    data = re.match(r'.*current status:\s+(?P<status>.*)', line.strip())
-    if data:
-      map_data['pass completed'] = bool(data.group('status') == 'finished')
+  for line in result.stdout.splitlines():
+    line = line.strip()
+    _r = REGEX_DDRESCUE_LOG.search(line)
+    if _r:
+      map_data[_r.group('key')] = convert_to_bytes('{size} {unit}B'.format(
+        **_r.groupdict()))
+    map_data['pass completed'] = 'current status: finished' in line
 
   # Check if 100% done
   try:
diff --git a/.bin/Scripts/settings/ddrescue.py b/.bin/Scripts/settings/ddrescue.py
index 559fdbe4..bc9bef93 100644
--- a/.bin/Scripts/settings/ddrescue.py
+++ b/.bin/Scripts/settings/ddrescue.py
@@ -38,6 +38,12 @@ DDRESCUE_SETTINGS = {
   '-vvvv':              {'Enabled': True,                 'Hidden': True, },
   }
 ETOC_REFRESH_RATE = 30 # in seconds
+REGEX_DDRESCUE_LOG = re.compile(
+  r'^\s*(?P<key>\S+):\s+'
+  r'(?P<size>\d+)\s+'
+  r'(?P<unit>[PTGMKB])i?B?',
+  re.IGNORECASE,
+  )
 REGEX_REMAINING_TIME = re.compile(
   r'remaining time:'
   r'\s*((?P<days>\d+)d)?'