From 30d699b64fbf298881c0de88cd6b727b1972538a Mon Sep 17 00:00:00 2001
From: Kenneth Reitz <me@kennethreitz.com>
Date: Sat, 23 Jul 2011 21:36:24 -0400
Subject: [PATCH] Forced Basic Auth for 404'd resources

---
 requests/monkeys.py | 61 +++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 59 insertions(+), 2 deletions(-)

diff --git a/requests/monkeys.py b/requests/monkeys.py
index 41cd370..909626c 100644
--- a/requests/monkeys.py
+++ b/requests/monkeys.py
@@ -9,7 +9,7 @@ Urllib2 Monkey patches.
 """
 
 import urllib2
-
+import re
 
 class Request(urllib2.Request):
     """Hidden wrapper around the urllib2.Request object. Allows for manual
@@ -26,8 +26,9 @@ class Request(urllib2.Request):
 
         return urllib2.Request.get_method(self)
 
-class HTTPRedirectHandler(urllib2.HTTPRedirectHandler):
 
+class HTTPRedirectHandler(urllib2.HTTPRedirectHandler):
+    """HTTP Redirect handler."""
     def http_error_301(self, req, fp, code, msg, headers):
         pass
 
@@ -36,17 +37,20 @@ class HTTPRedirectHandler(urllib2.HTTPRedirectHandler):
 
 
 class HTTPBasicAuthHandler(urllib2.HTTPBasicAuthHandler):
+    """HTTP Basic Auth Handler with authentication loop fixes."""
 
     def __init__(self, *args, **kwargs):
         urllib2.HTTPBasicAuthHandler.__init__(self, *args, **kwargs)
         self.retried_req = None
 
+
     def reset_retry_count(self):
         # Python 2.6.5 will call this on 401 or 407 errors and thus loop
         # forever. We disable reset_retry_count completely and reset in
         # http_error_auth_reqed instead.
         pass
 
+
     def http_error_auth_reqed(self, auth_header, host, req, headers):
         # Reset the retry counter once for each request.
         if req is not self.retried_req:
@@ -59,6 +63,59 @@ class HTTPBasicAuthHandler(urllib2.HTTPBasicAuthHandler):
 
 
 
+class HTTPForcedBasicAuthHandler(HTTPBasicAuthHandler):
+    """HTTP Basic Auth Handler with forced Authentication."""
+
+    auth_header = 'Authorization'
+    rx = re.compile('(?:.*,)*[ \t]*([^ \t]+)[ \t]+'
+                    'realm=(["\'])(.*?)\\2', re.I)
+
+    def __init__(self,  *args, **kwargs):
+        HTTPBasicAuthHandler.__init__(self, *args, **kwargs)
+
+
+    def http_error_401(self, req, fp, code, msg, headers):
+        url = req.get_full_url()
+        response = self._http_error_auth_reqed('www-authenticate', url, req, headers)
+        self.reset_retry_count()
+        return response
+
+    http_error_404 = http_error_401
+
+
+    def _http_error_auth_reqed(self, authreq, host, req, headers):
+
+        authreq = headers.get(authreq, None)
+
+        if self.retried > 5:
+            # retry sending the username:password 5 times before failing.
+            raise urllib2.HTTPError(req.get_full_url(), 401, "basic auth failed",
+                            headers, None)
+        else:
+            self.retried += 1
+
+        if authreq:
+
+            mo = self.rx.search(authreq)
+
+            if mo:
+                scheme, quote, realm = mo.groups()
+
+                if scheme.lower() == 'basic':
+                    response = self.retry_http_basic_auth(host, req, realm)
+
+                    if response and response.code not in (401, 404):
+                        self.retried = 0
+                    return response
+        else:
+            response = self.retry_http_basic_auth(host, req, 'Realm')
+
+            if response and response.code not in (401, 404):
+                self.retried = 0
+            return response
+
+
+
 class HTTPDigestAuthHandler(urllib2.HTTPDigestAuthHandler):
 
     def __init__(self, *args, **kwargs):
-- 
2.34.1