From 44bf5826b9a45bb9bff1e3bb1dd613e3d2779faf Mon Sep 17 00:00:00 2001
From: SavageCore <talk@savagecore.uk>
Date: Tue, 9 Jul 2024 02:46:32 +0100
Subject: [PATCH] Improve `get_game_install_path`

Thanks to @Beelzebub2

https://github.com/Beelzebub2/AutoBanana/commit/19ec41666c4c3edf3768f13f07bd834eddb19bd1
---
 main.py | 130 ++++++++++++++++++++++++++++++--------------------------
 1 file changed, 69 insertions(+), 61 deletions(-)

diff --git a/main.py b/main.py
index e724f67..59b92c2 100644
--- a/main.py
+++ b/main.py
@@ -15,92 +15,100 @@
 import os
 import sys
 import winreg
-import vdf
-import json
+import vdf  # type: ignore
 
 
-# Return the install path of the game
-def get_game_install_path(app_id):
+def get_steam_install_location():
+    """The function `get_steam_install_location` retrieves the installation location of Steam from the
+    Windows registry.
+
+    Returns
+    -------
+        The function `get_steam_install_location` returns the installation location of Steam as a
+    string.
 
-    # First check Steam
+    """
     steam_key = winreg.OpenKey(
         winreg.HKEY_LOCAL_MACHINE,
-        "SOFTWARE\Wow6432Node\Valve\Steam",
+        r"SOFTWARE\Wow6432Node\Valve\Steam",
     )
 
     steam_install_location = winreg.QueryValueEx(steam_key, "InstallPath")[0]
 
     winreg.CloseKey(steam_key)
 
-    # Check if the game is installed in Steam
-    steam_apps_path = os.path.join(steam_install_location, "steamapps")
+    return steam_install_location
 
-    for root, dirs, files in os.walk(steam_apps_path):
-        for file in files:
-            if file == "appmanifest_" + app_id + ".acf":
-                with open(os.path.join(root, file), "r") as f:
-                    manifest = vdf.load(f)
 
-                    install_location = os.path.join(
-                        steam_apps_path, "common", manifest["AppState"]["installdir"]
-                    )
+def get_game_install_path(app_id):
+    """
+    Searches for the installation path of a game using the provided app ID in Steam directories.
+
+    Parameters
+    ----------
+    app_id : str
+        The Steam application ID of the game.
+
+    Returns
+    -------
+    str or None
+        The installation path of the game if found, otherwise None.
+    """
+    steam_apps_path = os.path.join(steam_install_location, "steamapps")
 
-                # Check if install location exists
-                if os.path.exists(install_location):
-                    return install_location
+    # Check the main steamapps directory first
+    install_path = _check_app_manifest(steam_apps_path, app_id)
+    if install_path:
+        return install_path
 
-    # If not, check the library folders
-    library_folders = os.path.join(
-        steam_install_location, "steamapps", "libraryfolders.vdf"
-    )
+    # Check the library folders
+    library_folders_file = os.path.join(steam_apps_path, "libraryfolders.vdf")
+    if os.path.exists(library_folders_file):
+        with open(library_folders_file, "r") as f:
+            library_folders = vdf.load(f)["libraryfolders"]
 
-    # Find app id in library folders
-    with open(library_folders, "r") as f:
-        library = vdf.load(f)
-
-        for key in library["libraryfolders"]:
+        for key, library in library_folders.items():
             if key == "0":
                 continue
 
-            library_path = library["libraryfolders"][key]["path"]
-            apps = library["libraryfolders"][key]["apps"]
-
-            if app_id in apps:
-                for root, dirs, files in os.walk(
-                    os.path.join(library_path, "steamapps")
-                ):
-                    for file in files:
-                        if file == "appmanifest_" + app_id + ".acf":
-                            with open(os.path.join(root, file), "r") as f:
-                                manifest = vdf.load(f)
-
-                                install_location = os.path.join(
-                                    library_path,
-                                    "steamapps",
-                                    "common",
-                                    manifest["AppState"]["installdir"],
-                                )
-
-                                # Check if install location exists
-                                if os.path.exists(install_location):
-                                    return install_location
-
-    # If not installed in Steam, check the Epic Games Store
-    epic_manifests_path = os.path.join(
-        os.getenv("ProgramData"), "Epic", "EpicGamesLauncher", "Data", "Manifests"
-    )
+            library_path = os.path.join(library["path"], "steamapps")
+            install_path = _check_app_manifest(library_path, app_id)
+            if install_path:
+                return install_path
 
-    for root, dirs, files in os.walk(epic_manifests_path):
-        for file in files:
-            with open(os.path.join(root, file), "r") as f:
-                manifest = json.load(f)
+    return None
 
-                if manifest["DisplayName"] == "PAYDAY 3":
-                    return manifest["InstallLocation"]
 
+def _check_app_manifest(steam_apps_path, app_id):
+    """
+    Helper method to check for the app manifest and return the installation path if found.
+
+    Parameters
+    ----------
+    steam_apps_path : str
+        Path to the steamapps directory.
+    app_id : str
+        The Steam application ID of the game.
+
+    Returns
+    -------
+    str or None
+        The installation path of the game if found, otherwise None.
+    """
+    manifest_file = os.path.join(steam_apps_path, f"appmanifest_{app_id}.acf")
+    if os.path.exists(manifest_file):
+        with open(manifest_file, "r") as f:
+            manifest = vdf.load(f)
+            install_location = os.path.join(
+                steam_apps_path, "common", manifest["AppState"]["installdir"]
+            )
+            if os.path.exists(install_location):
+                return install_location
     return None
 
 
+steam_install_location = get_steam_install_location()
+
 cwd = os.getcwd()
 
 game_path = get_game_install_path("1272080")