OpenCore-Legacy-Patcher/resources/network_handler.py

# 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 NetworkUtilities:
    """
    Utilities for network related tasks, primarily used for downloading files
    """

    def __init__(self, url):
        self.url: str = url


    def verify_network_connection(self):
        """
        Verifies that the network is available

        :return: True if the network is available, False if not
        """

        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 DownloadObject:
    """
    Object for downloading files from the network

    Usage:
        >>> download_object = DownloadObject(url)
        >>> download_object.download(path, display_progress=True)

        >>> if download_object.is_active():
        >>>     print(download_object.get_percent())

        >>> if not download_object.download_complete:
        >>>     print("Download failed")

        >>> print("Download complete"")

    """

    def __init__(self, url):
        self.url:       str = url
        self.status:    str = "Inactive"
        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.download_complete: bool = False
        self.has_network:       bool = NetworkUtilities(self.url).verify_network_connection()

        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):
        """
        Download the file

        Spawns a thread to download the file, so that the main thread can continue
        Note sleep is disabled while the download is active
        """

        if self.active_thread:
            return
        self.status = "Downloading"
        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):
        """
        Get the filename from the URL

        :return: The filename
        """

        return Path(self.url).name


    def _populate_file_size(self):
        """
        Get the file size of the file to be downloaded

        If unable to get file size, set to zero
        """
        try:
            result = requests.head(self.url, allow_redirects=True, timeout=5)
            if 'Content-Length' in result.headers:
                self.total_file_size = float(result.headers['Content-Length'])
            else:
                raise Exception("Content-Length missing from headers")
        except Exception as e:
            logging.error(f"Error determining file size {self.url}: {str(e)}")
            logging.error("Assuming file size is 0")
            self.total_file_size = 0.0


    def _prepare_working_directory(self, path):
        """
        Delete the file if it already exists

        :param path: Path to the file
        :return:     True if successful, False if not
        """

        try:
            if Path(path).exists():
                logging.info(f"Deleting existing file: {path}")
                Path(path).unlink()
                return True
            if not Path(path).parent.exists():
                logging.info(f"Creating directory: {Path(path).parent}")
                Path(path).parent.mkdir(parents=True, exist_ok=True)
        except Exception as e:
            self.error = True
            self.error_msg = str(e)
            self.status = "Error"
            logging.error(f"Error preparing working directory {path}: {self.error_msg}")
            return False

        return True


    def _download(self, path, display_progress=False):
        utilities.disable_sleep_while_running()

        try:
            if not self.has_network:
                raise Exception("No network connection")

            if self._prepare_working_directory(path) is False:
                raise Exception(self.error_msg)

            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:
                            # Don't use logging here, as we'll be spamming the log file
                            if self.total_file_size == 0.0:
                                print(f"Downloaded {utilities.human_fmt(self.downloaded_file_size)} of {self.filename}")
                            else:
                                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"
        utilities.enable_sleep_after_running()


    def get_percent(self):
        if self.total_file_size == 0.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.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
        if self.active_thread.is_alive():
            time.sleep(1)