Add method for rebuilding proxy configuration.
authorCory Benfield <lukasaoz@gmail.com>
Wed, 12 Mar 2014 19:22:40 +0000 (19:22 +0000)
committerCory Benfield <lukasaoz@gmail.com>
Wed, 12 Mar 2014 19:22:40 +0000 (19:22 +0000)
This includes auth.

requests/sessions.py

index 4c24984eb781f2ceae8f7ac1ef498e0c24e0498b..db2fca392b3e768aecf667fe4c90e61d757da566 100644 (file)
@@ -12,6 +12,7 @@ import os
 from collections import Mapping
 from datetime import datetime
 
+from .auth import _basic_auth_str
 from .compat import cookielib, OrderedDict, urljoin, urlparse, builtin_str
 from .cookies import (
     cookiejar_from_dict, extract_cookies_to_jar, RequestsCookieJar, merge_cookies)
@@ -23,7 +24,10 @@ from .structures import CaseInsensitiveDict
 
 from .adapters import HTTPAdapter
 
-from .utils import requote_uri, get_environ_proxies, get_netrc_auth
+from .utils import (
+    requote_uri, get_environ_proxies, get_netrc_auth, should_bypass_proxies,
+    get_auth_from_url
+)
 
 from .status_codes import codes
 
@@ -154,6 +158,8 @@ class SessionRedirectMixin(object):
             prepared_request._cookies.update(self.cookies)
             prepared_request.prepare_cookies(prepared_request._cookies)
 
+            # Rebuild auth and proxy information.
+            proxies = self.rebuild_proxies(prepared_request, proxies)
             self.rebuild_auth(prepared_request, resp)
 
             resp = self.send(
@@ -196,6 +202,50 @@ class SessionRedirectMixin(object):
 
         return
 
+    def rebuild_proxies(self, prepared_request, proxies):
+        """
+        This method re-evaluates the proxy configuration by considering the
+        environment variables. If we are redirected to a URL covered by
+        NO_PROXY, we strip the proxy configuration. Otherwise, we set missing
+        proxy keys for this URL (in case they were stripped by a previous
+        redirect).
+
+        This method also replaces the Proxy-Authorization header where
+        necessary.
+        """
+        headers = prepared_request.headers
+        url = prepared_request.url
+        new_proxies = {}
+
+        # Consider proxies. First evaluate the new proxy config. If we are
+        # being redirected to a host on the NO_PROXY list then we want to
+        # remove the proxy dictionary entirely. Otherwise, if there's a relevant
+        # environment proxy, set it if we don't already have a proxy to go to.
+        if not should_bypass_proxies(url):
+            environ_proxies = get_environ_proxies(url)
+            scheme = urlparse(url).scheme
+
+            try:
+                new_proxies.setdefault(scheme, environ_proxies[scheme])
+            except KeyError:
+                pass
+
+        # If there's a proxy-authorization header present, remove it, then add
+        # a new one (potentially re-adding the one we just removed).
+        if 'Proxy-Authorization' in headers:
+            del headers['Proxy-Authorization']
+
+        try:
+            username, password = get_auth_from_url(new_proxies[scheme])
+            if username and password:
+                headers['Proxy-Authorization'] = _basic_auth_str(
+                    username, password
+                )
+        except KeyError:
+            pass
+
+        return new_proxies
+
 
 class Session(SessionRedirectMixin):
     """A Requests session.