integrity_verification.py: Add extra error handling

gui/gui_main.py

@@ -1452,30 +1452,33 @@ class wx_python_gui:
             # Fail gracefully and just head to installing the IA.
             apple_integrity_file = str(integrity_path)
             chunks = integrity_verification.generate_chunklist_dict(str(apple_integrity_file))
-            max_progress = len(chunks)
+            if chunks:
-            self.progress_bar.SetValue(0)
+                max_progress = len(chunks)
-            self.progress_bar.SetRange(max_progress)
+                self.progress_bar.SetValue(0)
+                self.progress_bar.SetRange(max_progress)
 
-            wx.App.Get().Yield()
+                wx.App.Get().Yield()
-            # See integrity_verification.py for more information on the integrity verification process
+                # See integrity_verification.py for more information on the integrity verification process
-            with Path(self.constants.payload_path / Path("InstallAssistant.pkg")).open("rb") as f:
+                with Path(self.constants.payload_path / Path("InstallAssistant.pkg")).open("rb") as f:
-                for chunk in chunks:
+                    for chunk in chunks:
-                    status = hashlib.sha256(f.read(chunk["length"])).digest()
+                        status = hashlib.sha256(f.read(chunk["length"])).digest()
-                    if not status == chunk["checksum"]:
+                        if not status == chunk["checksum"]:
-                        print(f"Chunk {chunks.index(chunk) + 1} checksum status FAIL: chunk sum {binascii.hexlify(chunk['checksum']).decode()}, calculated sum {binascii.hexlify(status).decode()}")
+                            print(f"Chunk {chunks.index(chunk) + 1} checksum status FAIL: chunk sum {binascii.hexlify(chunk['checksum']).decode()}, calculated sum {binascii.hexlify(status).decode()}")
-                        self.popup = wx.MessageDialog(
+                            self.popup = wx.MessageDialog(
-                        self.frame,
+                            self.frame,
-                            f"We've found that Chunk {chunks.index(chunk) + 1} of {len(chunks)} has failed the integrity check.\n\nThis generally happens when downloading on unstable connections such as WiFi or cellular.\n\nPlease try redownloading again on a stable connection (ie. Ethernet)",
+                                f"We've found that Chunk {chunks.index(chunk) + 1} of {len(chunks)} has failed the integrity check.\n\nThis generally happens when downloading on unstable connections such as WiFi or cellular.\n\nPlease try redownloading again on a stable connection (ie. Ethernet)",
-                            "Corrupted Installer!",
+                                "Corrupted Installer!",
-                            style = wx.OK | wx.ICON_EXCLAMATION
+                                style = wx.OK | wx.ICON_EXCLAMATION
-                        )
+                            )
-                        self.popup.ShowModal()
+                            self.popup.ShowModal()
-                        self.main_menu()
+                            self.main_menu()
-                        break
+                            break
-                    else:
+                        else:
-                        self.progress_bar.SetValue(self.progress_bar.GetValue() + 1)
+                            self.progress_bar.SetValue(self.progress_bar.GetValue() + 1)
-                        self.verifying_chunk_label.SetLabel(f"Verifying Chunk {self.progress_bar.GetValue()} of {max_progress}")
+                            self.verifying_chunk_label.SetLabel(f"Verifying Chunk {self.progress_bar.GetValue()} of {max_progress}")
-                        wx.App.Get().Yield()
+                            wx.App.Get().Yield()
+            else:
+                print("Invalid integrity file provided")
         else:
             print("Failed to download integrity file, skipping integrity check.")
     

resources/integrity_verification.py

@@ -13,6 +13,7 @@ CHUNK_LENGTH = 4 + 32
 def generate_chunklist_dict(chunklist):
     chunklist = Path(chunklist).read_bytes() if isinstance(chunklist, str) else chunklist
 
+    # Ref: https://github.com/apple-oss-distributions/xnu/blob/xnu-8020.101.4/bsd/kern/chunklist.h#L59-L69
     header = {
         "magic":       chunklist[:4],
         "length":      int.from_bytes(chunklist[4:8], "little"),
@@ -24,6 +25,9 @@ def generate_chunklist_dict(chunklist):
         "sigOffset":   int.from_bytes(chunklist[28:36], "little")
     }
 
+    if header["magic"] != b"CNKL":
+        return None
+
     all_chunks = chunklist[header["chunkOffset"]:header["chunkOffset"]+header["chunkCount"]*CHUNK_LENGTH]
     chunks = [{"length": int.from_bytes(all_chunks[i:i+4], "little"), "checksum": all_chunks[i+4:i+CHUNK_LENGTH]} for i in range(0, len(all_chunks), CHUNK_LENGTH)]
 
@@ -32,6 +36,8 @@ def generate_chunklist_dict(chunklist):
 
 def chunk(file_path, chunklist, verbose):
     chunks = generate_chunklist_dict(chunklist)
+    if chunks is None:
+        return False
     with Path(file_path).open("rb") as f:
         for chunk in chunks:
             status = hashlib.sha256(f.read(chunk["length"])).digest()