Merge c437fcd into 97c7821

tanner/config.py

@@ -33,7 +33,8 @@
                    'LOCALLOG': {'enabled': False, 'PATH': '/tmp/tanner_report.json'},
                    'CLEANLOG': {'enabled': False},
                    'REMOTE_DOCKERFILE': {'GITHUB': "https://raw.githubusercontent.com/mushorg/tanner/master/docker/"
-                                                   "tanner/template_injection/Dockerfile"}
+                                                   "tanner/template_injection/Dockerfile"},
+                   'SESSIONS': {"delete_timeout": 300}
                    }
 
 

tanner/server.py

@@ -1,12 +1,14 @@
 import asyncio
 import json
 import logging
+import concurrent
 
 import uvloop
 import yarl
 from aiohttp import web
 
-from tanner import dorks_manager, session_manager, redis_client
+from tanner import dorks_manager, redis_client
+from tanner.sessions import session_manager
 from tanner.config import TannerConfig
 from tanner.emulators import base
 from tanner.reporting.log_local import Reporting as local_report
@@ -23,6 +25,8 @@ def __init__(self):
         db_name = TannerConfig.get('SQLI', 'db_name')
 
         self.session_manager = session_manager.SessionManager()
+        self.delete_timeout = TannerConfig.get('SESSIONS', 'delete_timeout')
+
         self.dorks = dorks_manager.DorksManager()
         self.base_handler = base.BaseHandler(base_dir, db_name)
         self.logger = logging.getLogger(__name__)
@@ -56,7 +60,7 @@ async def handle_event(self, request):
             self.logger.exception('error parsing request: %s', data)
             response_msg = self._make_response(msg=type(error).__name__)
         else:
-            session = await self.session_manager.add_or_update_session(
+            session, _ = await self.session_manager.add_or_update_session(
                 data, self.redis_client
             )
             self.logger.info('Requested path %s', path)
@@ -99,6 +103,15 @@ async def handle_version(self, request):
     async def on_shutdown(self, app):
         await self.session_manager.delete_sessions_on_shutdown(self.redis_client)
         self.redis_client.close()
+        await self.redis_client.wait_closed()
+
+    async def delete_sessions(self):
+        try:
+            while True:
+                await self.session_manager.delete_old_sessions(self.redis_client)
+                await asyncio.sleep(self.delete_timeout)
+        except asyncio.CancelledError:
+            pass
 
     def setup_routes(self, app):
         app.router.add_route('*', '/', self.default_handler)
@@ -112,10 +125,21 @@ def create_app(self, loop):
         self.setup_routes(app)
         return app
 
+    async def start_background_delete(self, app):
+        app['session_delete'] = asyncio.create_task(self.delete_sessions())
+
+    async def cleanup_background_tasks(self, app):
+        app['session_delete'].cancel()
+        await app['session_delete']
+
     def start(self):
         loop = asyncio.get_event_loop()
         self.redis_client = loop.run_until_complete(redis_client.RedisClient.get_redis_client())
+
         app = self.create_app(loop)
+        app.on_startup.append(self.start_background_delete)
+        app.on_cleanup.append(self.cleanup_background_tasks)
+
         host = TannerConfig.get('TANNER', 'host')
         port = TannerConfig.get('TANNER', 'port')
         web.run_app(app, host=host, port=int(port))

tanner/session_manager.py → tanner/sessions/session_manager.py

@@ -1,35 +1,39 @@
 import logging
+import hashlib
+import asyncio
+import time
 
 import aioredis
 
-from tanner.session import Session
-from tanner.session_analyzer import SessionAnalyzer
+from tanner.sessions.session import Session
+from tanner.sessions.session_analyzer import SessionAnalyzer
 
 
 class SessionManager:
     def __init__(self, loop=None):
-        self.sessions = []
+        self.sessions = {}
         self.analyzer = SessionAnalyzer(loop=loop)
         self.logger = logging.getLogger(__name__)
 
     async def add_or_update_session(self, raw_data, redis_client):
-        # prepare the list of sessions
-        await self.delete_old_sessions(redis_client)
+
         # handle raw data
         valid_data = self.validate_data(raw_data)
         # push snare uuid into redis.
         await redis_client.sadd('snare_ids', *[valid_data['uuid']])
-        session = self.get_session(valid_data)
-        if session is None:
+        session_id = self.get_session_id(valid_data)
+        if session_id not in self.sessions:
             try:
                 new_session = Session(valid_data)
             except KeyError as key_error:
                 self.logger.exception('Error during session creation: %s', key_error)
                 return
-            self.sessions.append(new_session)
-            return new_session
-        session.update_session(valid_data)
-        return session
+            self.sessions[session_id] = new_session
+            return new_session, session_id
+        else:
+            self.sessions[session_id].update_session(valid_data)
+        # prepare the list of sessions
+        return self.sessions[session_id], session_id
 
     @staticmethod
     def validate_data(data):
@@ -53,33 +57,35 @@ def validate_data(data):
 
         return data
 
-    def get_session(self, data):
-        session = None
+    def get_session_id(self, data):
         ip = data['peer']['ip']
         user_agent = data['headers']['user-agent']
         sess_uuid = data['cookies']['sess_uuid']
-        for sess in self.sessions:
-            if sess.ip == ip and sess.user_agent == user_agent and sess_uuid == sess.get_uuid():
-                session = sess
-                break
-        return session
+
+        sess_id_string = "{ip}{user_agent}{sess_uuid}".format(ip=ip, user_agent=user_agent, sess_uuid=sess_uuid)
+
+        return hashlib.md5(sess_id_string.encode()).hexdigest()
 
     async def delete_old_sessions(self, redis_client):
-        for sess in self.sessions:
-            if not sess.is_expired():
+        id_for_deletion = []
+        for sess_id, session in self.sessions.items():
+            if not session.is_expired():
                 continue
-            is_deleted = await self.delete_session(sess, redis_client)
+            is_deleted = await self.delete_session(session, redis_client)
             if is_deleted:
-                try:
-                    self.sessions.remove(sess)
-                except ValueError:
-                    continue
+                id_for_deletion.append(sess_id)
+
+        for sess_id in id_for_deletion:
+            try:
+                del self.sessions[sess_id]
+            except ValueError:
+                continue
 
     async def delete_sessions_on_shutdown(self, redis_client):
-        for sess in self.sessions:
+        for sess_id, sess in self.sessions.items():
             is_deleted = await self.delete_session(sess, redis_client)
             if is_deleted:
-                self.sessions.remove(sess)
+                del self.sessions[sess_id]
 
     async def delete_session(self, sess, redis_client):
         await sess.remove_associated_db()

tanner/tests/test_base.py

@@ -3,7 +3,7 @@
 from unittest import mock
 
 from tanner.utils.asyncmock import AsyncMock
-from tanner import session
+from tanner.sessions import session
 from tanner.emulators import base
 from tanner import __version__ as tanner_version
 
@@ -182,7 +182,7 @@ def test_set_injectable_page(self):
         paths = [{'path': '/python.html', 'timestamp': 1465851064.2740946},
                  {'path': '/python.php/?foo=bar', 'timestamp': 1465851065.2740946},
                  {'path': '/python.html/?foo=bar', 'timestamp': 1465851065.2740946}]
-        with mock.patch('tanner.session.Session') as mock_session:
+        with mock.patch('tanner.sessions.session.Session') as mock_session:
             mock_session.return_value.paths = paths
             sess = session.Session(None)
         injectable_page = self.handler.set_injectable_page(sess)

tanner/tests/test_server.py

@@ -1,10 +1,12 @@
 import uuid
 from unittest import mock
+import hashlib
 
 from aiohttp.test_utils import AioHTTPTestCase, unittest_run_loop
 
 from tanner import server
 from tanner.config import TannerConfig
+from tanner.utils.asyncmock import AsyncMock
 from tanner import __version__ as tanner_version
 
 
@@ -23,17 +25,18 @@ def setUp(self):
 
         with mock.patch('tanner.dorks_manager.DorksManager', mock.Mock()):
             with mock.patch('tanner.emulators.base.BaseHandler', mock.Mock(), create=True):
-                with mock.patch('tanner.session_manager.SessionManager', mock.Mock(), create=True):
+                with mock.patch('tanner.sessions.session_manager.SessionManager', mock.Mock(), create=True):
                     self.serv = server.TannerServer()
 
         self.test_uuid = uuid.uuid4()
 
         async def _add_or_update_mock(data, client):
             sess = mock.Mock()
             sess.set_attack_type = mock.Mock()
+            sess_id = hashlib.md5(b"foo")
             test_uuid = uuid
             sess.get_uuid = mock.Mock(return_value=str(self.test_uuid))
-            return sess
+            return sess, sess_id
 
         async def _delete_sessions_mock(client):
             pass
@@ -48,8 +51,9 @@ async def choosed(client):
         dorks.choose_dorks = choosed
         dorks.extract_path = self._make_coroutine()
 
-        redis = mock.Mock()
+        redis = AsyncMock()
         redis.close = mock.Mock()
+        redis.wait_closed = AsyncMock()
         self.serv.dorks = dorks
         self.serv.redis_client = redis
 

tanner/tests/test_session_analyzer.py

@@ -5,7 +5,7 @@
 from unittest.mock import patch
 import geoip2
 import aioredis
-from tanner.session_analyzer import SessionAnalyzer
+from tanner.sessions.session_analyzer import SessionAnalyzer
 
 
 session = b'{"sess_uuid": "c546114f97f548f982756495f963e280", "start_time": 1466091813.4780173, ' \

tanner/tests/test_session_manager.py

@@ -1,8 +1,9 @@
 import asyncio
 import unittest
 from unittest import mock
+import hashlib
 
-from tanner import session, session_manager
+from tanner.sessions import session_manager, session
 
 
 class TestSessions(unittest.TestCase):
@@ -17,7 +18,8 @@ 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'
+                    '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,
@@ -28,7 +30,8 @@ def test_validate_missing_peer(self):
             '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'
+                    '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,
@@ -106,9 +109,9 @@ async def sess_sadd(key, value):
 
         redis_mock = mock.Mock()
         redis_mock.sadd = sess_sadd
-        sess = self.loop.run_until_complete(self.handler.add_or_update_session(data, redis_mock))
+        sess, sess_id = self.loop.run_until_complete(self.handler.add_or_update_session(data, redis_mock))
 
-        self.assertEquals([sess], self.handler.sessions)
+        self.assertDictEqual({sess_id: sess}, self.handler.sessions)
 
     def test_updating_session(self):
         async def sess_sadd(key, value):
@@ -126,12 +129,20 @@ async def sess_sadd(key, value):
             'cookies': {'sess_uuid': None}
         }
         sess = session.Session(data)
-        data['cookies']['sess_uuid'] = sess.get_uuid()
+
+        ip = data['peer']['ip']
+        user_agent = data['headers']['user-agent']
+        sess_uuid = data['cookies']['sess_uuid']
+
+        sess_id_string = "{ip}{user_agent}{sess_uuid}".format(ip=ip, user_agent=user_agent, sess_uuid=sess_uuid)
+
+        sess_id = hashlib.md5(sess_id_string.encode()).hexdigest()
+
         redis_mock = mock.Mock()
         redis_mock.sadd = sess_sadd
-        self.handler.sessions.append(sess)
+        self.handler.sessions[sess_id] = sess
         self.loop.run_until_complete(self.handler.add_or_update_session(data, redis_mock))
-        self.assertEqual(self.handler.sessions[0].count, 2)
+        self.assertEqual(self.handler.sessions[sess_id].count, 2)
 
     def test_deleting_sessions(self):
         async def analyze(session_key, redis_client):
@@ -155,11 +166,11 @@ async def sess_set(key, val):
         sess = session.Session(data)
         sess.is_expired = mock.MagicMock(name='expired')
         sess.is_expired.__bool__.reurned_value = True
-        self.handler.sessions.append(sess)
+        self.handler.sessions[sess] = sess
         redis_mock = mock.Mock()
         redis_mock.set = sess_set
         self.loop.run_until_complete(self.handler.delete_old_sessions(redis_mock))
-        self.assertListEqual(self.handler.sessions, [])
+        self.assertDictEqual(self.handler.sessions, {})
 
     def test_get_uuid(self):
         data = {