mkbloke · GitHub

Don't add hooks to the entire HTTP session. This is bad.

diff --git a/src/streamlink/plugins/vk.py b/src/streamlink/plugins/vk.py
index 33a6365c..c72e9458 100644
--- a/src/streamlink/plugins/vk.py
+++ b/src/streamlink/plugins/vk.py
@@ -14,6 +14,7 @@ from streamlink.plugin import Plugin, PluginError, pluginmatcher
 from streamlink.plugin.api import validate
 from streamlink.stream.dash import DASHStream
 from streamlink.stream.hls import HLSStream
+from streamlink.utils.url import update_qsd
 log = logging.getLogger(__name__)
@@ -29,24 +30,25 @@ class VK(Plugin):
     HASH_COOKIE = "hash429"
     SOLUTION_COOKIE = "solution429"
-    def check_for_waf(self, res, **kwargs):
-        if res.headers.get("X-WAF-Backend-Status") == "challenge_success" and hasattr(self, "original_request"):
-            self.session.http.cookies.update(res.cookies)
-            self.session.http.cookies.pop(self.HASH_COOKIE, None)
-            self.original_request.headers.pop("Cookie", None)
-            self.original_request.prepare_cookies(self.session.http.cookies)
-            res = self.session.http.send(self.original_request)
-            del self.original_request
-            self.save_cookies(cookie_filter=lambda c: c.name == self.SOLUTION_COOKIE)
-            return res
-
-        if res.status_code == 302 and res.cookies.get(self.HASH_COOKIE):
-            log.debug("Getting new WAF cookie")
-            cookie = res.cookies.get(self.HASH_COOKIE)
-            key = md5(cookie.encode("utf-8")).hexdigest()
-            res.headers["Location"] = f"{res.headers['Location']}&key={key}"
-            self.original_request = res.request
-            return res
+    def _get_cookies(self):
+        if self.session.http.cookies.get(self.SOLUTION_COOKIE, domain=".vk.com"):
+            log.debug("Using cached WAF cookie")
+            return
+
+        def on_response(res, **kwargs):
+            if res.headers.get("x-waf-redirect") == "1":
+                if not res.headers.get("X-WAF-Backend-Status"):
+                    log.debug("Getting new WAF cookie")
+                    cookie = res.cookies.get(self.HASH_COOKIE)
+                    key = md5(cookie.encode("utf-8")).hexdigest()
+                    res.headers["Location"] = update_qsd(res.headers["Location"], qsd={"key": key})
+                    return res
+                elif res.headers.get("X-WAF-Backend-Status") == "challenge_success":
+                    self.session.http.cookies.update(res.cookies)
+                    self.save_cookies(cookie_filter=lambda c: c.name == self.SOLUTION_COOKIE)
+                    return res
+
+        self.session.http.get("https://vk.com/", hooks={"response": on_response})
     def _has_video_id(self):
         return any(m for m in self.matches[:-1])
@@ -78,8 +80,7 @@ class VK(Plugin):
         raise NoStreamsError(self.url)
     def _get_streams(self):
-        self.session.http.hooks["response"] = self.check_for_waf
-
+        self._get_cookies()
         self.follow_vk_redirect()
         video_id = self.match.group("video_id")

Read the original on github.com ↗