From baaa897d104c3cbea2d1ae70449064cf18f72d2f Mon Sep 17 00:00:00 2001
From: Mykola Grymalyuk <48863253+khronokernel@users.noreply.github.com>
Date: Wed, 23 Jun 2021 21:46:42 -0600
Subject: [PATCH] Add disk type detection

---
 OCLP-CLI.command       |  3 ++-
 Resources/Build.py     | 25 +++++++++++++++++++++++++
 Resources/Constants.py |  1 +
 3 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/OCLP-CLI.command b/OCLP-CLI.command
index b93e8a103..c19297fad 100755
--- a/OCLP-CLI.command
+++ b/OCLP-CLI.command
@@ -72,7 +72,8 @@ class OpenCoreLegacyPatcher:
             print("- Using default payloads location")
 
         if args.disk:
-            print(f"Disk set: {args.disk}")
+            print(f"- Disk set: {args.disk}")
+            self.constants.disk = args.disk
         if args.verbose:
             print("- Set verbose configuration")
             self.constants.verbose_debug = True
diff --git a/Resources/Build.py b/Resources/Build.py
index 65c3db3c5..734d20396 100644
--- a/Resources/Build.py
+++ b/Resources/Build.py
@@ -98,6 +98,29 @@ class BuildOpenCore:
             fw_mask = b"\xff\x3f\x08\xc0\x00\x00\x00\x00"
         return fw_feature, fw_mask
 
+    def disk_type(self):
+        drive_host_info = plistlib.loads(subprocess.run(f"diskutil info -plist {self.constants.disk}".split(), stdout=subprocess.PIPE).stdout.decode().strip().encode())
+        sd_type = drive_host_info["MediaName"]
+        try:
+            ssd_type = drive_host_info["SolidState"]
+        except KeyError:
+            ssd_type = False
+        # Array filled with common SD Card names
+        # Note most USB-based SD Card readers generally report as "Storage Device", and no reliable way to detect further
+        if sd_type in ["SD Card Reader", "SD/MMC"]:
+            print("- Adding SD Card icon")
+            shutil.copy(self.constants.icon_path_sd, self.constants.opencore_release_folder)
+        elif ssd_type is True:
+            print("- Adding SSD icon")
+            shutil.copy(self.constants.icon_path_ssd, self.constants.opencore_release_folder)
+        elif drive_host_info["BusProtocol"] == "USB":
+            print("- Adding External USB Drive icon")
+            shutil.copy(self.constants.icon_path_external, self.constants.opencore_release_folder)
+        else:
+            print("- Adding Internal Drive icon")
+            shutil.copy(self.constants.icon_path_internal, self.constants.opencore_release_folder)
+
+
     def build_efi(self):
         Utilities.cls()
         if not self.constants.custom_model:
@@ -639,6 +662,8 @@ class BuildOpenCore:
         if self.model == self.constants.override_smbios:
             print("- Adding -no_compat_check")
             self.config["NVRAM"]["Add"]["7C436110-AB2A-4BBB-A880-FE41995C9F82"]["boot-args"] += " -no_compat_check"
+        if self.constants.disk != "":
+            self.disk_type()
 
     def set_smbios(self):
         spoofed_model = self.model
diff --git a/Resources/Constants.py b/Resources/Constants.py
index 954cdc569..63b63ed11 100644
--- a/Resources/Constants.py
+++ b/Resources/Constants.py
@@ -43,6 +43,7 @@ class Constants:
         self.innie_version = "1.3.0"
         self.fw_kext = "1.0.0"
         self.patcher_support_pkg_version = "0.0.10"  # PatcherSupportPkg
+        self.disk = ""
 
         # Get resource path
         self.current_path = Path(__file__).parent.parent.resolve()