codekingpro/portable-devtools

1# -*- coding: utf-8 -*-2"""The module containing HTTPProxyDigestAuth."""3import re4 5from requests import cookies, utils6 7from . import _digest_auth_compat as auth8 9 10class HTTPProxyDigestAuth(auth.HTTPDigestAuth):11    """HTTP digest authentication between proxy12 13    :param stale_rejects: The number of rejects indicate that:14        the client may wish to simply retry the request15        with a new encrypted response, without reprompting the user for a16        new username and password. i.e., retry build_digest_header17    :type stale_rejects: int18    """19    _pat = re.compile(r'digest ', flags=re.IGNORECASE)20 21    def __init__(self, *args, **kwargs):22        super(HTTPProxyDigestAuth, self).__init__(*args, **kwargs)23        self.stale_rejects = 024 25        self.init_per_thread_state()26 27    @property28    def stale_rejects(self):29        thread_local = getattr(self, '_thread_local', None)30        if thread_local is None:31            return self._stale_rejects32        return thread_local.stale_rejects33 34    @stale_rejects.setter35    def stale_rejects(self, value):36        thread_local = getattr(self, '_thread_local', None)37        if thread_local is None:38            self._stale_rejects = value39        else:40            thread_local.stale_rejects = value41 42    def init_per_thread_state(self):43        try:44            super(HTTPProxyDigestAuth, self).init_per_thread_state()45        except AttributeError:46            # If we're not on requests 2.8.0+ this method does not exist47            pass48 49    def handle_407(self, r, **kwargs):50        """Handle HTTP 407 only once, otherwise give up51 52        :param r: current response53        :returns: responses, along with the new response54        """55        if r.status_code == 407 and self.stale_rejects < 2:56            s_auth = r.headers.get("proxy-authenticate")57            if s_auth is None:58                raise IOError(59                    "proxy server violated RFC 7235:"60                    "407 response MUST contain header proxy-authenticate")61            elif not self._pat.match(s_auth):62                return r63 64            self.chal = utils.parse_dict_header(65                self._pat.sub('', s_auth, count=1))66 67            # if we present the user/passwd and still get rejected68            # https://tools.ietf.org/html/rfc2617#section-3.2.169            if ('Proxy-Authorization' in r.request.headers and70                    'stale' in self.chal):71                if self.chal['stale'].lower() == 'true':  # try again72                    self.stale_rejects += 173                # wrong user/passwd74                elif self.chal['stale'].lower() == 'false':75                    raise IOError("User or password is invalid")76 77            # Consume content and release the original connection78            # to allow our new request to reuse the same one.79            r.content80            r.close()81            prep = r.request.copy()82            cookies.extract_cookies_to_jar(prep._cookies, r.request, r.raw)83            prep.prepare_cookies(prep._cookies)84 85            prep.headers['Proxy-Authorization'] = self.build_digest_header(86                prep.method, prep.url)87            _r = r.connection.send(prep, **kwargs)88            _r.history.append(r)89            _r.request = prep90 91            return _r92        else:  # give up authenticate93            return r94 95    def __call__(self, r):96        self.init_per_thread_state()97        # if we have nonce, then just use it, otherwise server will tell us98        if self.last_nonce:99            r.headers['Proxy-Authorization'] = self.build_digest_header(100                r.method, r.url101            )102        r.register_hook('response', self.handle_407)103        return r104