Session ID middleware for Python

session_id_middleware.py

#!/usr/bin/env python

import base64
from Cookie import SimpleCookie
import hashlib
import hmac
import random
import string


class SessionIdMiddleware(object):
	id_chars = ''.join(('-_', string.digits, string.ascii_letters))

	def __init__(self, application, session_digest_type='sha1', session_environ_key='HTTP_SESSION_ID', session_cookie_name='session', session_cookie_max_age=60*60, session_secret=None):
		self.application = application

		self.environ_key = session_environ_key
		self.cookie_name = session_cookie_name
		self.cookie_max_age = session_cookie_max_age

		self.digest_type = session_digest_type
		self.digest_callable = lambda: hashlib.new(self.digest_type)

		if session_secret is None:
			session_secret = self.digest(self.random_id_chars())
		self.secret = session_secret

	def digest(self, message):
		hash = self.digest_callable()
		hash.update(message)
		return base64.b64encode(hash.digest(), '-_')

	def hmac(self, message):
		return base64.b64encode(hmac.new(self.secret, message, self.digest_callable).digest(), '-_')

	def random_id_chars(self, length=64):
		return ''.join(random.choice(self.id_chars) for i in xrange(0, length))

	def create_session_id(self, environ):
		id = self.digest(self.random_id_chars() + environ['REMOTE_ADDR']) 
		hash = self.hmac(id + environ['REMOTE_ADDR'])
		return id + '/' + hash

	def verify_session_id(self, environ, id):
		if '/' in id:
			id, hash = id.split('/', 1)
			if self.hmac(id + environ['REMOTE_ADDR']) == hash:
				return True
		return False

	def get_session_id(self, environ):
		# Did we get any cookies?
		if 'HTTP_COOKIE' in environ:
			cookies = SimpleCookie(environ['HTTP_COOKIE'])
			# containing our cookie?
			if self.cookie_name in cookies:
				id = str(cookies[self.cookie_name].value)
				# with a valid id?
				if self.verify_session_id(environ, id):
					return id
		
		return self.create_session_id(environ)

	def __call__(self, environ, start_response):
		# Make sure there is a session id in the environment
		if self.environ_key not in environ:
			environ[self.environ_key] = self.get_session_id(environ)

		def session_start_response(status, headers, exc_info=None):
			cookie = SimpleCookie()
			cookie[self.cookie_name] = environ[self.environ_key]
			headers.append(('Set-Cookie', cookie[self.cookie_name].OutputString()))
			return start_response(status, headers, exc_info)

		return self.application(environ, session_start_response)