Merge branch 'pr/127'

requests_toolbelt/auth/http_proxy_digest.py

@@ -53,12 +53,16 @@ def handle_407(self, r, **kwargs):
         :returns: responses, along with the new response
         """
         if r.status_code == 407 and self.stale_rejects < 2:
-            if "proxy-authenticate" not in r.headers:
+            s_auth = r.headers.get("proxy-authenticate")
+            if s_auth is None:
                 raise IOError(
                     "proxy server violated RFC 7235:"
                     "407 response MUST contain header proxy-authenticate")
+            elif not self._pat.match(s_auth):
+                return r
+
             self.chal = utils.parse_dict_header(
-                self._pat.sub('', r.headers['proxy-authenticate'], count=1))
+                self._pat.sub('', s_auth, count=1))
 
             # if we present the user/passwd and still get rejected
             # http://tools.ietf.org/html/rfc2617#section-3.2.1

tests/test_proxy_digest_auth.py

@@ -1,40 +1,112 @@
 # -*- coding: utf-8 -*-
-"""Test proxy digest authentication
-"""
+"""Test proxy digest authentication."""
 
 import unittest
+import mock
 
 import requests
 from requests_toolbelt.auth import http_proxy_digest
 
 
 class TestProxyDigestAuth(unittest.TestCase):
+    """Tests for the ProxyDigestAuth class."""
+
     def setUp(self):
+        """Set up variables for each test."""
         self.username = "username"
         self.password = "password"
         self.auth = http_proxy_digest.HTTPProxyDigestAuth(
             self.username, self.password
         )
-        self.auth.last_nonce = "bH3FVAAAAAAg74rL3X8AAI3CyBAAAAAA"
-        self.auth.chal = {
-            'nonce': self.auth.last_nonce,
-            'realm': 'testrealm@host.org',
-            'qop': 'auth'
-        }
         self.prepared_request = requests.Request(
             'GET',
             'http://host.org/index.html'
         ).prepare()
 
-    def test_proxy_digest(self):
-        """Test if it will generate Proxy-Authorization header
-        when nonce presents.
+    def test_with_existing_nonce(self):
+        """Test if it will generate Proxy-Auth header when nonce present.
+
         Digest authentication's correctness will not be tested here.
         """
+        self.auth.last_nonce = "bH3FVAAAAAAg74rL3X8AAI3CyBAAAAAA"
+        self.auth.chal = {
+            'nonce': self.auth.last_nonce,
+            'realm': 'testrealm@host.org',
+            'qop': 'auth'
+        }
+
         # prepared_request headers should be clear before calling auth
-        assert not self.prepared_request.headers.get('Proxy-Authorization')
+        assert self.prepared_request.headers.get('Proxy-Authorization') is None
         self.auth(self.prepared_request)
-        assert self.prepared_request.headers.get('Proxy-Authorization')
+        assert self.prepared_request.headers['Proxy-Authorization'] is not None
+
+    def test_no_challenge(self):
+        """Test that a response containing no auth challenge is left alone."""
+        connection = MockConnection()
+        first_response = connection.make_response(self.prepared_request)
+        first_response.status_code = 404
+
+        assert self.auth.last_nonce == ''
+        final_response = self.auth.handle_407(first_response)
+        headers = final_response.request.headers
+        assert self.auth.last_nonce == ''
+        assert first_response is final_response
+        assert headers.get('Proxy-Authorization') is None
+
+    def test_digest_challenge(self):
+        """Test a response with a digest auth challenge causes a new request.
+
+        This ensures that the auth class generates a new request with a
+        Proxy-Authorization header.
+
+        Digest authentication's correctness will not be tested here.
+        """
+        connection = MockConnection()
+        first_response = connection.make_response(self.prepared_request)
+        first_response.status_code = 407
+        first_response.headers['Proxy-Authenticate'] = (
+            'Digest'
+            ' realm="Fake Realm", nonce="oS6WVgAAAABw698CAAAAAHAk/HUAAAAA",'
+            ' qop="auth", stale=false'
+        )
+
+        assert self.auth.last_nonce == ''
+        final_response = self.auth.handle_407(first_response)
+        headers = final_response.request.headers
+        assert self.auth.last_nonce != ''
+        assert first_response is not final_response
+        assert headers.get('Proxy-Authorization') is not None
+
+    def test_ntlm_challenge(self):
+        """Test a response without a Digest auth challenge is left alone."""
+        connection = MockConnection()
+        first_response = connection.make_response(self.prepared_request)
+        first_response.status_code = 407
+        first_response.headers['Proxy-Authenticate'] = 'NTLM'
+
+        assert self.auth.last_nonce == ''
+        final_response = self.auth.handle_407(first_response)
+        headers = final_response.request.headers
+        assert self.auth.last_nonce == ''
+        assert first_response is final_response
+        assert headers.get('Proxy-Authorization') is None
+
+
+class MockConnection(object):
+    """Fake connection object."""
+
+    def send(self, request, **kwargs):
+        """Mock out the send method."""
+        return self.make_response(request)
+
+    def make_response(self, request):
+        """Make a response for us based on the request."""
+        response = requests.Response()
+        response.status_code = 200
+        response.request = request
+        response.raw = mock.MagicMock()
+        response.connection = self
+        return response
 
 if __name__ == '__main__':
     unittest.main()