Source code for onbasca.onbasca.http_session

# SPDX-FileCopyrightText: 2022 The Tor Project, Inc.
#
# SPDX-License-Identifier: BSD-3-Clause

import logging

import requests

from onbasca.onbasca import config, constants

logger = logging.getLogger(__name__)




class TimedSession(requests.Session):
    """Requests Session that sends timeout in the head and get methods."""

    def __init__(
        self,
        socks_address,
        timeout=config.DL_TIMEOUT_SECS,
        verify=False,
        **kwargs,
    ):
        logger.debug("Initializing session.")
        super().__init__()
        self._timeout = timeout
        # Because WebServer.verify is an string and requests accepts both
        # string or bool.
        if verify.lower() == "true":
            self._verify = True
        elif verify.lower() == "false":
            self._verify = False
        else:
            self._verify = verify
        # socks5h resolves DNS on the proxy
        self.proxies = {
            "http": "socks5h://{}:{}".format(*socks_address),
            "https": "socks5h://{}:{}".format(*socks_address),
        }
        # XXXX: mv nickname and uuid to user-agent too
        self.headers[
            "Tor-Bandwidth-Scanner-Nickname"
        ] = constants.HTTP_HEADERS["Tor-Bandwidth-Scanner-Nickname"].format(
            kwargs.get("nickname", "")
        )
        self.headers["Tor-Bandwidth-Scanner-UUID"] = constants.HTTP_HEADERS[
            "Tor-Bandwidth-Scanner-UUID"
        ].format(kwargs.get("uuid", ""))
        self.headers["User-Agent"] = constants.HTTP_HEADERS[
            "User-Agent"
        ].format(kwargs.get("tor_version", ""))
        logger.debug(self.headers)



    def get(self, url, **kwargs):
        logger.debug("HTTP GET %s. Verify: %s.", url, self._verify)
        bytes_range = kwargs.pop("range", None)
        # after get, we don't want connection to continue open
        headers = constants.HTTP_GET_HEADERS
        headers["Connection"] = "close"
        if bytes_range:
            # headers.format(bytes_range)
            headers["Range"] = bytes_range
        # Any dictionaries that you pass to a request method will be merged
        #  with the session-level values
        return super().get(
            url,
            headers=headers,
            timeout=self._timeout,
            verify=self._verify,
            allow_redirects=False,
            **kwargs,
        )




    def head(self, url, **kwargs):
        logger.debug("HTTP HEAD %s. Verify: %s.", url, self._verify)
        # By default, HEAD requests don't follow redirections
        return super().head(
            url, timeout=self._timeout, verify=self._verify, **kwargs
        )