From d9ed037d7ab70b5c20f3e190dc52606ce4ed323d Mon Sep 17 00:00:00 2001
From: Mykola Grymalyuk <khronokernel@icloud.com>
Date: Tue, 24 May 2022 17:13:39 -0600
Subject: [PATCH] gui_main.py: Prevent Idle sleep during long processes

---
 CHANGELOG.md           |  1 +
 gui/gui_main.py        |  9 +++++++++
 resources/utilities.py | 31 ++++++++++++++++++++++---------
 3 files changed, 32 insertions(+), 9 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index ba0134d49..63c0d1b78 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -11,6 +11,7 @@
     - Avoids UI elements getting under the dock
   - Add return to disk when selecting partitions
   - Add "Search for disks again" option during OpenCore Install
+  - Prevent Idle Sleep while running long processes (ie. downloading, flashing)
 - Resolve failing to find binaries with `--patch_sys_vol` argument
 
 ## 0.4.5
diff --git a/gui/gui_main.py b/gui/gui_main.py
index d33bdfbec..c5950b9bf 100644
--- a/gui/gui_main.py
+++ b/gui/gui_main.py
@@ -105,6 +105,9 @@ class wx_python_gui:
         sys.stdout = self.stock_stdout
         sys.stderr = self.stock_stderr
         self.reset_frame_modal()
+
+        # Re-enable sleep if we failed to do so before returning to the main menu
+        utilities.enable_sleep_after_running()
     
     def reset_frame_modal(self):
         if not self.frame_modal:
@@ -1551,6 +1554,7 @@ class wx_python_gui:
             # If we're unable to download the integrity file immediately after downloading the IA, there's a legitmate issue
             # on Apple's end.
             # Fail gracefully and just head to installing the IA.
+            utilities.disable_sleep_while_running()
             apple_integrity_file = str(integrity_path)
             chunks = integrity_verification.generate_chunklist_dict(str(apple_integrity_file))
             if chunks:
@@ -1600,6 +1604,7 @@ class wx_python_gui:
         self.verifying_chunk_label.Centre(wx.HORIZONTAL)
         self.return_to_main_menu.Bind(wx.EVT_BUTTON, self.flash_installer_menu)
         self.return_to_main_menu.Centre(wx.HORIZONTAL)
+        utilities.enable_sleep_after_running()
 
 
     def flash_installer_menu(self, event=None):
@@ -1871,6 +1876,7 @@ class wx_python_gui:
         self.return_to_main_menu.Enable()
 
     def start_script(self):
+        utilities.disable_sleep_while_running()
         args = [self.constants.oclp_helper_path, "/bin/sh", self.constants.installer_sh_path]
         output, error, returncode = run.Run()._stream_output(comm=args)
         if "Install media now available at" in output:
@@ -1886,6 +1892,7 @@ class wx_python_gui:
             print("- Failed to create macOS installer")
             popup = wx.MessageDialog(self.frame, f"Failed to create macOS installer\n\nOutput: {output}\n\nError: {error}", "Error", wx.OK | wx.ICON_ERROR)
             popup.ShowModal()
+        utilities.enable_sleep_after_running()
 
 
     def download_and_unzip_pkg(self):
@@ -1908,6 +1915,8 @@ class wx_python_gui:
             path = self.constants.installer_pkg_path
 
         if utilities.download_file(link, path):
+            # Download thread will re-enable Idle Sleep after downloading
+            utilities.disable_sleep_while_running()
             if str(path).endswith(".zip"):
                 if Path(self.constants.installer_pkg_path).exists():
                     subprocess.run(["rm", self.constants.installer_pkg_path])
diff --git a/resources/utilities.py b/resources/utilities.py
index 8cb7a78ff..cb33247b4 100644
--- a/resources/utilities.py
+++ b/resources/utilities.py
@@ -12,15 +12,8 @@ import binascii
 import argparse
 from ctypes import CDLL, c_uint, byref
 import time
-
-try:
-    import requests
-except ImportError:
-    subprocess.run(["pip3", "install", "requests"], stdout=subprocess.PIPE)
-    try:
-        import requests
-    except ImportError:
-        raise Exception("Missing requests library!\nPlease run the following before starting OCLP:\npip3 install requests")
+import atexit
+import requests
 
 from resources import constants, ioreg
 from data import sip_data, os_data
@@ -135,6 +128,23 @@ def csr_decode(os_sip):
 def friendly_hex(integer: int):
     return "{:02X}".format(integer)
 
+sleep_process = None
+
+def disable_sleep_while_running():
+    global sleep_process
+    print("- Disabling Idle Sleep")
+    if sleep_process is None:
+        # If sleep_process is active, we'll just keep it running
+        sleep_process = subprocess.Popen(["caffeinate", "-d", "-i", "-s"], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+    # Ensures that if we don't properly close the process, 'atexit' will for us
+    atexit.register(enable_sleep_after_running)
+
+def enable_sleep_after_running():
+    global sleep_process
+    if sleep_process:
+        print("- Re-enabling Idle Sleep")
+        sleep_process.kill()
+        sleep_process = None
 
 def amfi_status():
     amfi_1 = "amfi_get_out_of_my_way=0x1"
@@ -362,6 +372,7 @@ def verify_network_connection(url):
 
 def download_file(link, location, is_gui=None, verify_checksum=False):
     if verify_network_connection(link):
+        disable_sleep_while_running()
         short_link = os.path.basename(link)
         if Path(location).exists():
             Path(location).unlink()
@@ -421,7 +432,9 @@ def download_file(link, location, is_gui=None, verify_checksum=False):
                 while chunk:
                     checksum.update(chunk)
                     chunk = file.read(1024 * 1024 * 16)
+            enable_sleep_after_running()
             return checksum
+        enable_sleep_after_running()
         return True
     else:
         cls()