Ticket #4922: cookieagent.patch

File cookieagent.patch, 5.0 KB (added by gabrtv, 3 years ago)

CookieAgent implementation

  • twisted/web/client.py

     
    11# -*- test-case-name: twisted.web.test.test_webclient -*- 
    22# Copyright (c) 2001-2010 Twisted Matrix Laboratories. 
    33# See LICENSE for details. 
     4import urllib2 
    45 
    56""" 
    67HTTP client. 
     
    848849                p.transport.loseConnection() 
    849850        self._protocolCache = {} 
    850851 
     852class CookieAgent(Agent): 
     853    """ 
     854    L{CookieAgent} extends the basic L{Agent} to add RFC-compliant 
     855    handling of HTTP cookies.  Cookies are written to and extracted 
     856    from a C{cookielib.CookieJar} instance. 
     857    """ 
     858    def __init__(self, cookieJar, *args, **kwargs): 
     859        Agent.__init__(self, reactor, *args, **kwargs) 
     860        self.jar = cookieJar 
    851861 
     862    def request(self, method, uri, headers=None, bodyProducer=None): 
     863        """ 
     864        Issue a new request. 
     865 
     866        @param method: The request method to send. 
     867        @type method: C{str} 
     868 
     869        @param uri: The request URI send. 
     870        @type uri: C{str} 
     871 
     872        @param headers: The request headers to send.  If no I{Host} header is 
     873            included, one will be added based on the request URI. 
     874        @type headers: L{Headers} 
     875 
     876        @param bodyProducer: An object which will produce the request body or, 
     877            if the request body is to be empty, L{None}. 
     878        @type bodyProducer: L{IBodyProducer} provider 
     879 
     880        @return: A L{Deferred} which fires with the result of the request (a 
     881            L{Response} instance), or fails if there is a problem setting up a 
     882            connection over which to issue the request.  It may also fail with 
     883            L{SchemeNotSupported} if the scheme of the given URI is not 
     884            supported. 
     885        @rtype: L{Deferred} 
     886        """ 
     887        scheme, host, port, path = _parse(uri) 
     888        if headers is None: 
     889            headers = Headers() 
     890        if not headers.hasHeader('host'): 
     891            # This is a lot of copying.  It might be nice if there were a bit 
     892            # less. 
     893            headers = Headers(dict(headers.getAllRawHeaders())) 
     894            headers.addRawHeader( 
     895                'host', self._computeHostValue(scheme, host, port)) 
     896        # setting cookie header explicitly will disable automatic request cookies 
     897        if not headers.hasHeader('cookie'): 
     898            last_req = self._urllib2Request(uri) 
     899            self.jar.add_cookie_header(last_req) 
     900            cookie_header = last_req.get_header('Cookie', None) 
     901            if cookie_header is not None: 
     902                headers.addRawHeader('cookie', cookie_header) 
     903        if self.persistent: 
     904            sem = self._semaphores.get((scheme, host, port)) 
     905            if sem is None: 
     906                sem = defer.DeferredSemaphore(self.maxConnections) 
     907                self._semaphores[scheme, host, port] = sem 
     908            d = sem.run(self._request, method, scheme, host, port, path, 
     909                           headers, bodyProducer) 
     910            d.addCallback(self._extractCookies, last_req) 
     911            return d 
     912        else: 
     913            d = self._request( 
     914                method, scheme, host, port, path, headers, bodyProducer) 
     915            d.addCallback(self._extractCookies) 
     916            return d 
     917     
     918    def _urllib2Request(self, uri): 
     919        """ 
     920        Given a URI, return a urllib2.Request instance for use with cookielib 
     921         
     922        @param uri: The request URI send. 
     923        @type uri: C{str} 
     924 
     925        @return: A C{urllib2.Request} instance for use with cookielib 
     926        @rtype: C{urllib2.Request} 
     927        """ 
     928        return urllib2.Request(uri) 
     929 
     930    def _extractCookies(self, response, request): 
     931        """ 
     932        Extract response cookies and store them in the cookie jar 
     933 
     934        @param response: A urllib2-style response object 
     935        @type uri: C{urllib.Response} 
     936 
     937        @param response: A urllib2-style response object 
     938        @type uri: C{urllib.Response} 
     939 
     940        @return: A L{Deferred} which fires with the result of the request (a 
     941            L{Response} instance), or fails if there is a problem setting up a 
     942            connection over which to issue the request.  It may also fail with 
     943            L{SchemeNotSupported} if the scheme of the given URI is not 
     944            supported. 
     945        @rtype: L{Deferred} 
     946        """ 
     947        # construct a urllib2-style response object 
     948        class _Response(object): 
     949            def info(self): 
     950                class _Meta(object): 
     951                    def getheaders(self, name): 
     952                        return response.headers.getRawHeaders(name.lower(), []) 
     953                return _Meta() 
     954        resp = _Response() 
     955        self.jar.extract_cookies(resp, request) 
     956        return response 
     957 
    852958__all__ = [ 
    853959    'PartialDownloadError', 
    854960    'HTTPPageGetter', 'HTTPPageDownloader', 'HTTPClientFactory', 'HTTPDownloader', 
    855961    'getPage', 'downloadPage', 
    856962 
    857     'ResponseDone', 'Response', 'Agent'] 
     963    'ResponseDone', 'Response', 'Agent', 'CookieAgent']