Ticket #4922: cookieagent_composed2.patch

twisted/web/client.py

@@ -1 +1 @@
 # -*- test-case-name: twisted.web.test.test_webclient -*-
 # Copyright (c) 2001-2010 Twisted Matrix Laboratories.
 # See LICENSE for details.
+import urllib2
 
 """
 HTTP client.
@@ -848 +849 @@
                 p.transport.loseConnection()
         self._protocolCache = {}
 
+class CookieAgent(object):
+    """
+    L{CookieAgent} extends the basic L{Agent} to add RFC-compliant
+    handling of HTTP cookies.  Cookies are written to and extracted
+    from a C{cookielib.CookieJar} instance.
+    """
+    def __init__(self, agent, cookieJar, *args, **kwargs):
+        self._agent = agent
+        self.jar = cookieJar
 
+    def request(self, method, uri, headers=None, bodyProducer=None):
+        """
+        Issue a new request.
+
+        @param method: The request method to send.
+        @type method: C{str}
+
+        @param uri: The request URI send.
+        @type uri: C{str}
+
+        @param headers: The request headers to send.  If no I{Host} header is
+            included, one will be added based on the request URI.
+        @type headers: L{Headers}
+
+        @param bodyProducer: An object which will produce the request body or,
+            if the request body is to be empty, L{None}.
+        @type bodyProducer: L{IBodyProducer} provider
+
+        @return: A L{Deferred} which fires with the result of the request (a
+            L{Response} instance), or fails if there is a problem setting up a
+            connection over which to issue the request.  It may also fail with
+            L{SchemeNotSupported} if the scheme of the given URI is not
+            supported.
+        @rtype: L{Deferred}
+        """
+        # setting cookie header explicitly will disable automatic request cookies
+        if not headers.hasHeader('cookie'):
+            last_req = self._urllib2Request(uri)
+            self.jar.add_cookie_header(last_req)
+            cookie_header = last_req.get_header('Cookie', None)
+            if cookie_header is not None:
+                headers.addRawHeader('cookie', cookie_header)        
+        d = self._agent.request(method, uri, headers, bodyProducer)
+        d.addCallback(self._extractCookies, last_req)
+        return d
+    
+    def _urllib2Request(self, uri):
+        """
+        Given a URI, return a urllib2.Request instance for use with cookielib
+        
+        @param uri: The request URI send.
+        @type uri: C{str}
+
+        @return: A C{urllib2.Request} instance for use with cookielib
+        @rtype: C{urllib2.Request}
+        """
+        return urllib2.Request(uri)
+
+    def _extractCookies(self, response, request):
+        """
+        Extract response cookies and store them in the cookie jar
+
+        @param response: A urllib2-style response object
+        @type uri: C{urllib.Response}
+
+        @param response: A urllib2-style response object
+        @type uri: C{urllib.Response}
+
+        @return: A L{Deferred} which fires with the result of the request (a
+            L{Response} instance), or fails if there is a problem setting up a
+            connection over which to issue the request.  It may also fail with
+            L{SchemeNotSupported} if the scheme of the given URI is not
+            supported.
+        @rtype: L{Deferred}
+        """
+        # construct a urllib2-style response object
+        class _Response(object):
+            def info(self):
+                class _Meta(object):
+                    def getheaders(self, name):
+                        return response.headers.getRawHeaders(name.lower(), [])
+                return _Meta()
+        resp = _Response()
+        self.jar.extract_cookies(resp, request)
+        return response
+
 __all__ = [
     'PartialDownloadError',
     'HTTPPageGetter', 'HTTPPageDownloader', 'HTTPClientFactory', 'HTTPDownloader',
     'getPage', 'downloadPage',
 
-    'ResponseDone', 'Response', 'Agent']
+    'ResponseDone', 'Response', 'Agent', 'CookieAgent']