Permalink
Browse files

Completed session attributes and tests.

  • Loading branch information...
1 parent fb9e70a commit ba10e957d96b9e8c80ed5c796e3d89110584ef3f Zaur Nasibov committed with Oct 5, 2012
View
@@ -13,6 +13,8 @@
'projects'))
WORKER_SCRIPT = '/static/js/kaylee/klworker.js'
+SECRET_KEY = '1234ABCabc!@{}xyz&%*'
+
## User applications ##
#######################
View
1 demo/run.py 100644 → 100755
@@ -1,3 +1,4 @@
+#!/usr/bin/env python
# -*- coding: utf-8 -*-
import os
import sys
View
@@ -249,6 +249,7 @@ def to_dict(self):
def SECRET_KEY(self):
if self._secret_key is None:
raise KayleeError('SECRET_KEY is not defined.')
+ return self._secret_key
class Applications(object):
"""A container for active Kaylee applications.
View
@@ -9,13 +9,15 @@
:license: MIT, see LICENSE for more details.
"""
+import random
import cPickle as pickle
from base64 import b64encode, b64decode
from hmac import new as hmac
from hashlib import sha1, sha256
from abc import abstractmethod
from copy import copy
from functools import wraps
+from copy import deepcopy
from Crypto.Cipher import AES
@@ -218,45 +220,64 @@ def serialize(self, attributes = None):
if attributes is None:
attributes = self.serializable
- res = { attr : getattr(self, attr) for attr in attributes }
+ res = { attr : getattr(self, attr) for attr in attributes
+ if not attr.startswith('#') }
# process session attributes, if any
- sess_attributes = [attr for attr in attributes if attr.startswith('#')]
+ sess_attributes = [attr[1:] for attr in attributes if attr.startswith('#')]
if len(sess_attributes) > 0:
res['__kaylee_task_session__'] = self.encrypt(sess_attributes)
return res
def encrypt(self, attributes):
from . import kl
mac = hmac(kl.config.SECRET_KEY, None, sha1)
- encryption_key = hashlib.sha256(kl.config.SECRET_KEY).digest()
- encryptor = AES.new(encryption_key, AES.MODE_CBC)
+ encryption_key = sha256(kl.config.SECRET_KEY).digest()
- result = []
- for attr, value in sorted(attributes):
- result.append(self._encrypt_attr(attr, value, encryptor))
+ iv = ''.join(chr(random.randint(0, 0xFF)) for i in xrange(16))
+ encryptor = AES.new(encryption_key, AES.MODE_CBC, iv)
+
+ b64_iv = b64encode(iv)
+ result = [b64_iv] # store initialization vector
+ for attr in sorted(attributes):
+ value = getattr(self, attr)
+ result.append(Task._encrypt_attr(attr, value, encryptor))
mac.update('|' + result[-1])
return '{}?{}'.format(b64encode(mac.digest()), '&'.join(result))
- @classmethod
- def decrypt(d):
+ @staticmethod
+ def deserialize(d):
+ result = {}
+
+ if isinstance(d, dict):
+ s = d['__kaylee_task_session__']
+ result = deepcopy(d)
+ del result['__kaylee_task_session__']
+ else:
+ s = d
+
from . import kl
- base64_hash, data = string.split('?', 1)
+ base64_hash, data = s.split('?', 1)
mac = hmac(kl.config.SECRET_KEY, None, sha1)
- decryption_key = hashlib.sha256(kl.config.SECRET_KEY).digest()
- decryptor = AES.new(decryption_key, AES.MODE_CBC)
+ iv, data = data.split('&', 1)
+ iv = b64decode(iv)
+
+ decryption_key = sha256(kl.config.SECRET_KEY).digest()
+ decryptor = AES.new(decryption_key, AES.MODE_CBC, iv)
- result = {}
for item in data.split('&'):
mac.update('|' + item)
- attr, val = self._decrypt_attr(item, decryptor)
+ attr, val = Task._decrypt_attr(item, decryptor)
result[attr] = val
- safe_str_cmp(client_hash, mac.digest()):
+ if b64decode(base64_hash) == mac.digest():
+ return result
+ else:
+ raise KayleeError('Session attribute signature verification failed.')
- @classmethod
+ @staticmethod
def _encrypt_attr(attr, val, encryptor):
BLOCK_SIZE = 32
PADDING = ' '
@@ -269,10 +290,10 @@ def _encrypt_attr(attr, val, encryptor):
val = b64encode(val)
return val
- @classmethod
+ @staticmethod
def _decrypt_attr(data, decryptor):
data = b64decode(data)
- data = decryptor.decrypt(data).rstript(' ')
+ data = decryptor.decrypt(data).rstrip(' ')
attr, val = data.split('=', 1)
val = pickle.loads(val)
return attr, val
@@ -10,16 +10,16 @@
'projects')
sys.path.insert(0, PROJECTS_DIR)
+class KayleeTest(unittest.TestCase):
+ """Base class for all Kaylee tests."""
+
def load_tests(test_cases):
suite = unittest.TestSuite()
for tcase in test_cases:
loaded_suite = unittest.defaultTestLoader.loadTestsFromTestCase(tcase)
suite.addTest(loaded_suite)
return suite
-class KayleeTest(unittest.TestCase):
- """Base class for all Kaylee tests."""
-
def suite():
"""A testsuite that has all Kaylee tests. You can use this
@@ -0,0 +1,87 @@
+import unittest
+import kaylee
+from kaylee.testsuite import KayleeTest, load_tests
+from kaylee.errors import KayleeError
+
+class TaskSimpleFields(kaylee.Task):
+ serializable = ['f1', 'f2']
+
+ def __init__(self, task_id, f1, f2):
+ super(TaskSimpleFields, self).__init__(task_id)
+ self.f1 = f1
+ self.f2 = f2
+
+
+class TaskSessionFields(kaylee.Task):
+ serializable = ['#f1', '#f2']
+
+ def __init__(self, task_id, f1, f2):
+ super(TaskSessionFields, self).__init__(task_id)
+ self.f1 = f1
+ self.f2 = f2
+
+
+class TaskTests(KayleeTest):
+ def setUp(self):
+ pass
+
+ def test_simple_serialize(self):
+ t1 = TaskSimpleFields(1, 10, 'someval')
+ d = t1.serialize()
+ self.assertEqual(d['id'], '1')
+ self.assertEqual(d['f1'], 10)
+ self.assertEqual(d['f2'], 'someval')
+
+ def test_secret_key_config_check(self):
+ # Checks if KayleeError raised when SECRET_KEY is not defined
+ import test_config
+ from kaylee import setup, kl
+ secret_key = test_config.SECRET_KEY
+
+ del test_config.SECRET_KEY
+ setup(test_config)
+
+ t1 = TaskSessionFields(1, 10, 'someval')
+ self.assertRaises(KayleeError, t1.serialize)
+
+ def test_session_serialize(self):
+ import test_config
+ from kaylee import setup, kl
+ setup(test_config)
+
+ t1 = TaskSessionFields(1, 10, 'someval')
+ d = t1.serialize()
+ self.assertIn('__kaylee_task_session__', d)
+ val = d['__kaylee_task_session__']
+ self.assertGreater(len(val), 0)
+ self.assertEqual(val.count('&'), 2)
+
+ def test_session_deserialize(self):
+ import test_config
+ from kaylee import setup, kl
+ setup(test_config)
+
+ t1 = TaskSessionFields(1, 10, 'someval')
+ d1 = t1.serialize()
+
+ # check if deserialization works in general
+ d2 = TaskSessionFields.deserialize(d1)
+ self.assertEqual(len(d2), 3)
+ self.assertEqual(d2['f1'], t1.f1)
+ self.assertEqual(d2['f2'], t1.f2)
+ self.assertEqual(d2['id'], t1.id)
+
+ # test if deserialization works for a string
+ s = d1['__kaylee_task_session__']
+ d3 = TaskSessionFields.deserialize(s)
+ self.assertEqual(len(d3), 2)
+ self.assertEqual(d3['f1'], t1.f1)
+ self.assertEqual(d3['f2'], t1.f2)
+ self.assertNotIn('id', d3)
+
+ # test if incorrect signature raises KayleeError
+ s2 = s
+ s2 = s[3:] # pad the signature
+ self.assertRaises(KayleeError, TaskSessionFields.deserialize, s2)
+
+kaylee_suite = load_tests([TaskTests])
@@ -5,6 +5,7 @@
PROJECTS_DIR = PROJECTS_DIR
+SECRET_KEY = 'aJD2fn;1340913)*(!!&$)(#&<AHFB12b'
REGISTRY = {
'name' : 'MemoryNodesRegistry',
View
0 run-tests.py 100644 → 100755
No changes.
View
2 setup.py 100644 → 100755
@@ -1,3 +1,5 @@
+#!/usr/bin/env python
+
import os
from distutils.command.install_data import install_data
from distutils.command.install import INSTALL_SCHEMES

0 comments on commit ba10e95

Please sign in to comment.