From c308bcb993d8d4bbbf48d707658b8ccc44d7f25d Mon Sep 17 00:00:00 2001
From: Mykola Grymalyuk <khronokernel@icloud.com>
Date: Thu, 4 May 2023 12:43:06 -0600
Subject: [PATCH] payloads: USe variable name for mounting

---
 resources/constants.py        | 10 +++++++--
 resources/reroute_payloads.py | 40 ++++++++++++++++++-----------------
 2 files changed, 29 insertions(+), 21 deletions(-)

diff --git a/resources/constants.py b/resources/constants.py
index 23cec43fe..131b851fd 100644
--- a/resources/constants.py
+++ b/resources/constants.py
@@ -106,8 +106,9 @@ class Constants:
         self.kdkless_version: str = "1.0.0"
 
         # Get resource path
-        self.current_path: Path = Path(__file__).parent.parent.resolve()
-        self.payload_path: Path = self.current_path / Path("payloads")
+        self.current_path:  Path = Path(__file__).parent.parent.resolve()
+        self.original_path: Path = Path(__file__).parent.parent.resolve()
+        self.payload_path:  Path = self.current_path / Path("payloads")
 
         # Patcher Settings
         ## Internal settings
@@ -225,6 +226,11 @@ class Constants:
         ]
 
     # Payload Location
+
+    # Support Disk Images
+    @property
+    def payload_path_dmg(self):
+        return self.original_path / Path("payloads.dmg")
     # OpenCore
     @property
     def opencore_zip_source(self):
diff --git a/resources/reroute_payloads.py b/resources/reroute_payloads.py
index 9877e341f..760e359fb 100644
--- a/resources/reroute_payloads.py
+++ b/resources/reroute_payloads.py
@@ -36,7 +36,7 @@ class RoutePayloadDiskImage:
             self._unmount_active_dmgs(unmount_all_active=False)
             output = subprocess.run(
                 [
-                    "hdiutil", "attach", "-noverify", f"{self.constants.payload_path}.dmg",
+                    "hdiutil", "attach", "-noverify", f"{self.constants.payload_path_dmg}",
                     "-mountpoint", Path(self.temp_dir.name / Path("payloads")),
                     "-nobrowse",
                     "-shadow", Path(self.temp_dir.name / Path("payloads_overlay")),
@@ -55,7 +55,7 @@ class RoutePayloadDiskImage:
                 logging.info(f"Return Code: {output.returncode}")
 
 
-    def _unmount_active_dmgs(self, unmount_all_active=True) -> None:
+    def _unmount_active_dmgs(self, unmount_all_active: bool = True) -> None:
         """
         Unmounts disk images associated with OCLP
 
@@ -70,20 +70,22 @@ class RoutePayloadDiskImage:
         dmg_info = subprocess.run(["hdiutil", "info", "-plist"], stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
         dmg_info = plistlib.loads(dmg_info.stdout)
 
-        for image in dmg_info["images"]:
-            if image["image-path"].endswith("payloads.dmg"):
-                if unmount_all_active is False:
-                    # Check that only our personal payloads.dmg is unmounted
-                    if "shadow-path" in image:
-                        if self.temp_dir.name in image["shadow-path"]:
-                            logging.info("- Unmounting personal payloads.dmg")
-                            subprocess.run(
-                                ["hdiutil", "detach", image["system-entities"][0]["dev-entry"], "-force"],
-                                stdout=subprocess.PIPE, stderr=subprocess.STDOUT
-                            )
-                else:
-                    logging.info(f"- Unmounting payloads.dmg at: {image['system-entities'][0]['dev-entry']}")
-                    subprocess.run(
-                        ["hdiutil", "detach", image["system-entities"][0]["dev-entry"], "-force"],
-                        stdout=subprocess.PIPE, stderr=subprocess.STDOUT
-                    )
\ No newline at end of file
+
+        for variant in ["Universal-Binaries.dmg", "payloads.dmg"]:
+            for image in dmg_info["images"]:
+                if image["image-path"].endswith(variant):
+                    if unmount_all_active is False:
+                        # Check that only our personal payloads.dmg is unmounted
+                        if "shadow-path" in image:
+                            if self.temp_dir.name in image["shadow-path"]:
+                                logging.info(f"- Unmounting personal {variant}")
+                                subprocess.run(
+                                    ["hdiutil", "detach", image["system-entities"][0]["dev-entry"], "-force"],
+                                    stdout=subprocess.PIPE, stderr=subprocess.STDOUT
+                                )
+                    else:
+                        logging.info(f"- Unmounting {variant} at: {image['system-entities'][0]['dev-entry']}")
+                        subprocess.run(
+                            ["hdiutil", "detach", image["system-entities"][0]["dev-entry"], "-force"],
+                            stdout=subprocess.PIPE, stderr=subprocess.STDOUT
+                        )
\ No newline at end of file