Implement cookie session

python2/httplib2/__init__.py

@@ -62,6 +62,8 @@
     except (ImportError, AttributeError):
         socks = None
 
+from cookie import CookieJar, NullCookieJar
+
 # Build the appropriate socket wrapper for ssl
 ssl = None
 ssl_SSLError = None
@@ -1603,6 +1605,7 @@ class Http(object):
     - Basic
     - Digest
     - WSSE
+    - Cookies
 
     and more.
     """
@@ -1615,6 +1618,7 @@ def __init__(
         ca_certs=None,
         disable_ssl_certificate_validation=False,
         ssl_version=None,
+        cookie_jar=None,
     ):
         """If 'cache' is a string then it is used as a directory name for
         a disk cache. Otherwise it must be an object that supports the
@@ -1640,6 +1644,11 @@ def __init__(
         not be performed.
 
         By default, ssl.PROTOCOL_SSLv23 will be used for the ssl version.
+
+        cookie_jar is an object compatible with httplib2.CookieJar. It will
+        enable the cookie manager which extracts the cookies from the HTTP
+        response and sets policy matched cookies into the subsequent requests
+        automatically.
         """
         self.proxy_info = proxy_info
         self.ca_certs = ca_certs
@@ -1688,13 +1697,18 @@ def __init__(
         # Keep Authorization: headers on a redirect.
         self.forward_authorization_headers = False
 
+        self.cookie_jar = cookie_jar or NullCookieJar()
+
     def close(self):
         """Close persistent connections, clear sensitive data.
         Not thread-safe, requires external synchronization against concurrent requests.
         """
         existing, self.connections = self.connections, {}
         for _, c in existing.iteritems():
             c.close()
+
+        self.cookie_jar.clear()
+
         self.certificates.clear()
         self.clear_credentials()
 
@@ -1740,6 +1754,19 @@ def clear_credentials(self):
         self.credentials.clear()
         self.authorizations = []
 
+    def _set_cookie_header(self, headers, host, uri, is_https, cookies):
+        # Add one-time cookies provided by request parameter
+        cookie_list = ['{0}={1}'.format(n, v) for n, v in cookies.iteritems()] \
+            if cookies else []
+
+        # Add the cookies matched in cookie jar
+        cookie_header = self.cookie_jar.get_header(host, uri, is_https)
+        if cookie_header:
+            cookie_list.append(cookie_header)
+
+        if cookie_list:
+            headers['cookie'] = '; '.join(cookie_list)
+
     def _conn_request(self, conn, request_uri, method, body, headers):
         i = 0
         seen_bad_status_line = False
@@ -1815,6 +1842,10 @@ def _conn_request(self, conn, request_uri, method, body, headers):
                 if method != "HEAD":
                     content = _decompressContent(response, content)
             break
+
+        self.cookie_jar.extract_header(
+            conn.host, request_uri, response.get('set-cookie'))
+
         return (response, content)
 
     def _request(
@@ -1951,6 +1982,7 @@ def request(
         headers=None,
         redirections=DEFAULT_MAX_REDIRECTS,
         connection_type=None,
+        cookies=None,
     ):
         """ Performs a single HTTP request.
 
@@ -1969,6 +2001,12 @@ def request(
         The maximum number of redirect to follow before raising an
         exception is 'redirections. The default is 5.
 
+        The 'cookies' is a one time cookie list which is valid in this request
+        only. It's a dict of {<name>: <value>}. These cookies will be sent out
+        in the request even the cookie manager is not enabled. If the cookie
+        manager is enabled, these cookies will be sent out with the cookies
+        from cookie jar, but will never be stored into cookie jar.
+
         The return value is a tuple of (response, content), the first
         being and instance of the 'Response' class, the second being
         a string that contains the response entity body.
@@ -2030,6 +2068,8 @@ def request(
             if "range" not in headers and "accept-encoding" not in headers:
                 headers["accept-encoding"] = "gzip, deflate"
 
+            self._set_cookie_header(headers, conn.host, request_uri, scheme == 'https', cookies)
+
             info = email.Message.Message()
             cachekey = None
             cached_value = None
@@ -2284,5 +2324,13 @@ def __init__(self, info):
     def __getattr__(self, name):
         if name == "dict":
             return self
+        elif name == "cookies":
+            # Provide a way to access the cookie info of this reponse
+            if 'set-cookie' not in self:
+                cookies = {}
+            else:
+                cookies = {ck['name']: ck['value'] for ck in CookieJar.parse_iter(self['set-cookie'])}
+            setattr(self, name, cookies)
+            return cookies
         else:
             raise AttributeError(name)