From 544ffb1aff144f4dded45a75a033ec3e1b2ce8b3 Mon Sep 17 00:00:00 2001
From: 2Shirt <2xShirt@gmail.com>
Date: Fri, 20 May 2022 17:58:07 -0700
Subject: [PATCH] Refactor partition table type lookup sections

---
 scripts/wk/clone/ddrescue.py | 30 ++++++++++++++++++++++++++----
 1 file changed, 26 insertions(+), 4 deletions(-)

diff --git a/scripts/wk/clone/ddrescue.py b/scripts/wk/clone/ddrescue.py
index b46cadaa..93724443 100644
--- a/scripts/wk/clone/ddrescue.py
+++ b/scripts/wk/clone/ddrescue.py
@@ -526,7 +526,7 @@ class State():
           settings['Table Type'] = 'MBR'
         else:
           # Match source type
-          settings['Table Type'] = get_table_type(self.source)
+          settings['Table Type'] = get_table_type(self.source.path)
         if std.ask('Create an empty Windows boot partition on the clone?'):
           settings['Create Boot Partition'] = True
           offset = 2 if settings['Table Type'] == 'GPT' else 1
@@ -1728,14 +1728,36 @@ def get_percent_color(percent):
   return color
 
 
-def get_table_type(disk):
+def get_table_type(disk_path):
   """Get disk partition table type, returns str.
 
   NOTE: If resulting table type is not GPT or MBR
         then an exception is raised.
   """
-  table_type = str(disk.raw_details.get('pttype', '')).upper()
-  table_type = table_type.replace('DOS', 'MBR')
+  disk_path = str(disk_path)
+  table_type = None
+
+  # Linux
+  if std.PLATFORM == 'Linux':
+    cmd = f'lsblk --json --output=pttype --nodeps {disk_path}'.split()
+    json_data = exe.get_json_from_command(cmd)
+    table_type = json_data['blockdevices'][0].get('pttype', '').upper()
+    table_type = table_type.replace('DOS', 'MBR')
+
+  # macOS
+  if std.PLATFORM == 'Darwin':
+    cmd = ['diskutil', 'list', '-plist', disk_path]
+    proc = exe.run_program(cmd, check=False, encoding=None, errors=None)
+    try:
+      plist_data = plistlib.loads(proc.stdout)
+    except (TypeError, ValueError):
+      # Invalid / corrupt plist data? return empty dict to avoid crash
+      pass
+    else:
+      disk_details = plist_data.get('AllDisksAndPartitions', [{}])[0]
+      table_type = disk_details['Content']
+      table_type = table_type.replace('FDisk_partition_scheme', 'MBR')
+      table_type = table_type.replace('GUID_partition_scheme', 'GPT')
 
   # Check type
   if table_type not in ('GPT', 'MBR'):