Skip to content

Commit

Permalink
Merge pull request #28 from afeena/feature/session-tests
Browse files Browse the repository at this point in the history
Handling missing data in session
  • Loading branch information
glaslos committed Jun 6, 2016
2 parents 110a1d2 + 99747b2 commit 9384e86
Show file tree
Hide file tree
Showing 3 changed files with 167 additions and 17 deletions.
25 changes: 17 additions & 8 deletions session.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand All @@ -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):
Expand All @@ -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()
35 changes: 26 additions & 9 deletions session_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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())
self.r.set(sess.get_key(), sess.to_json())
124 changes: 124 additions & 0 deletions tests/test_session_manager.py
Original file line number Diff line number Diff line change
@@ -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)

0 comments on commit 9384e86

Please sign in to comment.