From c5f72c10ff283fd71cb435f6b514ca86f36bdf64 Mon Sep 17 00:00:00 2001
From: Mykola Grymalyuk <khronokernel@icloud.com>
Date: Wed, 8 Feb 2023 14:31:54 -0700
Subject: [PATCH] kdk_handler.py: Clear out backup KDKs when unused

---
 CHANGELOG.md             |  2 ++
 resources/kdk_handler.py | 16 +++++++++++++---
 2 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index f488f3e5d..77c6bc4cd 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -10,6 +10,8 @@
 - Implement Kernel Debug Kit installation during OS installs
   - Avoids network requirement for first time installs
   - Paired along side AutoPkgInstaller
+- Implement Kernel Debug Kit backup system
+  - Allows for easy restoration of KDKs if OS updates corrupted installed KDKs
 - Backend Changes:
   - Refactored kdk_handler.py
     - Prioritizes KdkSupportPkg repository for downloads
diff --git a/resources/kdk_handler.py b/resources/kdk_handler.py
index 0476b00eb..f2b4b8f96 100644
--- a/resources/kdk_handler.py
+++ b/resources/kdk_handler.py
@@ -399,7 +399,15 @@ class KernelDebugKitObject:
             logging.warning("- Cannot remove KDK, not running as root")
             return
 
-        result = utilities.elevated(["rm", "-rf", kdk_path], stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
+        if not Path(kdk_path).exists():
+            logging.warning(f"- KDK does not exist: {kdk_path}")
+            return
+
+        rm_args = ["rm", "-f", kdk_path]
+        if Path(kdk_path).is_dir():
+            rm_args = ["rm", "-rf", kdk_path]
+
+        result = utilities.elevated(rm_args, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
         if result.returncode != 0:
             logging.warning(f"- Failed to remove KDK: {kdk_path}")
             logging.warning(f"- {result.stdout.decode('utf-8')}")
@@ -434,10 +442,12 @@ class KernelDebugKitObject:
         logging.info("- Cleaning unused KDKs")
         for kdk_folder in Path(KDK_INSTALL_PATH).iterdir():
             if kdk_folder.is_dir():
-                if kdk_folder.name.endswith(".kdk"):
+                if kdk_folder.name.endswith(".kdk") or kdk_folder.name.endswith(".pkg"):
                     should_remove = True
                     for build in exclude_builds:
-                        if build != "" and kdk_folder.name.endswith(f"{build}.kdk"):
+                        if build != "":
+                            continue
+                        if kdk_folder.name.endswith(f"{build}.kdk") or kdk_folder.name.endswith(f"{build}.pkg"):
                             should_remove = False
                             break
                     if should_remove is False: