/
clientsessioncache.py
55 lines (45 loc) · 2.18 KB
/
clientsessioncache.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
from toto.session import *
from copy import copy
from base64 import urlsafe_b64encode, urlsafe_b64decode
class AESCipher(object):
'''A convenient cipher implementation for AES encryption and decryption.
Create a new ``AESCipher`` with the given ``key`` and ``iv`` that wraps
``Crypto.AES`` but is reusable and thread safe. For convenience, both
the ``key`` and ``iv`` may be provided as one string, in which case the
last ``AES.block_size`` (16) bytes will be used for ``iv``.
'''
def __init__(self, key, iv=None):
from Crypto.Cipher import AES
self.block_size = AES.block_size
if not iv:
iv = key[-self.block_size:]
key = key[:-self.block_size]
self.aes = lambda:AES.new(key, AES.MODE_CBC, iv)
def encrypt(self, data):
diff = self.block_size - (len(data) % self.block_size)
return self.aes().encrypt(data + chr(diff) * diff)
def decrypt(self, data):
decrypted = self.aes().decrypt(data)
return decrypted[:-ord(decrypted[-1])]
class ClientCache(TotoSessionCache):
'''A ``TotoSessionCache`` implementation that stores all session data with the
client. Depending on use, the session may be sent as a header or cookie.
``ClientCache`` works by storing the encrypted session state in
the session ID and decrypting it on each request. When using this method,
it is important to keep session state small as it can add significant
overhead to each request.
``cipher`` will be used to encrypt and decrypt the session data. It should
be identical between all servers in a deployment to allow proper request
balancing. ``cipher`` is expected to implement ``encrypt(data)`` and
the reverse ``decrypt(data)`` both accepting and returning ``str`` objects.
'''
def __init__(self, cipher):
self.cipher = cipher
def store_session(self, session_data):
persisted_data = copy(session_data)
del persisted_data['session_id']
return urlsafe_b64encode(self.cipher.encrypt(TotoSession.dumps(session_data)))
def load_session(self, session_id):
session_data = TotoSession.loads(self.cipher.decrypt(urlsafe_b64decode(session_id)))
session_data['session_id'] = session_id
return session_data