install.py: refactor adding disk data requests

resources/install.py

@@ -14,23 +14,7 @@ class tui_disk_installation:
     def __init__(self, versions):
         self.constants: constants.Constants = versions
 
-    def copy_efi(self):
-        utilities.cls()
-        utilities.header(["Installing OpenCore to Drive"])
-
-        if not self.constants.opencore_release_folder.exists():
-            utilities.TUIOnlyPrint(
-                ["Installing OpenCore to Drive"],
-                "Press [Enter] to go back.\n",
-                [
-                    """OpenCore folder missing!
-Please build OpenCore first!"""
-                ],
-            ).start()
-            return
-
-        print("\nDisk picker is loading...")
-
+    def list_disks(self):
         all_disks = {}
         # TODO: AllDisksAndPartitions is not supported in Snow Leopard and older
         try:
@@ -54,7 +38,60 @@ Please build OpenCore first!"""
             except KeyError:
                 # Avoid crashing with CDs installed
                 continue
-        # TODO: Advanced mode
+
+        supported_disks = {}
+        for disk in all_disks:
+            if not any(all_disks[disk]["partitions"][partition]["fs"] in ("msdos", "EFI") for partition in all_disks[disk]["partitions"]):
+                continue
+            supported_disks.update({
+                disk: {
+                    "disk": disk,
+                    "name": all_disks[disk]["name"],
+                    "size": utilities.human_fmt(all_disks[disk]['size']),
+                    "partitions": all_disks[disk]["partitions"]
+                }
+            })
+        return supported_disks
+
+    def list_partitions(self, disk_response, supported_disks):
+        # Takes disk UUID as well as diskutil dataset generated by list_disks
+        # Returns list of FAT32 partitions
+        disk_identifier = disk_response
+        selected_disk = supported_disks[disk_identifier]
+
+        supported_partitions = {}
+        
+        for partition in selected_disk["partitions"]:
+            if selected_disk["partitions"][partition]["fs"] not in ("msdos", "EFI"):
+                continue
+            supported_partitions.update({
+                partition: {
+                    "partition": partition,
+                    "name": selected_disk["partitions"][partition]["name"],
+                    "size": utilities.human_fmt(selected_disk["partitions"][partition]["size"])
+                }
+            })
+        return supported_partitions
+
+
+    def copy_efi(self):
+        utilities.cls()
+        utilities.header(["Installing OpenCore to Drive"])
+
+        if not self.constants.opencore_release_folder.exists():
+            utilities.TUIOnlyPrint(
+                ["Installing OpenCore to Drive"],
+                "Press [Enter] to go back.\n",
+                [
+                    """OpenCore folder missing!
+Please build OpenCore first!"""
+                ],
+            ).start()
+            return
+
+        print("\nDisk picker is loading...")
+
+        all_disks = self.list_disks()
         menu = utilities.TUIMenu(
             ["Select Disk"],
             "Please select the disk you would like to install OpenCore to: ",
@@ -63,9 +100,7 @@ Please build OpenCore first!"""
             loop=True,
         )
         for disk in all_disks:
-            if not any(all_disks[disk]["partitions"][partition]["fs"] in ("msdos", "EFI") for partition in all_disks[disk]["partitions"]):
-                continue
-            menu.add_menu_option(f"{disk}: {all_disks[disk]['name']} ({utilities.human_fmt(all_disks[disk]['size'])})", key=disk[4:])
+            menu.add_menu_option(f"{disk}: {all_disks[disk]['name']} ({all_disks[disk]['size']})", key=disk[4:])
 
         response = menu.start()
 
@@ -96,9 +131,9 @@ Please build OpenCore first!"""
 
         if response == -1:
             return        
-        self.install_opencore(disk_identifier, response)
+        self.install_opencore(f"{disk_identifier}s{response}")
 
-    def install_opencore(self, disk_identifier, response):
+    def install_opencore(self, full_disk_identifier):
         def determine_sd_card(media_name):
             # Array filled with common SD Card names
             # Note most USB-based SD Card readers generally report as "Storage Device"
@@ -117,7 +152,7 @@ Please build OpenCore first!"""
         args = [
             "osascript",
             "-e",
-            f'''do shell script "diskutil mount {disk_identifier}s{response}"'''
+            f'''do shell script "diskutil mount {full_disk_identifier}"'''
             ' with prompt "OpenCore Legacy Patcher needs administrator privileges to mount your EFI."'
             " with administrator privileges"
             " without altering line endings",
@@ -126,7 +161,7 @@ Please build OpenCore first!"""
         if self.constants.detected_os >= os_data.os_data.el_capitan and not self.constants.recovery_status:
             result = subprocess.run(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
         else:
-            result = subprocess.run(f"diskutil mount {disk_identifier}s{response}".split(), stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+            result = subprocess.run(f"diskutil mount {full_disk_identifier}".split(), stdout=subprocess.PIPE, stderr=subprocess.PIPE)
 
         if result.returncode != 0:
             if "execution error" in result.stderr.decode() and result.stderr.decode().strip()[-5:-1] == "-128":
@@ -137,8 +172,9 @@ Please build OpenCore first!"""
                     ["Copying OpenCore"], "Press [Enter] to go back.\n", ["An error occurred!"] + result.stderr.decode().split("\n") + ["", "Please report this to the devs at GitHub."]
                 ).start()
                 return
-        drive_host_info = plistlib.loads(subprocess.run(f"diskutil info -plist {disk_identifier}".split(), stdout=subprocess.PIPE).stdout.decode().strip().encode())
-        partition_info = plistlib.loads(subprocess.run(f"diskutil info -plist {disk_identifier}s{response}".split(), stdout=subprocess.PIPE).stdout.decode().strip().encode())
+        partition_info = plistlib.loads(subprocess.run(f"diskutil info -plist {full_disk_identifier}".split(), stdout=subprocess.PIPE).stdout.decode().strip().encode())
+        parent_disk = partition_info["ParentWholeDisk"]
+        drive_host_info = plistlib.loads(subprocess.run(f"diskutil info -plist {parent_disk}".split(), stdout=subprocess.PIPE).stdout.decode().strip().encode())
         sd_type = drive_host_info["MediaName"]
         try:
             ssd_type = drive_host_info["SolidState"]
@@ -150,7 +186,7 @@ Please build OpenCore first!"""
         utilities.header(["Copying OpenCore"])
 
         if mount_path.exists():
-            if (mount_path / Path("EFI/Microsoft")).exists():
+            if (mount_path / Path("EFI/Microsoft")).exists() and self.constants.gui_mode is False:
                 print("- Found Windows Boot Loader")
                 print("\nWould you like to continue installing OpenCore?")
                 print("Installing OpenCore onto this drive may make Windows unbootable until OpenCore")
@@ -202,8 +238,9 @@ Please build OpenCore first!"""
                 print("- Unmounting EFI partition")
                 subprocess.run(["diskutil", "umount", mount_path], stdout=subprocess.PIPE).stdout.decode().strip().encode()
             print("- OpenCore transfer complete")
-            print("\nPress [Enter] to continue.\n")
-            input()
+            if self.constants.gui_mode is False:
+                print("\nPress [Enter] to continue.\n")
+                input()
         else:
             utilities.TUIOnlyPrint(["Copying OpenCore"], "Press [Enter] to go back.\n", ["EFI failed to mount!", "Please report this to the devs at GitHub."]).start()