From ef426c95b503517407e410c218484326f1e8b315 Mon Sep 17 00:00:00 2001
From: 2Shirt <twoshirt@2Shirts-MacBook-Pro.local>
Date: Thu, 18 Mar 2021 21:29:58 -0600
Subject: [PATCH] Fix get_core_storage_volumes()

---
 scripts/wk/os/mac.py | 40 ++++++++++++++++++++++++----------------
 1 file changed, 24 insertions(+), 16 deletions(-)

diff --git a/scripts/wk/os/mac.py b/scripts/wk/os/mac.py
index 31b1f5de..d5735005 100644
--- a/scripts/wk/os/mac.py
+++ b/scripts/wk/os/mac.py
@@ -64,6 +64,7 @@ def get_apfs_volumes(device_path):
 
 
 def get_core_storage_volumes(device_path):
+  # pylint: disable=too-many-branches
   """Get CoreStorage volumes contained in device_path, returns list."""
   disks = []
   volumes = []
@@ -85,12 +86,13 @@ def get_core_storage_volumes(device_path):
     for p_v in l_vg['CoreStoragePhysicalVolumes']:
       uuid = p_v['CoreStorageUUID']
       cmd = ['diskutil', 'coreStorage', 'info', '-plist', uuid]
+      proc = run_program(cmd, check=False, encoding=None, errors=None)
       try:
-        uuid_data = plistlib.loads(proc.stdout)
+        plist_data = plistlib.loads(proc.stdout)
       except (TypeError, ValueError):
         LOG.error('failed to get diskutil corestorage info for %s', uuid)
         continue
-      if uuid_data['DeviceIdentifier'] == device_path:
+      if plist_data['DeviceIdentifier'] == device_path:
         related = True
         break
 
@@ -99,27 +101,33 @@ def get_core_storage_volumes(device_path):
       continue
 
     # Add logical disks to list
-    for l_v in l_vg['CoreStorageLogicalVolumes']:
-      uuid = l_v['CoreStorageUUID']
-      cmd = ['diskutil', 'coreStorage', 'info', '-plist', uuid]
-      try:
-        uuid_data = plistlib.loads(proc.stdout)
-      except (TypeError, ValueError):
-        LOG.error('failed to get diskutil corestorage info for %s', uuid)
-        continue
-      disks.append(uuid_data['DeviceIdentifier'])
+    for l_vf in l_vg['CoreStorageLogicalVolumeFamilies']:
+      for l_v in l_vf['CoreStorageLogicalVolumes']:
+        uuid = l_v['CoreStorageUUID']
+        cmd = ['diskutil', 'coreStorage', 'info', '-plist', uuid]
+        proc = run_program(cmd, check=False, encoding=None, errors=None)
+        try:
+          plist_data = plistlib.loads(proc.stdout)
+        except (TypeError, ValueError):
+          LOG.error('failed to get diskutil corestorage info for %s', uuid)
+          continue
+        disks.append(plist_data['DeviceIdentifier'])
 
   # Get volumes from logical disks
   for disk in disks:
     cmd = ['diskutil', 'list', '-plist', f'/dev/{disk}']
     proc = run_program(cmd, check=False, encoding=None, errors=None)
     try:
-      disk_data = plistlib.loads(proc.stdout)
+      plist_data = plistlib.loads(proc.stdout)
+      disk_data = plist_data['AllDisksAndPartitions'][0]
     except (TypeError, ValueError):
       LOG.error('Failed to get diskutil list for %s', disk)
       continue
-    for part in disk_data['AllDisksAndPartitions'][0]['Partitions']:
-      volumes.append(Disk(f'/dev/{part["DeviceIdentifier"]}'))
+    if 'Partitions' in disk_data:
+      for part in disk_data['Partitions']:
+        volumes.append(Disk(f'/dev/{part["DeviceIdentifier"]}'))
+    elif 'VolumeName' in disk_data:
+      volumes.append(Disk(f'/dev/{disk_data["DeviceIdentifier"]}'))
 
   # Done
   return volumes
@@ -186,13 +194,13 @@ def mount_volumes(device_path=None, read_write=False):
 
     # Volumes
     result += f'{"Mounted on "+vol.details.get("mountpoint", "?"):<40}'
-    size = vol.details.get('VolumeSize', -1)
+    size = vol.details.get('TotalSize', -1)
     free = vol.details.get('FreeSpace', -1)
     used = size - free
     result = (
       f'{result} ({vol.details.get("fstype", "Unknown FS")+",":<5} '
       f'{std.bytes_to_string(used, decimals=1):>9} used, '
-      f'{std.bytes_to_string(size, decimals=1):>9} size, '
+      f'{std.bytes_to_string(free, decimals=1):>9} free) '
       )
     report.append(result)