Wednesday, December 29, 2010

Setting cookies from Google App Engine

I was disappointed to realize that the google.appengine.ext.webapp.Response class, while in many ways similar to the WebOb Response class, doesn't support the latter's set_cookie method.

I thought about using python's Cookie module directly, but but the WebOb method looked much more convenient.  So I decided the smartest thing to do was to just get the source code for the cookie-related operations from WebOb and adapt it to Google App Engine.  Feel free to use this little hack in your projects.

Example Usage:
   class FooHandler(webapp.RequestHandler):  
    def get(self):  
     cookies = Cookies(self.response)  
     cookies.add_cookie('foo', 'bar', max_age=360, path='/',
                        domain='example.org', secure=True) 


Source Code (cookies.py):
 from Cookie import BaseCookie  
 from datetime import datetime, date, timedelta  
 import time  
   
 # Implements WebOb's set_cookie, unset_cookie, and delete_cookie methods which GAE's webapp framework is missing  
 # modified from original WebOB code found at http://python-webob.sourcearchive.com/documentation/0.9.8-1/response_8py-source.html  
   
 # author: Alex Epshteyn ( http://myprogblog.blogspot.com/ )  
   
 # Original license from WebOb framework, which also applies to this file:  
   
 # Copyright (c) 2007 Ian Bicking and Contributors  
 # Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:  
 # The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.  
 # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.  
   
 class Cookies(object):  
  """  
  Wraps the 'headers' field of a google.appengine.ext.webapp.Response instance with cookie methods probided for a WebOb response object.  
  Example usage:  
   class FooHandler(webapp.RequestHandler):  
    def get(self):  
     cookies = Cookies(self.response)  
     cookies.add_cookie('foo', 'bar')  
  """  
   
  def __init__(self, response):  
   """  
   Takes a GAE webapp Response object (i.e. the value of the 'response' field of a google.appengine.ext.webapp.RequestHandler instance)  
   """  
   self.headers = response.headers  
   
  def set_cookie(self, key, value='', max_age=None,  
          path='/', domain=None, secure=None, httponly=False,  
          version=None, comment=None, expires=None, overwrite=False):  
    """  
    Set (add) a cookie for the response  
    """  
    if isinstance(value, unicode) and self.charset is not None:  
      value = '"%s"' % value.encode(self.charset)  
    if overwrite:  
      self.unset_cookie(key, strict=False)  
    cookies = BaseCookie()  
    cookies[key] = value  
    if isinstance(max_age, timedelta):  
      max_age = max_age.seconds + max_age.days*24*60*60  
    if max_age is not None and expires is None:  
      expires = datetime.utcnow() + timedelta(seconds=max_age)  
    if isinstance(expires, timedelta):  
      expires = datetime.utcnow() + expires  
    if isinstance(expires, datetime):  
      expires = '"'+_serialize_cookie_date(expires)+'"'  
    for var_name, var_value in [  
      ('max-age', max_age),  
      ('path', path),  
      ('domain', domain),  
      ('secure', secure),  
      ('HttpOnly', httponly),  
      ('version', version),  
      ('comment', comment),  
      ('expires', expires),  
    ]:  
      if var_value is not None and var_value is not False:  
        cookies[key][var_name] = str(var_value)  
    self._add_cookie(cookies)  
   
  def _add_cookie(self, cookie):  
    if not isinstance(cookie, str):  
      cookie = cookie.output(header='').lstrip()  
      if cookie.endswith(';'):  
        # Python 2.4 adds a trailing ; to the end, strip it to be  
        # consistent with 2.5  
        cookie = cookie[:-1]  
    if cookie:  
      self.headers.add_header('Set-Cookie', cookie)  
   
   
  def delete_cookie(self, key, path='/', domain=None):  
    """  
    Delete a cookie from the client. Note that path and domain must match  
    how the cookie was originally set.  
   
    This sets the cookie to the empty string, and max_age=0 so  
    that it should expire immediately.  
    """  
    self.set_cookie(key, '', path=path, domain=domain,  
            max_age=0, expires=timedelta(days=-5))  
   
  def unset_cookie(self, key, strict=True):  
    """  
    Unset a cookie with the given name (remove it from the  
    response). If there are multiple cookies (e.g., two cookies  
    with the same name and different paths or domains), all such  
    cookies will be deleted.  
    """  
    existing = self.headers.get_all('Set-Cookie')  
    if not existing:  
      if not strict:  
        return  
      raise KeyError("No cookies at all have been set")  
    del self.headers['Set-Cookie']  
    found = False  
    for header in existing:  
      cookies = BaseCookie()  
      cookies.load(header)  
      if key in cookies:  
        found = True  
        del cookies[key]  
        self._add_cookie(cookies)  
      else:  
        # this branching is required because Cookie.Morsel.output()  
        # strips quotes from expires= parameter, so better use  
        # it as is, if it hasn't changed  
        self._add_cookie(header)  
    if strict and not found:  
      raise KeyError(  
        "No cookie has been set with the name %r" % key)  
   
 # this function is from WebOb's datetime_utils.py (http://python-webob.sourcearchive.com/documentation/0.9.8-1/datetime__utils_8py-source.html)  
 def _serialize_cookie_date(dt):  
   if dt is None:  
     return None  
   if isinstance(dt, unicode):  
     dt = dt.encode('ascii')  
   if isinstance(dt, timedelta):  
     dt = datetime.now() + dt  
   if isinstance(dt, (datetime, date)):  
     dt = dt.timetuple()  
   return time.strftime('%a, %d-%b-%Y %H:%M:%S GMT', dt)