Go back to flaresolverr

commands/pygg.py

@@ -1,5 +1,7 @@
 #!/usr/bin/env python3
 import argparse
+import io
+import json
 import os
 import re
 import urllib
@@ -9,9 +11,82 @@ import requests
 
 
 YGG_DOMAIN = "www6.yggtorrent.lol"
-YGG_IP = os.getenv("YGG_IP")
+FLARE_ENDPOINT = os.getenv("FLARE_ENDPOINT")
 
 
+class FlareRequests(requests.Session):
+    def request(self, method, url, params=None, data=None, **kwargs):
+        if not FLARE_ENDPOINT:
+            return super().request(method, url, params, data, **kwargs)
+
+        sessions = requests.post(FLARE_ENDPOINT, json={"cmd": "sessions.list"}).json()
+
+        if "sessions" in sessions and len(sessions["sessions"]) > 0:
+            FLARE_SESSION = sessions["sessions"][0]
+        else:
+            response = requests.post(FLARE_ENDPOINT, json={"cmd": "sessions.create"})
+            session = response.json()
+
+            if "session" in session:
+                FLARE_SESSION = session["session"]
+            else:
+                raise requests.RequestException(response)
+
+        if params:
+            url += "&" if len(url.split("?")) > 1 else "?"
+            url = f"{url}{urllib.parse.urlencode(params)}"
+
+        post_data = {
+            "cmd": f"request.{method.lower()}",
+            "session": FLARE_SESSION,
+            "url": url,
+        }
+
+        if data:
+            post_data["postData"] = urllib.parse.urlencode(data)
+
+        try:
+            response = requests.post(
+                FLARE_ENDPOINT,
+                json=post_data,
+            )
+
+            content = response.json()
+
+            if "solution" in content:
+                encoding = None
+                solution = content["solution"]
+                if "content-type" in solution["headers"]:
+                    content_type = solution["headers"]["content-type"].split(";")
+                    if len(content_type) > 1:
+                        charset = content_type[1].split("=")
+                        if len(charset) > 1:
+                            encoding = charset[1]
+
+                resolved = requests.Response()
+
+                resolved.status_code = solution["status"]
+                resolved.headers = solution["headers"]
+                resolved.raw = io.BytesIO(solution["response"].encode())
+                resolved.url = url
+                resolved.encoding = encoding
+                resolved.reason = content["status"]
+                resolved.cookies = solution["cookies"]
+
+                return resolved
+
+            raise requests.RequestException(response)
+        except requests.RequestException:
+            session = requests.post(
+                FLARE_ENDPOINT,
+                json={"cmd": "sessions.destroy", "session": FLARE_SESSION},
+            )
+
+            raise requests.RequestException(solution)
+
+
+session = FlareRequests()
+
 parser = argparse.ArgumentParser()
 parser.add_argument("-u", "--uploader", action="append")
 parser.add_argument("-b", "--blacklist", action="append", default=["dvd", "iso"])
@@ -31,12 +106,8 @@ def parse_size(size):
 
 
 def check_files(id):
-    req = requests.get(
+    req = requests.get(f"https://{YGG_DOMAIN}/engine/get_files", params={"torrent": id})
-        f"http://{YGG_IP}/engine/get_files",
+    res = json.loads(req.text)
-        params={"torrent": id},
-        headers={"Host": YGG_DOMAIN},
-    )
-    res = req.json()
     html = bs4.BeautifulSoup(res["html"], "html.parser")
     trs = html.select("tr")
     return len(trs) == 1 and "mkv" in trs[0].get_text().lower()
@@ -58,11 +129,7 @@ def search_ygg(query, multi, full):
     if multi:
         ygg_params["option_langue:multiple[]"] = "4"
 
-    req = requests.get(
+    req = session.get(f"https://{YGG_DOMAIN}/engine/search", params=ygg_params)
-        f"http://{YGG_IP}/engine/search",
-        params=ygg_params,
-        headers={"Host": YGG_DOMAIN},
-    )
     html = bs4.BeautifulSoup(
         req.text,
         "html.parser",