From 1a576c72a2080eeb0139cb3e6dbc0140ff371778 Mon Sep 17 00:00:00 2001
From: Jazzzny <75343012+Jazzzny@users.noreply.github.com>
Date: Fri, 9 Aug 2024 18:13:53 -0400
Subject: [PATCH] Provide additional resilience in USB detection code (#1144)

* Add fallback, don't bail out

* Part 2

* Part 3

* Fix import

* Move encoding
---
 opencore_legacy_patcher/support/install.py          | 13 ++++++++++---
 .../support/macos_installer_handler.py              | 11 +++++++++--
 2 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/opencore_legacy_patcher/support/install.py b/opencore_legacy_patcher/support/install.py
index e9bbf1cb7..f1a6aebf1 100644
--- a/opencore_legacy_patcher/support/install.py
+++ b/opencore_legacy_patcher/support/install.py
@@ -5,6 +5,7 @@ install.py: Installation of OpenCore files to ESP
 import logging
 import plistlib
 import subprocess
+import re
 
 from pathlib import Path
 
@@ -27,9 +28,15 @@ class tui_disk_installation:
             # Sierra and older
             disks = plistlib.loads(subprocess.run(["/usr/sbin/diskutil", "list", "-plist"], stdout=subprocess.PIPE).stdout.decode().strip().encode())
         for disk in disks["AllDisksAndPartitions"]:
-            disk_info = plistlib.loads(subprocess.run(["/usr/sbin/diskutil", "info", "-plist", disk["DeviceIdentifier"]], stdout=subprocess.PIPE).stdout.decode().strip().encode())
             try:
-                all_disks[disk["DeviceIdentifier"]] = {"identifier": disk_info["DeviceNode"], "name": disk_info["MediaName"], "size": disk_info["TotalSize"], "partitions": {}}
+                disk_info = plistlib.loads(subprocess.run(["/usr/sbin/diskutil", "info", "-plist", disk["DeviceIdentifier"]], stdout=subprocess.PIPE).stdout.decode().strip().encode())
+            except:
+                # Chinesium USB can have garbage data in MediaName
+                diskutil_output = subprocess.run(["/usr/sbin/diskutil", "info", "-plist", disk["DeviceIdentifier"]], stdout=subprocess.PIPE).stdout.decode().strip()
+                ungarbafied_output = re.sub(r'(<key>MediaName</key>\s*<string>).*?(</string>)', r'\1\2', diskutil_output).encode()
+                disk_info = plistlib.loads(ungarbafied_output)
+            try:
+                all_disks[disk["DeviceIdentifier"]] = {"identifier": disk_info["DeviceNode"], "name": disk_info.get("MediaName", "Disk"), "size": disk_info["TotalSize"], "partitions": {}}
                 for partition in disk["Partitions"]:
                     partition_info = plistlib.loads(subprocess.run(["/usr/sbin/diskutil", "info", "-plist", partition["DeviceIdentifier"]], stdout=subprocess.PIPE).stdout.decode().strip().encode())
                     all_disks[disk["DeviceIdentifier"]]["partitions"][partition["DeviceIdentifier"]] = {
@@ -98,7 +105,7 @@ class tui_disk_installation:
         partition_info = plistlib.loads(subprocess.run(["/usr/sbin/diskutil", "info", "-plist", full_disk_identifier], stdout=subprocess.PIPE).stdout.decode().strip().encode())
         parent_disk = partition_info["ParentWholeDisk"]
         drive_host_info = plistlib.loads(subprocess.run(["/usr/sbin/diskutil", "info", "-plist", parent_disk], stdout=subprocess.PIPE).stdout.decode().strip().encode())
-        sd_type = drive_host_info["MediaName"]
+        sd_type = drive_host_info.get("MediaName", "Disk")
         try:
             ssd_type = drive_host_info["SolidState"]
         except KeyError:
diff --git a/opencore_legacy_patcher/support/macos_installer_handler.py b/opencore_legacy_patcher/support/macos_installer_handler.py
index fc764c2d3..2475493d1 100644
--- a/opencore_legacy_patcher/support/macos_installer_handler.py
+++ b/opencore_legacy_patcher/support/macos_installer_handler.py
@@ -6,6 +6,7 @@ import logging
 import plistlib
 import tempfile
 import subprocess
+import re
 
 from pathlib import Path
 
@@ -171,9 +172,15 @@ fi
             disks = plistlib.loads(subprocess.run(["/usr/sbin/diskutil", "list", "-plist"], stdout=subprocess.PIPE).stdout.decode().strip().encode())
 
         for disk in disks["AllDisksAndPartitions"]:
-            disk_info = plistlib.loads(subprocess.run(["/usr/sbin/diskutil", "info", "-plist", disk["DeviceIdentifier"]], stdout=subprocess.PIPE).stdout.decode().strip().encode())
             try:
-                all_disks[disk["DeviceIdentifier"]] = {"identifier": disk_info["DeviceNode"], "name": disk_info["MediaName"], "size": disk_info["TotalSize"], "removable": disk_info["Internal"], "partitions": {}}
+                disk_info = plistlib.loads(subprocess.run(["/usr/sbin/diskutil", "info", "-plist", disk["DeviceIdentifier"]], stdout=subprocess.PIPE).stdout.decode().strip().encode())
+            except:
+                # Chinesium USB can have garbage data in MediaName
+                diskutil_output = subprocess.run(["/usr/sbin/diskutil", "info", "-plist", disk["DeviceIdentifier"]], stdout=subprocess.PIPE).stdout.decode().strip()
+                ungarbafied_output = re.sub(r'(<key>MediaName</key>\s*<string>).*?(</string>)', r'\1\2', diskutil_output).encode()
+                disk_info = plistlib.loads(ungarbafied_output)
+            try:
+                all_disks[disk["DeviceIdentifier"]] = {"identifier": disk_info["DeviceNode"], "name": disk_info.get("MediaName", "Disk"), "size": disk_info["TotalSize"], "removable": disk_info["Internal"], "partitions": {}}
             except KeyError:
                 # Avoid crashing with CDs installed
                 continue