Cleanup docker containers on shutdown (#171)

* fix #164 Cleanup docker containers on shutdown * fix tests
mushorg · Jul 10, 2017 · 150c05c · 150c05c
1 parent fc9dc6f
commit 150c05c
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 12 deletions.
diff --git a/tanner/server.py b/tanner/server.py
@@ -79,8 +79,9 @@ async def handle_dorks(self, request):
         return web.json_response(response_msg)
 
     async def on_shutdown(self, app):
+        await self.session_manager.delete_sessions_on_shutdown(self.redis_client)
         self.redis_client.close()
-        
+
     def setup_routes(self, app):
         app.router.add_route('*', '/', self.default_handler)
         app.router.add_post('/event', self.handle_event)

diff --git a/tanner/session_manager.py b/tanner/session_manager.py
@@ -70,13 +70,25 @@ async def delete_old_sessions(self, redis_client):
         for sess in self.sessions:
             if not sess.is_expired():
                 continue
-            await sess.remove_associated_db()
-            if sess.associated_env is not None:
-                await sess.remove_associated_env()
-            self.sessions.remove(sess)
-            try:
-                await redis_client.set(sess.get_uuid(), sess.to_json())
-                await self.analyzer.analyze(sess.get_uuid(), redis_client)
-            except asyncio_redis.NotConnectedError as redis_error:
-                self.logger.error('Error connect to redis, session stay in memory. %s', redis_error)
-                self.sessions.append(sess)
+            is_deleted = await self.delete_session(sess, redis_client)
+            if is_deleted:
+                self.sessions.remove(sess)
+
+    async def delete_sessions_on_shutdown(self, redis_client):
+        for sess in self.sessions:
+            is_deleted = await self.delete_session(sess, redis_client)
+            if is_deleted:
+                self.sessions.remove(sess)
+
+    async def delete_session(self, sess, redis_client):
+        await sess.remove_associated_db()
+        if sess.associated_env is not None:
+            await sess.remove_associated_env()
+        try:
+            await redis_client.set(sess.get_uuid(), sess.to_json())
+            await self.analyzer.analyze(sess.get_uuid(), redis_client)
+        except asyncio_redis.NotConnectedError as redis_error:
+            self.logger.error('Error connect to redis, session stay in memory. %s', redis_error)
+            return False
+        else:
+            return True
diff --git a/tanner/tests/test_server.py b/tanner/tests/test_server.py
@@ -32,8 +32,12 @@ async def _add_or_update_mock(data, client):
             sess.get_uuid = mock.Mock(return_value=str(self.test_uuid))
             return sess
 
-        self.serv.session_manager.add_or_update_session = _add_or_update_mock
+        async def _delete_sessions_mock(client):
+            pass
 
+        self.serv.session_manager.add_or_update_session = _add_or_update_mock
+        self.serv.session_manager.delete_sessions_on_shutdown = _delete_sessions_mock
+
         async def choosed(client):
             return [x for x in range(10)]