network_handler.py: Fix checksum support

opencore_legacy_patcher/support/network_handler.py

@@ -13,7 +13,7 @@ import enum
 import hashlib
 import atexit
 
-from typing import Union
+from typing import Optional, Union
 from pathlib import Path
 
 from . import utilities
@@ -162,7 +162,7 @@ class DownloadObject:
 
     """
 
-    def __init__(self, url: str, path: str) -> None:
+    def __init__(self, url: str, path: str, checksum_algo: Optional["hashlib._Hash"] = None) -> None:
         self.url:       str = url
         self.status:    str = DownloadStatus.INACTIVE
         self.error_msg: str = ""
@@ -181,10 +181,8 @@ class DownloadObject:
 
         self.active_thread: threading.Thread = None
 
-        self.should_checksum: bool = False
-
         self.checksum = None
-        self._checksum_storage: hash = None
+        self._checksum_storage: Optional[hashlib._Hash] = checksum_algo
 
         if self.has_network:
             self._populate_file_size()
@@ -194,7 +192,7 @@ class DownloadObject:
         self.stop()
 
 
-    def download(self, display_progress: bool = False, spawn_thread: bool = True, verify_checksum: bool = False) -> None:
+    def download(self, display_progress: bool = False, spawn_thread: bool = True) -> None:
         """
         Download the file
 
@@ -204,7 +202,7 @@ class DownloadObject:
         Parameters:
             display_progress (bool): Display progress in console
             spawn_thread (bool): Spawn a thread to download the file, otherwise download in the current thread
-            verify_checksum (bool): Calculate checksum of downloaded file if True
+            verify_checksum (Optional[hashlib._Hash]): Checksum algorithm to use for verifying the download, optional
 
         """
         self.status = DownloadStatus.DOWNLOADING
@@ -213,12 +211,10 @@ class DownloadObject:
             if self.active_thread:
                 logging.error("Download already in progress")
                 return
-            self.should_checksum = verify_checksum
             self.active_thread = threading.Thread(target=self._download, args=(display_progress,))
             self.active_thread.start()
             return
 
-        self.should_checksum = verify_checksum
         self._download(display_progress)
 
 
@@ -235,15 +231,14 @@ class DownloadObject:
         """
 
         if verify_checksum:
-            self.should_checksum = True
+            self._checksum_storage = hashlib.sha256()
-            self.checksum = hashlib.sha256()
 
         self.download(spawn_thread=False)
 
         if not self.download_complete:
             return False
 
-        return self.checksum.hexdigest() if self.checksum else True
+        return self._checksum_storage.hexdigest() if self._checksum_storage else True
 
 
     def _get_filename(self) -> str:
@@ -283,7 +278,8 @@ class DownloadObject:
         Parameters:
             chunk (bytes): Chunk to update checksum with
         """
-        self._checksum_storage.update(chunk)
+        if self._checksum_storage:
+            self._checksum_storage.update(chunk)
 
 
     def _prepare_working_directory(self, path: Path) -> bool:
@@ -353,7 +349,7 @@ class DownloadObject:
                     if chunk:
                         file.write(chunk)
                         self.downloaded_file_size += len(chunk)
-                        if self.should_checksum:
+                        if self._checksum_storage:
                             self._update_checksum(chunk)
                         if display_progress and i % 100:
                             # Don't use logging here, as we'll be spamming the log file
@@ -368,6 +364,9 @@ class DownloadObject:
                 logging.info(f"- Time elapsed: {(time.time() - self.start_time):.2f} seconds")
                 logging.info(f"- Speed: {utilities.human_fmt(self.downloaded_file_size / (time.time() - self.start_time))}/s")
                 logging.info(f"- Location: {self.filepath}")
+                if self._checksum_storage:
+                    self.checksum = self._checksum_storage.hexdigest()
+                    logging.info(f"Checksum: {self.checksum}")
         except Exception as e:
             self.error = True
             self.error_msg = str(e)