From a9769b14b6dbbdebecb263f5d861077eaa45afa3 Mon Sep 17 00:00:00 2001
From: mastercoms <mastercoms@tuta.io>
Date: Sun, 11 Aug 2024 01:35:05 -0400
Subject: [PATCH] fix issues with steam drm downloader

---
 .vscode/launch.json |  7 +++-
 pbpy/pbsteamcmd.py  | 89 +++++++++++++++++++++------------------------
 pbsync/__init__.py  |  0
 pbsync/__main__.py  |  4 ++
 4 files changed, 51 insertions(+), 49 deletions(-)
 create mode 100644 pbsync/__init__.py

diff --git a/.vscode/launch.json b/.vscode/launch.json
index e093717..662e201 100644
--- 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
index e3e69c4..2add4a1 100644
--- 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
new file mode 100644
index 0000000..e69de29
diff --git a/pbsync/__main__.py b/pbsync/__main__.py
index 3dd8785..1b5987e 100644
--- 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