push analyzed sessions into redis (#53)

* push analyzed sessions into redis
mushorg · Jul 11, 2016 · 03315b4 · 03315b4
1 parent a572b15
commit 03315b4
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 11 deletions.
diff --git a/session_analyzer.py b/session_analyzer.py
@@ -9,6 +9,7 @@
 class SessionAnalyzer:
     def __init__(self):
         self.r = redis.StrictRedis(host='localhost', port=6379)
+        self.queue = asyncio.Queue()
 
     @asyncio.coroutine
     def analyze(self, session_key):
@@ -21,7 +22,21 @@ def analyze(self, session_key):
             print("Can't get session for analyze", e)
         else:
             result = self.create_stats(session)
-            return result
+            yield from self.queue.put(result)
+            yield from self.save_session()
+
+    @asyncio.coroutine
+    def save_session(self):
+        while not self.queue.empty():
+            session = yield from self.queue.get()
+            s_key = session['sensor_uuid']
+            del_key = session['uuid']
+            try:
+                self.r.lpush(s_key, json.dumps(session))
+                self.r.delete(del_key)
+            except redis.ConnectionError as e:
+                print('Error with redis. Session will be returned to the queue', e)
+                self.queue.put(session)
 
     def create_stats(self, session):
         sess_duration = session['end_time'] - session['start_time']
@@ -51,7 +66,7 @@ def create_stats(self, session):
 
     def analyze_paths(self, paths):
         tbr = []
-        attack_types = set()
+        attack_types = []
         current_path = paths[0]
         dorks = self.r.smembers(DorksManager.dorks_key)
 
@@ -68,7 +83,7 @@ def analyze_paths(self, paths):
             if path['path'] in dorks:
                 hidden_links += 1
             if 'attack_type' in path:
-                attack_types.add(path['attack_type'])
+                attack_types.append(path['attack_type'])
         return tbr_average, errors, hidden_links, attack_types
 
     def choose_possible_owner(self, stats):
@@ -99,7 +114,7 @@ def choose_possible_owner(self, stats):
         if stats['hidden_links'] > 0:
             possible_owners['crawler'] += 1
             possible_owners['attacker'] += 1
-        if stats['attack_types'].intersection(attacks):
+        if set(stats['attack_types']).intersection(attacks):
             possible_owners['attacker'] += 1
 
         maxval = max(possible_owners.items(), key=operator.itemgetter(1))[1]

diff --git a/tests/test_session_analyzer.py b/tests/test_session_analyzer.py
@@ -26,7 +26,7 @@ def setUp(self):
         self.session = json.loads(session.decode('utf-8'))
         with mock.patch('redis.StrictRedis', mock.Mock()):
             self.handler = SessionAnalyzer()
-        attrs = {'get.return_value': session, 'smembers.return_value': set()}
+        attrs = {'get.return_value': session, 'smembers.return_value': set(),'lpush.return_value':''}
         self.handler.r = mock.Mock(**attrs)
 
     def tests_load_session_fail(self):
@@ -38,12 +38,6 @@ def tests_load_session_fail(self):
             loop.run_until_complete(self.handler.analyze(None))
         self.assertRaises(redis.ConnectionError)
 
-    def test_create_analyze_fail(self):
-        res = None
-        loop = asyncio.get_event_loop()
-        res = loop.run_until_complete(self.handler.analyze(None))
-        self.assertIsNotNone(res)
-
     def test_create_stats(self):
         stats = self.handler.create_stats(self.session)
         self.assertEqual(stats['possible_owners'], ['attacker'])