kdk_handler: Enumerate KDKs from Apple developer portal

resources/kdk_handler.py

@@ -1,19 +1,37 @@
 # Kernel Debug Kit downloader
 
+import datetime
+from typing import cast
+import urllib.parse
 from pathlib import Path
-import requests, urllib
+
+import packaging.version
+import requests
+
 from resources import utilities
+from resources.constants import Constants
 
 SESSION = requests.Session()
 
+
 class kernel_debug_kit_handler:
-
-
-    def __init__(self, constants):
+    def __init__(self, constants: Constants):
         self.constants = constants
 
+    def get_available_kdks(self):
+        KDK_API_LINK = "https://kdk-api.dhinak.net/v1"
 
-    def get_closest_match(self, host_version: str, host_build: str):
+        print("- Fetching available KDKs")
+
+        results = SESSION.get(KDK_API_LINK, headers={"User-Agent": f"OCLP/{self.constants.patcher_version}"})
+
+        if results.status_code != 200:
+            print(" - Could not fetch KDK list")
+            return None
+
+        return results.json().sort(key=lambda x: (packaging.version.parse(x["version"]), datetime.datetime.fromisoformat(x["date"])), reversed=True)
+
+    def get_closest_match_legacy(self, host_version: str, host_build: str):
         # Get the closest match to the provided version
         # KDKs are generally a few days late, so we'll rely on N-1 matching
 
@@ -21,12 +39,7 @@ class kernel_debug_kit_handler:
 
         OS_DATABASE_LINK = "https://api.appledb.dev/main.json"
 
-        newest_build = {
-            "name": "",
-            "version": "",
-            "build": "",
-            "date": "",
-        }
+        parsed_host_version = cast(packaging.version.Version, packaging.version.parse(host_version))
 
         print(f"- Checking closest match for: {host_version} build {host_build}")
 
@@ -36,65 +49,27 @@ class kernel_debug_kit_handler:
             print(" - Could not fetch database")
             return None, ""
 
-        results = results.json()
+        macos_builds = [i for i in results.json()["ios"] if i["osType"] == "macOS"]
+        macos_builds.sort(key=lambda x: (packaging.version.parse(x["version"]), datetime.datetime.fromisoformat(x["released"])), reverse=True)
 
         # Iterate through, find build that is closest to the host version
         # Use date to determine which is closest
-        if "ios" in results:
-            for variant in results["ios"]:
-                if variant["osStr"] == "macOS":
-
-                    name = variant["version"]
-                    version = name.split(" ")[0]
-                    build = variant["build"]
-                    date = variant["released"]
-
-                    if version != host_version:
-                        # Check if this is a security update (ie. 13.0.1)
-                        version_split = version.split(".")
-                        host_version_split = host_version.split(".")
-                        if len(version_split) >= 2 and len(host_version_split) >= 2:
-                            if not (version_split[0] == host_version_split[0] and version_split[1] == host_version_split[1]):
-                                continue
-                        else:
-                            continue
-
-                    if build == host_build:
-                        # Skip, as we want the next closest match
-                        continue
-
-                    # Check if this is the newest build
-                    if newest_build["date"] == "":
-                        newest_build = {
-                            "name": name,
-                            "version": version,
-                            "build": build,
-                            "date": date,
-                        }
-
-                    else:
-                        if date > newest_build["date"]:
-                            newest_build = {
-                                "name": name,
-                                "version": version,
-                                "build": build,
-                                "date": date,
-                            }
-
-        if newest_build["date"] != "":
-            print(f"- Closest match: {newest_build['version']} build {newest_build['build']}")
-            return self.generate_kdk_link(newest_build["version"], newest_build["build"])
+        for build_info in macos_builds:
+            if build_info["osType"] == "macOS":
+                version = cast(packaging.version.Version, packaging.version.parse(build_info["version"]))
+                if version == parsed_host_version:
+                    # Skip, as we want the next closest match
+                    continue
+                elif version <= parsed_host_version and version.major == parsed_host_version.major and version.minor == parsed_host_version.minor:
+                    # The KDK list is already sorted by date then version, so the first match is the closest
+                    print(f"- Closest match: {build_info['version']} build {build_info['build']}")
+                    return self.generate_kdk_link(build_info["version"], build_info["build"]), build_info["build"]
 
         print(" - Could not find a match")
         return None, ""
 
-
     def generate_kdk_link(self, version: str, build: str):
-        # Note: cannot do lazy matching as we don't store old version/build numbers nor can we enumerate KDKs from the portal
-        URL_TEMPLATE = f"https://download.developer.apple.com/macOS/Kernel_Debug_Kit_{version}_build_{build}/Kernel_Debug_Kit_{version}_build_{build}.dmg"
-
-        return URL_TEMPLATE, build
-
+        return f"https://download.developer.apple.com/macOS/Kernel_Debug_Kit_{version}_build_{build}/Kernel_Debug_Kit_{version}_build_{build}.dmg"
 
     def verify_apple_developer_portal(self, link):
         # Determine whether Apple Developer Portal is up
@@ -105,7 +80,7 @@ class kernel_debug_kit_handler:
         # 1: Portal is up but file is not available
         # 2: Portal is down
 
-        TOKEN_URL_BASE = "https://developerservices2.apple.com/services/download?path="
+        TOKEN_URL_BASE = "https://developerservices2.apple.com/services/download"
         remote_path = urllib.parse.urlparse(link).path
         token_url = urllib.parse.urlunparse(urllib.parse.urlparse(TOKEN_URL_BASE)._replace(query=urllib.parse.urlencode({"path": remote_path})))
 
@@ -126,7 +101,6 @@ class kernel_debug_kit_handler:
                 return 2
         return 0
 
-
     def download_kdk(self, version: str, build: str):
         error_msg = ""
         detected_build = build
@@ -135,37 +109,60 @@ class kernel_debug_kit_handler:
             print(" - KDK is already installed")
             return True, error_msg, detected_build
 
-        # Note: cannot do lazy matching as we don't store old version/build numbers nor can we enumerate KDKs from the portal
-        URL_TEMPLATE, detected_build = self.generate_kdk_link(version, build)
+        download_link = None
+        closest_match_download_link = None
+        closest_build = ""
+
+        kdk_list = self.get_available_kdks()
+
+        parsed_version = cast(packaging.version.Version, packaging.version.parse(version))
+
+        if kdk_list:
+            for kdk in kdk_list:
+                kdk_version = cast(packaging.version.Version, packaging.version.parse(kdk["version"]))
+                if kdk["build"] == build:
+                    download_link = kdk["download_url"]
+                elif not closest_match_download_link and kdk_version <= parsed_version and kdk_version.major == parsed_version.major and kdk_version.minor == parsed_version.minor:
+                    # The KDK list is already sorted by date then version, so the first match is the closest
+                    closest_match_download_link = kdk["download_url"]
+                    closest_build = kdk["build"]
+        else:
+            print(" - Could not fetch KDK list, falling back to brute force")
+            download_link = self.generate_kdk_link(version, build)
+            closest_match_download_link, closest_build = self.get_closest_match_legacy(version, build)
 
         print(f"- Downloading Apple KDK for macOS {version} build {build}")
-        result = self.verify_apple_developer_portal(URL_TEMPLATE)
-        if result == 2:
-            error_msg = "Could not contact Apple download servers"
-            return False, error_msg, ""
-        elif result == 0:
+        # download_link is None if no matching KDK is found, so we'll fall back to the closest match
+        result = self.verify_apple_developer_portal(download_link) if download_link else 1
+        if result == 0:
             print(" - Downloading KDK")
         elif result == 1:
             print(" - Could not find KDK, finding closest match")
-            URL_TEMPLATE, detected_build = self.get_closest_match(version, build)
 
-            if self.is_kdk_installed(detected_build) is True:
-                return True, error_msg, detected_build
-            if URL_TEMPLATE is None:
+            if self.is_kdk_installed(closest_build) is True:
+                return True, error_msg, closest_build
+
+            if closest_match_download_link is None:
                 error_msg = "Could not find KDK for host, nor closest match"
                 return False, error_msg, ""
-            result = self.verify_apple_developer_portal(URL_TEMPLATE)
-            if result == 2:
-                error_msg = "Could not contact Apple download servers"
-                return False, error_msg, ""
-            elif result == 0:
+
+            result = self.verify_apple_developer_portal(closest_match_download_link)
+
+            if result == 0:
                 print(" - Downloading KDK")
+                download_link = closest_match_download_link
             elif result == 1:
                 print(" - Could not find KDK")
                 error_msg = "Could not find KDK for host on Apple's servers, nor closest match"
                 return False, error_msg, ""
+            elif result == 2:
+                error_msg = "Could not contact Apple download servers"
+                return False, error_msg, ""
+        elif result == 2:
+            error_msg = "Could not contact Apple download servers"
+            return False, error_msg, ""
 
-        if utilities.download_apple_developer_portal(URL_TEMPLATE, self.constants.kdk_download_path):
+        if utilities.download_apple_developer_portal(download_link, self.constants.kdk_download_path):
             return True, error_msg, detected_build
         error_msg = "Failed to download KDK"
         return False, error_msg, ""
@@ -176,4 +173,4 @@ class kernel_debug_kit_handler:
                 if file.is_dir():
                     if file.name.endswith(f"{build}.kdk"):
                         return True
-        return False
+        return False