Add an option to rebuild vendorfw bundle

Signed-off-by: Sasha Finkelstein <[email protected]>

Author: WhatAmISupposedToPutHere

src/main.py

@@ -437,6 +437,74 @@ def action_repair_or_upgrade(self, oses, upgrade):
         # Go for step2 again
         self.step2()
 
+    def action_rebuild_vendorfw(self, oses):
+        choices = {str(i): f"{p.desc}\n      {str(o)}" for i, (p, o) in enumerate(oses)}
+
+        if len(choices) > 1:
+            print()
+            p_question("Choose an install to rebuild firmware for:")
+            idx = self.choice("Installed OS", choices)
+        else:
+            idx = list(choices.keys())[0]
+
+        self.part, osi = oses[int(idx)]
+        self.ins = stub.StubInstaller(self.sysinfo, self.dutil, self.osinfo)
+        if not self.ins.check_existing_install(osi):
+            p_error(   "The existing installation is missing files.")
+            p_message(f"This tool can only rebuild firmware on installations that completed the first")
+            p_message( "stage of the installation process. If it was interrupted, please")
+            p_message( "delete the partitions manually and reinstall from scratch.")
+            return False
+
+        vars = m1n1.extract_vars(self.ins.boot_obj_path)
+        if vars is None:
+            p_error("Could not get variables from the installed m1n1")
+            p_message(f"Path: {self.ins.boot_obj_path}")
+            return False
+
+        esp_id = None
+        for var in vars:
+            k, v = var.split('=')
+            if k == "chosen.asahi,efi-system-partition":
+                esp_id = v
+                break
+        if esp_id is None:
+            p_error("Unable to determine the associated EFI system partition for this install.")
+            p_message("Is first stage m1n1 corrupted?")
+            return False
+
+        target = self.dutil.get_partition_info(esp_id)
+        mountpoint = self.dutil.mount(target.name)
+
+        os.makedirs("vendorfw", exist_ok=True)
+        fw_pkg = asahi_firmware.core.FWPackage("vendorfw")
+        ipsw = None
+        for ver in IPSW_VERSIONS:
+            if ver.version == osi.version:
+                ipsw = ver
+                break
+        if ipsw is None:
+            p_error(f"This install is using an unsupported macOS version {osi.version}")
+            p_message("Unable to rebuild firmware")
+            return False
+
+        self.ins.load_ipsw(ipsw)
+        self.ins.load_identity()
+        self.ins.collect_firmware(fw_pkg)
+        fw_pkg.close()
+
+        p_plain(f"  Copying firmware into {target.name} partition...")
+        base = os.path.join(mountpoint, "vendorfw")
+        logging.info(f"Firmware -> {base}")
+        shutil.copytree(fw_pkg.path, base, dirs_exist_ok=True)
+
+        print()
+        p_success(f"Firmware rebuild complete. Press enter to continue.")
+        self.input()
+        print()
+
+        return True
+
     def do_install(self, total_size=None):
         p_progress(f"Installing stub macOS into {self.part.name} ({self.part.label})")
 
@@ -907,6 +975,7 @@ def main_loop(self):
         parts_resizable = []
         oses_incomplete = []
         oses_upgradable = []
+        oses_vendorfw = []
 
         for i, p in enumerate(self.parts):
             p.index = i
@@ -981,6 +1050,8 @@ def main_loop(self):
                 p_plain(f"    OS: [{state}] {os}")
                 if os.attached_partitions:
                     p_plain(f"      Extra partitions: {' '.join('#' + str(i.index) for i in os.attached_partitions)}")
+                if os.stub and os.m1n1_ver:
+                    oses_vendorfw.append((p, os))
                 if os.stub and os.m1n1_ver and os.m1n1_ver != self.m1n1_ver:
                     oses_upgradable.append((p, os))
                 elif os.stub and not (os.bp and os.bp.get("coih", None)):
@@ -1018,6 +1089,9 @@ def main_loop(self):
         if oses_upgradable:
             actions["m"] = "Upgrade m1n1 on an existing OS"
             default = default or "m"
+        if oses_vendorfw:
+            actions["v"] = "Rebuild vendor firmware package"
+            default = default or "v"
 
         if not actions:
             p_error("No actions available on this system.")
@@ -1046,6 +1120,8 @@ def main_loop(self):
             return self.action_select_disk()
         elif act == "w":
             return self.action_wipe()
+        elif act == "v":
+            return self.action_rebuild_vendorfw(oses_vendorfw)
         elif act == "q":
             return False
 

src/stub.py

@@ -209,26 +209,20 @@ def repair(self, cur_os):
                 else:
                     logging.info("'asahi' dir not found in ESP")
 
-    def install_files(self, cur_os):
-        logging.info("StubInstaller.install_files()")
-        logging.info(f"VGID: {self.osi.vgid}")
-        logging.info(f"OS info: {self.osi}")
-
-        p_progress("Beginning stub OS install...")
+    def load_identity(self):
         self.get_paths()
 
         logging.info("Parsing metadata...")
 
-        sysver = plistlib.load(self.open("SystemVersion.plist"))
         manifest = plistlib.load(self.open("BuildManifest.plist"))
-        bootcaches = plistlib.load(self.open("usr/standalone/bootcaches.plist"))
+        self.bootcaches = plistlib.load(self.open("usr/standalone/bootcaches.plist"))
         self.flush_progress()
 
         if self.is_ota:
-            variant = "macOS Customer Software Update"
+            self.variant = "macOS Customer Software Update"
             behavior = "Update"
         else:
-            variant = "macOS Customer"
+            self.variant = "macOS Customer"
             behavior = "Erase"
 
         self.manifest = manifest
@@ -237,7 +231,7 @@ def install_files(self, cur_os):
                 identity["ApChipID"] != f'0x{self.sysinfo.chip_id:04X}' or
                 identity["Info"]["DeviceClass"] != self.sysinfo.device_class or
                 identity["Info"]["RestoreBehavior"] != behavior or
-                identity["Info"]["Variant"] != variant):
+                identity["Info"]["Variant"] != self.variant):
                 continue
             break
         else:
@@ -248,6 +242,16 @@ def install_files(self, cur_os):
         self.all_identities = manifest["BuildIdentities"]
         self.identity = identity
         manifest["BuildIdentities"] = [identity]
+        return identity
+
+    def install_files(self, cur_os):
+        logging.info("StubInstaller.install_files()")
+        logging.info(f"VGID: {self.osi.vgid}")
+        logging.info(f"OS info: {self.osi}")
+
+        p_progress("Beginning stub OS install...")
+        identity = self.load_identity()
+        sysver = plistlib.load(self.open("SystemVersion.plist"))
 
         self.stub_info.update({
             "vgid": self.osi.vgid,
@@ -292,13 +296,13 @@ def install_files(self, cur_os):
 
         os.makedirs(self.pb_vgid, exist_ok=True)
 
-        bless2 = bootcaches["bless2"]
+        bless2 = self.bootcaches["bless2"]
 
         restore_bundle = os.path.join(self.pb_vgid, bless2["RestoreBundlePath"])
         os.makedirs(restore_bundle, exist_ok=True)
         restore_manifest = os.path.join(restore_bundle, "BuildManifest.plist")
         with open(restore_manifest, "wb") as fd:
-            plistlib.dump(manifest, fd)
+            plistlib.dump(self.manifest, fd)
         self.copy_idata.append((restore_manifest, "BuildManifest.plist"))
         self.extract("SystemVersion.plist", restore_bundle)
         self.extract("RestoreVersion.plist", restore_bundle)
@@ -311,10 +315,9 @@ def install_files(self, cur_os):
         self.extract_file("BootabilityBundle/Restore/Firmware/Bootability.dmg.trustcache",
                           os.path.join(restore_bundle, "Bootability/Bootability.trustcache"))
 
-        self.extract_tree(f"Firmware/Manifests/restore/{variant}/", restore_bundle)
+        self.extract_tree(f"Firmware/Manifests/restore/{self.variant}/", restore_bundle)
 
         copied = set()
-        self.kernel_path = None
         for key, val in identity["Manifest"].items():
             if key in ("BaseSystem", "OS", "Ap,SystemVolumeCanonicalMetadata",
                        "RestoreRamDisk", "RestoreTrustCache"):
@@ -328,8 +331,6 @@ def install_files(self, cur_os):
             if path.startswith("kernelcache."):
                 name = os.path.basename(path)
                 self.copy_idata.append((os.path.join(restore_bundle, name), name))
-                if self.kernel_path is None:
-                    self.kernel_path = os.path.join(restore_bundle, name)
             copied.add(path)
 
         self.flush_progress()
@@ -404,7 +405,10 @@ def collect_firmware(self, pkg):
             shutil.rmtree("fud_firmware")
 
         os.makedirs("fud_firmware", exist_ok=True)
+        bless2 = self.bootcaches["bless2"]
+        restore_bundle = os.path.join(self.pb_vgid, bless2["RestoreBundlePath"])
         copied = set()
+        kernel_path = None
         for identity in [self.identity]:
             if (identity["Info"]["RestoreBehavior"] != "Erase" or
                 identity["Info"]["Variant"] != "macOS Customer"):
@@ -419,6 +423,10 @@ def collect_firmware(self, pkg):
                            "StaticTrustCache", "SystemVolume"):
                     continue
                 path = val["Info"]["Path"]
+                if path.startswith("kernelcache."):
+                    name = os.path.basename(path)
+                    if kernel_path is None:
+                        kernel_path = os.path.join(restore_bundle, name)
                 if (not val["Info"].get("IsFUDFirmware", False)
                     or val["Info"].get("IsLoadedByiBoot", False)
                     or val["Info"].get("IsLoadedByiBootStage1", False)
@@ -450,7 +458,7 @@ def collect_firmware(self, pkg):
         col = ISPFWCollection("recovery/usr/sbin/")
         pkg.add_files(sorted(col.files()))
         logging.info("Collecting Kernel firmware")
-        col = KernelFWCollection(self.kernel_path)
+        col = KernelFWCollection(kernel_path)
         pkg.add_files(sorted(col.files()))
         logging.info("Making fallback firmware archive")
         subprocess.run(["tar", "czf", "all_firmware.tar.gz",