PyNyaaTa/pynyaata/session.py

from io import BytesIO
from os import getenv
from urllib import parse

from requests import Response, Session, post


CLOUDPROXY_ENDPOINT = getenv("CLOUDPROXY_ENDPOINT")


class FlareRequests(Session):
    def request(self, method, url, params=None, data=None, **kwargs):
        if not CLOUDPROXY_ENDPOINT:
            return super().request(method, url, params, data, **kwargs)

        if params:
            url += "&" if len(url.split("?")) > 1 else "?"
            url = f"{url}{parse.urlencode(params)}"

        post_data = {
            "cmd": f"request.{method.lower()}",
            "url": url,
        }

        if data:
            post_data["postData"] = parse.urlencode(data)

        response = post(
            CLOUDPROXY_ENDPOINT,
            json=post_data,
        )

        content = response.json()

        if "solution" in content:
            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 = Response()

            resolved.status_code = solution["status"]
            resolved.headers = solution["headers"]
            resolved.raw = BytesIO(solution["response"].encode())
            resolved.url = url
            resolved.encoding = encoding or None
            resolved.reason = content["status"]
            resolved.cookies = solution["cookies"]

            return resolved