fix issues with steam drm downloader

ProjectBorealis · Aug 11, 2024 · a9769b1 · a9769b1
1 parent 907da82
commit a9769b1
Show file tree

Hide file tree

Showing 4 changed files with 51 additions and 49 deletions.
diff --git a/.vscode/launch.json b/.vscode/launch.json
@@ -1,5 +1,5 @@
 {
-    "version": "0.2.0",
+    "version": "0.7.0",
     "inputs": [
         {
             "id": "debugPath",
@@ -176,7 +176,10 @@
             "request": "launch",
             "module": "pbsync",
             "console": "integratedTerminal",
-            "justMyCode": false
+            "justMyCode": false,
+            "env": {
+                "GEVENT_SUPPORT": "True"
+            },
         },
         {
             "name": "Migrate Assets",

diff --git a/pbpy/pbsteamcmd.py b/pbpy/pbsteamcmd.py
@@ -6,6 +6,7 @@
 import urllib.request
 from pathlib import Path
 
+import gevent
 import steam.protobufs.steammessages_partnerapps_pb2  # don't remove
 from steam.client import SteamClient
 
@@ -181,29 +182,6 @@ def handle_drm_file():
 
         if drm_proc.returncode != 0:
 
-            def finish_handling_drm_fail():
-                if not drm_output.exists() and drm_download_failed:
-                    input(
-                        f"DRM download failed, download the file from https://partner.steamgames.com/apps/drm/{drm_app_id} and place it at {drm_output}, then press enter"
-                    )
-                handled_drm = False
-                if drm_output.exists():
-                    handle_drm_file()
-                    handled_drm = True
-                if (
-                    drm_proc.returncode != 0
-                    and not drm_download_failed
-                    or drm_download_failed
-                    and not handled_drm
-                ):
-                    if drm_output.exists():
-                        pbtools.remove_file(str(drm_output))
-                    pbtools.error_state(
-                        f"DRM wrapping failed: exit code {drm_proc.returncode}",
-                        hush=True,
-                        term=True,
-                    )
-
             if drm_download_failed and drm_id:
                 steamclient = SteamWorker()
                 steamclient.login(
@@ -212,30 +190,47 @@ def finish_handling_drm_fail():
                     True,
                 )
 
-                def wait_for_drm_download():
-                    time.sleep(5)
-                    resp = steamclient.steam.send_um_and_wait(
-                        "PartnerApps.Download#1",
-                        {
-                            "file_id": f"/{drm_app_id}/{drm_id}/{drm_exe_path.name}_{drm_id}",
-                            "app_id": int(drm_app_id),
-                        },
-                    )
-                    url = resp.body.download_url
-                    if url:
-                        with urllib.request.urlopen(url) as response, open(
-                            str(drm_output), "wb"
-                        ) as out_file:
-                            shutil.copyfileobj(response, out_file)
-                    steamclient.close()
-                    finish_handling_drm_fail()
-
-                if steamclient.logged_on_once:
-                    wait_for_drm_download()
-                else:
-                    steamclient.steam.once("logged_on", wait_for_drm_download)
-            else:
-                finish_handling_drm_fail()
+                while not steamclient.logged_on_once:
+                    # TODO: doesn't work, "this operation would block forever"
+                    # steamclient.steam.wait_event("logged_on")
+                    gevent.sleep(1)
+
+                resp = steamclient.steam.send_um_and_wait(
+                    "PartnerApps.Download#1",
+                    {
+                        "file_id": f"/{drm_app_id}/{drm_id}/{drm_exe_path.name}_{drm_id}",
+                        "app_id": int(drm_app_id),
+                    },
+                )
+                url = resp.body.download_url
+                if url:
+                    with urllib.request.urlopen(url) as response, open(
+                        str(drm_output), "wb"
+                    ) as out_file:
+                        shutil.copyfileobj(response, out_file)
+                steamclient.close()
+
+            if not drm_output.exists() and drm_download_failed:
+                input(
+                    f"DRM download failed, download the file from https://partner.steamgames.com/apps/drm/{drm_app_id} and place it at {drm_output}, then press enter"
+                )
+            handled_drm = False
+            if drm_output.exists():
+                handle_drm_file()
+                handled_drm = True
+            if (
+                drm_proc.returncode != 0
+                and not drm_download_failed
+                or drm_download_failed
+                and not handled_drm
+            ):
+                if drm_output.exists():
+                    pbtools.remove_file(str(drm_output))
+                pbtools.error_state(
+                    f"DRM wrapping failed: exit code {drm_proc.returncode}",
+                    hush=True,
+                    term=True,
+                )
         else:
             handle_drm_file()
     else:

diff --git a/pbsync/__init__.py b/pbsync/__init__.py
diff --git a/pbsync/__main__.py b/pbsync/__main__.py
@@ -1,3 +1,7 @@
+from gevent import monkey
+
+monkey.patch_all()
+
 import argparse
 import multiprocessing
 import os