From 3a6f87c9c8e0cb7fc9ec6c9eb2ed553cffc5aa90 Mon Sep 17 00:00:00 2001 From: Mykola Grymalyuk Date: Thu, 2 Feb 2023 10:50:29 -0700 Subject: [PATCH] network_handler.py: Initial implementation --- resources/gui/gui_main.py | 65 ++++++++++++--- resources/network_handler.py | 150 +++++++++++++++++++++++++++++++++++ 2 files changed, 203 insertions(+), 12 deletions(-) create mode 100644 resources/network_handler.py diff --git a/resources/gui/gui_main.py b/resources/gui/gui_main.py index ed834ff53..b8a15ce3f 100644 --- a/resources/gui/gui_main.py +++ b/resources/gui/gui_main.py @@ -19,8 +19,8 @@ from datetime import datetime import py_sip_xnu import logging -from resources import constants, defaults, install, installer, utilities, run, generate_smbios, updates, integrity_verification, global_settings, kdk_handler -from resources.sys_patch import sys_patch_download, sys_patch_detect, sys_patch, sys_patch_auto +from resources import constants, defaults, install, installer, utilities, run, generate_smbios, updates, integrity_verification, global_settings, kdk_handler, network_handler +from resources.sys_patch import sys_patch_download, sys_patch_detect, sys_patch from resources.build import build from data import model_array, os_data, smbios_data, sip_data, cpu_data from resources.gui import menu_redirect, gui_help @@ -1750,14 +1750,32 @@ class wx_python_gui: ) ) self.download_label.Centre(wx.HORIZONTAL) - # Redirect stdout to label - logging.getLogger().handlers[1].stream = menu_redirect.RedirectLabel(self.download_label) + + self.download_label_2 = wx.StaticText(self.frame, label="") + self.download_label_2.SetFont(wx.Font(12, wx.DEFAULT, wx.NORMAL, wx.BOLD)) + self.download_label_2.SetPosition( + wx.Point( + self.download_label.GetPosition().x, + self.download_label.GetPosition().y + self.download_label.GetSize().height + 5 + ) + ) + self.download_label_2.Centre(wx.HORIZONTAL) + + # Progress Bar + self.download_progress = wx.Gauge(self.frame, range=100, size=(self.frame.GetSize().width - 100, 20)) + self.download_progress.SetPosition( + wx.Point( + self.download_label_2.GetPosition().x, + self.download_label_2.GetPosition().y + self.download_label_2.GetSize().height + 5 + ) + ) + self.download_progress.Centre(wx.HORIZONTAL) self.return_to_main_menu = wx.Button(self.frame, label="Return to Main Menu") self.return_to_main_menu.SetPosition( wx.Point( - self.download_label.GetPosition().x, - self.download_label.GetPosition().y + self.download_label.GetSize().height + 30 + self.download_progress.GetPosition().x, + self.download_progress.GetPosition().y + self.download_progress.GetSize().height + 15 ) ) self.return_to_main_menu.Bind(wx.EVT_BUTTON, self.main_menu) @@ -1765,16 +1783,32 @@ class wx_python_gui: self.frame.SetSize(-1, self.return_to_main_menu.GetPosition().y + self.return_to_main_menu.GetSize().height + 40) wx.GetApp().Yield() + + ia_download = network_handler.download_object(app_dict['Link']) + ia_download.download(self.constants.payload_path / "InstallAssistant.pkg") + + while ia_download.is_active(): + wx.GetApp().Yield() + self.download_label.SetLabel(f"{utilities.human_fmt(ia_download.downloaded_file_size)} downloaded of {utilities.human_fmt(ia_download.total_file_size)} ({ia_download.get_percent():.2f}%)") + self.download_label.Centre(wx.HORIZONTAL) + self.download_label_2.SetLabel( + f"Average download speed: {utilities.human_fmt(ia_download.get_speed())}/s" + ) + self.download_label_2.Centre(wx.HORIZONTAL) + + self.download_progress.SetValue(ia_download.get_percent()) + + wx.GetApp().Yield() + time.sleep(0.1) + + # Download macOS install data - if installer.download_install_assistant(self.constants.payload_path, app_dict['Link']): - # Fix stdout - logging.getLogger().handlers[1].stream = self.stock_stream + if ia_download.download_complete: self.download_label.SetLabel(f"Finished Downloading {installer_name}") self.download_label.Centre(wx.HORIZONTAL) wx.App.Get().Yield() - self.installer_validation(apple_integrity_file_link= app_dict['integrity']) + self.installer_validation(apple_integrity_file_link=app_dict['integrity']) else: - logging.getLogger().handlers[1].stream = self.stock_stream self.download_label.SetLabel(f"Failed to download {installer_name}") self.download_label.Centre(wx.HORIZONTAL) @@ -2245,7 +2279,14 @@ class wx_python_gui: else: path = self.constants.installer_pkg_path - if utilities.download_file(link, path): + + autopkg_download = network_handler.download_object(link) + autopkg_download.download(path, display_progress=False) + + while autopkg_download.is_active(): + time.sleep(0.1) + + if autopkg_download.download_complete: # Download thread will re-enable Idle Sleep after downloading utilities.disable_sleep_while_running() if str(path).endswith(".zip"): diff --git a/resources/network_handler.py b/resources/network_handler.py new file mode 100644 index 000000000..60289f7b7 --- /dev/null +++ b/resources/network_handler.py @@ -0,0 +1,150 @@ +# Download files from the network +# Implements an object, where other libraries can use to query download status + +import time +import requests +import threading +import logging +from pathlib import Path + +from resources import utilities + + +SESSION = requests.Session() + +class network_utilities: + + def __init__(self, url): + self.url: str = url + + + def verify_network_connection(self): + try: + response = requests.head(self.url, timeout=5, allow_redirects=True) + return True + except ( + requests.exceptions.Timeout, + requests.exceptions.TooManyRedirects, + requests.exceptions.ConnectionError, + requests.exceptions.HTTPError + ): + return False + + +class download_object: + + def __init__(self, url): + self.url: str = url + self.status: str = "Downloading" + self.error_msg: str = "" + self.filename: str = self._get_filename() + + self.total_file_size: float = 0.0 + self.downloaded_file_size: float = 0.0 + self.start_time: float = time.time() + + self.error: bool = False + self.should_stop: bool = False + self.has_network: bool = network_utilities(self.url).verify_network_connection() + self.download_complete: bool = False + + self.active_thread: threading.Thread = None + + if self.has_network: + self._populate_file_size() + + + def __del__(self): + self.stop() + + + def download(self, path, display_progress=False): + if self.active_thread: + return + logging.info(f"Starting download: {self.filename}") + self.active_thread = threading.Thread(target=self._download, args=(path,display_progress,)) + self.active_thread.start() + + + def _get_filename(self): + return Path(self.url).name + + + def _populate_file_size(self): + try: + self.total_file_size = int(requests.head(self.url, allow_redirects=True, timeout=5).headers['Content-Length']) + except Exception as e: + self.error = True + self.error_msg = str(e) + self.status = "Error" + logging.error(f"Error determining file size {self.url}: {self.error_msg}") + + + def _prepare_working_directory(self, path): + if Path(path).exists(): + Path(path).unlink() + + + def _download(self, path, display_progress=False): + try: + if not self.has_network: + raise Exception("No network connection") + + self._prepare_working_directory(path) + + response = SESSION.get(self.url, stream=True) + + with open(path, 'wb') as file: + for i, chunk in enumerate(response.iter_content(1024 * 1024 * 4)): + if self.should_stop: + raise Exception("Download stopped") + if chunk: + file.write(chunk) + self.downloaded_file_size += len(chunk) + if display_progress and i % 100: + print(f"Downloaded {self.get_percent():.2f}% of {self.filename} ({utilities.human_fmt(self.get_speed())}/s) ({self.get_time_remaining():.2f} seconds remaining)") + self.download_complete = True + logging.info(f"Download complete: {self.filename}") + except Exception as e: + self.error = True + self.error_msg = str(e) + self.status = "Error" + logging.error(f"Error downloading {self.url}: {self.error_msg}") + self.status = "Done" + + + def get_percent(self): + if self.total_file_size == 0: + logging.error("File size is 0, cannot calculate percent") + return -1 + return self.downloaded_file_size / self.total_file_size * 100 + + + def get_speed(self): + return self.downloaded_file_size / (time.time() - self.start_time) + + + def get_time_remaining(self): + if self.total_file_size == 0: + logging.error("File size is 0, cannot calculate time remaining") + return -1 + return (self.total_file_size - self.downloaded_file_size) / self.get_speed() + + + def get_file_size(self): + return self.total_file_size + + + def is_active(self): + if self.status == "Downloading": + return True + return False + + + def stop(self): + self.should_stop = True + + + + +