From 99747b2a61a4dfae7707260c2aa3b049045644de Mon Sep 17 00:00:00 2001 From: Evgeniya Tokarchuk Date: Sun, 5 Jun 2016 01:45:23 +0500 Subject: [PATCH] validate session data, fix bugs, test for sessions --- session.py | 25 ++++--- session_manager.py | 35 +++++++--- tests/test_session_manager.py | 124 ++++++++++++++++++++++++++++++++++ 3 files changed, 167 insertions(+), 17 deletions(-) create mode 100644 tests/test_session_manager.py diff --git a/session.py b/session.py index f6595b32..c3266db5 100644 --- a/session.py +++ b/session.py @@ -7,13 +7,17 @@ class Session: KEEP_ALIVE_TIME = 75 def __init__(self, data): - self.ip = data['peer']['ip'] - self.port = data['peer']['port'] - self.user_agent = data['headers']['USER-AGENT'] - self.uuid = data['uuid'] + try: + self.ip = data['peer']['ip'] + self.port = data['peer']['port'] + self.user_agent = data['headers']['user-agent'] + self.uuid = data['uuid'] + self.paths = [{'path': data['path'], 'timestamp': time.time()}] + except KeyError as e: + raise + self.timestamp = time.time() self.count = 1 - self.paths = [{'path': data['path'], 'timestamp': time.time()}] def update_session(self, path): self.timestamp = time.time() @@ -22,7 +26,7 @@ def update_session(self, path): def is_expired(self): exp_time = self.timestamp + self.KEEP_ALIVE_TIME - if (time.time() - exp_time > 0): + if time.time() - exp_time > 0: return True def to_json(self): @@ -36,5 +40,10 @@ def to_json(self): return json.dumps(s) def get_key(self): - bstr = (self.ip + self.user_agent).encode('utf-8') - return hashlib.md5(bstr).digest() + bstr = b'' + try: + bstr = (str(self.ip) + str(self.user_agent)).encode('utf-8') + except ValueError as e: + print('can\'t create byte string for hash', e) + finally: + return hashlib.md5(bstr).digest() diff --git a/session_manager.py b/session_manager.py index a727c519..62579a79 100644 --- a/session_manager.py +++ b/session_manager.py @@ -10,25 +10,42 @@ def __init__(self): self.r = redis.StrictRedis(host='localhost', port=6379) @asyncio.coroutine - def add_or_update_session(self, data): + def add_or_update_session(self, raw_data): # prepare the list of sessions self.delete_old_sessions() - if 'peer' not in data: - peer = dict(ip=None, port=None) - data['peer'] = peer - session = self.get_session(data) + # handle raw data + valid_data = self.validate_data(raw_data) + session = self.get_session(valid_data) if session is None: - new_session = Session(data) + try: + new_session = Session(valid_data) + except KeyError: + print('Bad session') + return self.sessions.append(new_session) return new_session else: - session.update_session(data['path']) + session.update_session(valid_data['path']) return session + def validate_data(self, data): + if 'peer' not in data: + peer = dict(ip=None, port=None) + data['peer'] = peer + + data['headers'] = dict((k.lower(), v) for k, v in data['headers'].items()) + if 'user-agent' not in data['headers']: + data['headers']['user-agent'] = None + if 'path' not in data: + data['path'] = None + if 'uuid' not in data: + data['uuid'] = None + return data + def get_session(self, data): session = None ip = data['peer']['ip'] - user_agent = data['headers']['USER-AGENT'] + user_agent = data['headers']['user-agent'] for sess in self.sessions: if sess.ip == ip and sess.user_agent == user_agent: session = sess @@ -40,4 +57,4 @@ def delete_old_sessions(self): if not sess.is_expired(): continue self.sessions.remove(sess) - self.r.set(sess.get_key(), sess.to_json()) \ No newline at end of file + self.r.set(sess.get_key(), sess.to_json()) diff --git a/tests/test_session_manager.py b/tests/test_session_manager.py new file mode 100644 index 00000000..d9ecd783 --- /dev/null +++ b/tests/test_session_manager.py @@ -0,0 +1,124 @@ +import unittest +import session_manager +import session +import hashlib +from unittest import mock + + +class TestSessions(unittest.TestCase): + def setUp(self): + with mock.patch('redis.StrictRedis', mock.Mock(), create=True): + self.handler = session_manager.SessionManager() + + def test_validate_missing_peer(self): + data = { + 'headers': { + 'USER-AGENT': 'Mozilla/5.0 (Windows NT 6.1) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/41.0.2228.0 Safari/537.36' + }, + 'path': '/foo', + 'uuid': None + } + + assertion_data = { + 'peer': {'ip': None, 'port': None}, + 'headers': { + 'user-agent': 'Mozilla/5.0 (Windows NT 6.1) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/41.0.2228.0 Safari/537.36' + }, + 'path': '/foo', + 'uuid': None + } + data = self.handler.validate_data(data) + self.assertDictEqual(data, assertion_data) + + def test_validate_missing_user_agent(self): + data = { + 'peer': { + 'ip': '127.0.0.1', + 'port': 80 + }, + 'headers': {}, + 'path': '/foo', + 'uuid': None + } + + assertion_data = { + 'peer': { + 'ip': '127.0.0.1', + 'port': 80 + }, + 'headers': {'user-agent': None}, + 'path': '/foo', + 'uuid': None + } + data = self.handler.validate_data(data) + self.assertDictEqual(data, assertion_data) + + def test_adding_new_session(self): + data = { + 'peer': { + }, + 'headers': {}, + 'path': '/foo', + 'uuid': None + } + sess = yield from self.handler.add_or_update_session(data) + assertion_data = { + 'peer': { + 'ip': None, + 'port': None + }, + 'headers': {'user-agent':None}, + 'path': '/foo', + 'uuid': None + } + assertion_session = session.Session(assertion_data) + self.assertEquals(session,assertion_session) + + def test_updating_session(self): + data = { + 'peer': { + 'ip': None, + 'port': None + }, + 'headers': {'user-agent': None}, + 'path': '/foo', + 'uuid': None + } + sess = session.Session(data) + self.handler.sessions.append(sess) + yield from self.handler.add_or_update_session(data) + self.assertEqual(self.handler.sessions[0].count,2) + def test_deleting_sessions(self): + data = { + 'peer': { + 'ip': None, + 'port': None + }, + 'headers': {'user-agent': None}, + 'path': '/foo', + 'uuid': None + } + sess = session.Session(data) + sess.is_expired = mock.MagicMock(name='expired') + sess.is_expired.__bool__.reurned_value = True + self.handler.sessions.append(sess) + experied = mock.Mock() + experied.return_value = True + self.handler.delete_old_sessions() + self.assertListEqual(self.handler.sessions,[]) + + def test_get_key(self): + data = { + 'peer': { + 'ip': None, + 'port': None + }, + 'headers': {'user-agent': None}, + 'path': '/foo', + 'uuid': None + } + sess = session.Session(data) + key = sess.get_key() + bstr = ('None'+'None').encode('utf-8') + hash = hashlib.md5(bstr).digest() + self.assertEqual(key,hash) \ No newline at end of file