Source code for tanium_kit.wequests

from __future__ import absolute_import, division, print_function, unicode_literals

import datetime

from . import log_tools


[docs]class WequestsError(Exception):
    pass


[docs]class Wequests(object):

    DEFAULTS = {}
    DEFAULTS["connect_secs"] = 5
    DEFAULTS["response_secs"] = 15
    DEFAULTS["method"] = "get"
    DEFAULTS["headers"] = {}
    DEFAULTS["proxies"] = {}
    DEFAULTS["http_proxy"] = ""
    DEFAULTS["https_proxy"] = ""
    DEFAULTS["data"] = None
    DEFAULTS["verify"] = False
    DEFAULTS["params"] = None
    DEFAULTS["cookies"] = None
    DEFAULTS["files"] = None
    DEFAULTS["allow_redirects"] = True
    DEFAULTS["hooks"] = None
    DEFAULTS["stream"] = None
    DEFAULTS["cert"] = None
    DEFAULTS["json"] = None
    DEFAULTS["auth_failures"] = [401, 403]

    USE_SESSION = True
    LOUD = True
    REQUESTS_PKG = None
    KWARGS = {}
    _RS = None

    def __init__(self, **kwargs):
        import requests  # EXTERNAL DEPENDENCY
        self.REQUESTS_PKG = requests
        self.KWARGS = kwargs

        self.LOUD = kwargs.get("loud", False)
        self.USE_SESSION = kwargs.get("use_session", True)

        for k, v in self.DEFAULTS.items():
            if k in kwargs:
                self.DEFAULTS[k] = kwargs[k]

        if self.USE_SESSION:
            self._RS = self.REQUESTS_PKG.Session()
        else:
            self._RS = self.REQUESTS_PKG

        if self.LOUD:
            try:
                self.REQUESTS_PKG.packages.urllib3.disable_warnings()
            except:
                pass

        log_tools.shutup_requests(loud=self.LOUD)

[docs]    def request(self, url, **kwargs):
        # get our request args with defaults from self
        supplied_headers = kwargs.get("headers", self.DEFAULTS["headers"])
        connect_secs = kwargs.get("connect_secs", self.DEFAULTS["connect_secs"])
        response_secs = kwargs.get("response_secs", self.DEFAULTS["response_secs"])
        method = kwargs.get("method", self.DEFAULTS["method"]).lower()
        proxies = kwargs.get("proxies", self.DEFAULTS["proxies"])
        http_proxy = kwargs.get("http_proxy", self.DEFAULTS["http_proxy"])
        https_proxy = kwargs.get("https_proxy", self.DEFAULTS["https_proxy"])
        verify = kwargs.get("verify", self.DEFAULTS["verify"])
        data = kwargs.get("data", self.DEFAULTS["data"])
        params = kwargs.get("params", self.DEFAULTS["params"])
        cookies = kwargs.get("cookies", self.DEFAULTS["cookies"])
        files = kwargs.get("files", self.DEFAULTS["files"])
        allow_redirects = kwargs.get("allow_redirects", self.DEFAULTS["allow_redirects"])
        hooks = kwargs.get("hooks", self.DEFAULTS["hooks"])
        stream = kwargs.get("stream", self.DEFAULTS["stream"])
        cert = kwargs.get("cert", self.DEFAULTS["cert"])
        json = kwargs.get("json", self.DEFAULTS["json"])

        # construct our headers
        headers = {}
        headers.update(self.DEFAULTS["headers"])
        headers.update(supplied_headers)

        # create our request args
        req_args = {}
        req_args["url"] = url
        req_args["method"] = method
        req_args["headers"] = headers
        req_args["timeout"] = (connect_secs, response_secs)
        req_args["verify"] = verify
        req_args["params"] = params
        req_args["cookies"] = cookies
        req_args["files"] = files
        req_args["allow_redirects"] = allow_redirects
        req_args["hooks"] = hooks
        req_args["stream"] = stream
        req_args["cert"] = cert
        req_args["json"] = json
        req_args["data"] = data
        req_args["proxies"] = proxies

        if http_proxy:
            req_args["proxies"]["http"] = http_proxy

        if https_proxy:
            req_args["proxies"]["https"] = https_proxy

        err = "HTTP {} to URL {} failed: {}"

        sent = datetime.datetime.utcnow()

        try:
            ret = self._RS.request(**req_args)
        except Exception as e:
            raise WequestsError(err.format(method, url, e))

        received = datetime.datetime.utcnow()
        ret.http_sent = sent
        ret.http_received = received
        ret.http_req_args = req_args
        ret.http_kwargs = kwargs

        if ret.status_code in self.DEFAULTS["auth_failures"]:
            e = "Response code {} is one of authorization failed codes {}"
            e = e.format(ret.status_code, self.DEFAULTS["auth_failures"])
            raise WequestsError(err.format(method, url, e))

        if not ret.ok:
            e = "Response is not ok, reason: {}"
            e = e.format(ret.reason)
            raise WequestsError(err.format(method, url, e))
        return ret