add rearchitect patch from Michael Garvin <gar@bridgedanalytics.com>

Thanks!

1 files changed, 69 insertions, 37 deletions

diff --git a/urllib2_kerberos.py b/urllib2_kerberos.py
index fb06f01..3b4eb6f 100644
--- a/urllib2_kerberos.py
+++ b/urllib2_kerberos.py
@@ -34,25 +34,25 @@ def getLogger():
 
 log = getLogger()
 
-class HTTPKerberosAuthHandler(u2.BaseHandler):
+class AbstractKerberosAuthHandler(u2.AbstractDigestAuthHandler):
     """auth handler for urllib2 that does Kerberos HTTP Negotiate Authentication
     """
 
-    rx = re.compile('(?:.*,)*\s*Negotiate\s*([^,]*),?', re.I)
-    handler_order = 480  # before Digest auth
-
-    def negotiate_value(self, headers):
-        authreq = headers.get('www-authenticate', None)
+    def negotiate_value(self, auth_header, headers):
+        """checks for "Negotiate" in proper auth header
+        """
+        authreq = headers.get(auth_header, None)
 
         if authreq:
-            mo = HTTPKerberosAuthHandler.rx.search(authreq)
+            rx = re.compile('(?:.*,)*\s*Negotiate\s*([^,]*),?', re.I)
+            mo = rx.search(authreq)
             if mo:
                 return mo.group(1)
             else:
                 log.debug("regex failed on: %s" % authreq)
 
         else:
-            log.debug("www-authenticate header not found")
+            log.debug("%s header not found" % auth_header)
 
         return None
 
@@ -60,22 +60,13 @@ class HTTPKerberosAuthHandler(u2.BaseHandler):
         self.retried = 0
         self.context = None
 
-    def generate_request_header(self, req, headers):
-        neg_value = self.negotiate_value(headers)
-        if neg_value is None:
-            self.retried = 0
-            return None
-
-        if self.retried > 0:
-            raise u2.HTTPError(req.get_full_url(), 401, "kerberos negotiate auth failed",
-                               headers, None)
-
+    def generate_request_header(self, req, headers, neg_value):
         self.retried += 1
 
         log.debug("retry count: %d" % self.retried)
 
         log.debug("req.get_host() returned %s" % req.get_host())
-        result, self.context = k.authGSSClientInit("HTTP@%s" % req.get_host())
+        result, self.context = k.authGSSClientInit("HTTP@%s" % req.get_host().rpartition(":")[0])
 
         if result < 1:
             log.warning("authGSSClientInit returned result %d" % result)
@@ -103,7 +94,9 @@ class HTTPKerberosAuthHandler(u2.BaseHandler):
             return None
 
         if k.authGSSClientStep(self.context, neg_value) < 1:
-            log.critical("mutual auth failed: authGSSClientStep returned result %d" % result)
+            #TODO pyflakes flagged this so I commented it out --Gar
+            #log.critical("mutual auth failed: authGSSClientStep returned result %d" % result)
+            pass
 
     def clean_context(self):
         if self.context is not None:
@@ -111,29 +104,68 @@ class HTTPKerberosAuthHandler(u2.BaseHandler):
             k.authGSSClientClean(self.context)
             self.context = None
 
-    def http_error_401(self, req, fp, code, msg, headers):
-        log.debug("inside http_error_401")
-        try:
-            neg_hdr = self.generate_request_header(req, headers)
+    def http_error_auth_reqed(self, auth_header, host, req, headers):
+        neg_value = self.negotiate_values(auth_header, headers) #Check for appropriate auth_header
+        if neg_value is not None:
+            if not self.retried > 0:
+                return self.retry_http_kerberos_auth(req, headers, auth_header, neg_value)
+            else:
+                return None
+        else:
+            self.retried = 0
+
+        def retry_http_kerberos_auth(self, req, headers, auth_header, neg_value):
+            try:
+                neg_hdr = self.generate_request_header(req, headers, neg_value)
+
+                if neg_hdr is None:
+                    log.debug("neg_hdr was None")
+                    return None
 
-            if neg_hdr is None:
-                log.debug("neg_hdr was None")
+                req.add_unredirected_header(self.auth_header, neg_hdr)
+                resp = self.parent.open(req)
+
+                self.authenticate_server(resp.info())
+
+                return resp
+
+            except k.GSSError, e:
+                log.critical("GSSAPI Error: %s/%s" % (e[0][0], e[1][0]))
                 return None
 
-            req.add_unredirected_header('Authorization', neg_hdr)
-            resp = self.parent.open(req)
+            finally:
+                self.clean_context()
+                self.retried = 0
+
+class ProxyKerberosAuthHandler(u2.BaseHandler, AbstractKerberosAuthHandler):
+    """Kerberos Negotiation handler for HTTP proxy auth
+    """
 
-            self.authenticate_server(resp.info())
+    auth_header = 'Proxy-Authorization'
 
-            return resp
+    handler_order = 480 # before Digest auth
 
-        except k.GSSError, e:
-            log.critical("GSSAPI Error: %s/%s" % (e[0][0], e[1][0]))
-            return None
-        
-        finally:
-            self.clean_context()
-            self.retried = 0
+    def http_error_407(self, req, fp, code, msg, headers):
+        log.debug("inside http_error_407")
+        host = req.get_host()
+        retry = self.http_error_auth_reqed('proxy_authenticate', host, req, headers)
+        self.reset_retry_count()
+        return retry
+
+class HTTPKerberosAuthHandler(u2.BaseHandler, AbstractKerberosAuthHandler):
+    """Kerberos Negotiation handler for HTTP auth
+    """
+
+    auth_header = 'Authorization'
+
+    handler_order = 480 # before Digest auth
+
+    def http_error_401(self, req, fp, code, msg, headers):
+        log.debug("inside http_error_401")
+        host = req.get_host()
+        retry = self.http_error_auth_reqed('www-authenticate', host, req, headers)
+        self.reset_retry_count()
+        return retry
 
 def test():
     log.setLevel(logging.DEBUG)