optional postgress database #322

bin/tanner

@@ -17,7 +17,9 @@ def main():
     """)
     parser = argparse.ArgumentParser()
     parser.add_argument("--config", help="tanner config")
+    parser.add_argument("--database", help="choose a database postgres/redis")
     args = parser.parse_args()
+    database=args.database
     if args.config:
         TannerConfig.set_config(args.config)
     debug_log_file_name = TannerConfig.get("LOGGER", "log_debug")
@@ -27,7 +29,7 @@ def main():
     print("Error logs will be stored in", error_log_file_name)
     if TannerConfig.get('LOCALLOG', 'enabled') is True:
         print("Data logs will be stored in", TannerConfig.get('LOCALLOG', 'PATH'))
-    tanner = server.TannerServer()
+    tanner = server.TannerServer(database=database)
     tanner.start()
 
 

requirements.txt

@@ -1,5 +1,6 @@
 aiohttp
 aiomysql
+aiopg
 aiohttp_jinja2==1.1.0
 docker<2.6
 mimesis<3.0.0

tanner/config.py

@@ -18,6 +18,7 @@
                    'API': {'host': '0.0.0.0', 'port': 8092, 'auth': False, 'auth_signature': 'tanner_api_auth'},
                    'PHPOX': {'host': '0.0.0.0', 'port': 8088},
                    'REDIS': {'host': 'localhost', 'port': 6379, 'poolsize': 80, 'timeout': 1},
+		           'POSTGRES': {'host': 'localhost', 'port': 6379, 'poolsize': 80, 'timeout': 1, 'db_name':'tanner_db', 'user':'postgres', 'password':''},
                    'EMULATORS': {'root_dir': '/opt/tanner'},
                    'EMULATOR_ENABLED': {'sqli': True, 'rfi': True, 'lfi': True, 'xss': True, 'cmd_exec': True,
                                         'php_code_injection': True, 'php_object_injection': True, "crlf": True,

tanner/postgres_client.py

@@ -0,0 +1,26 @@
+import asyncio
+import logging
+import sqlalchemy import create_engine
+import aiopg
+from tanner.config import TannerConfig
+
+LOGGER = logging.getLogger(__name__)
+
+class PostgresClient:
+    async def get_postgres_client():
+        postgres_client=None
+        try:
+            host=TannerConfig.get('POSTGRES', 'host')
+            port=TannerConfig.get('POSTGRES', 'port')
+            dbname=TannerConfig.get('POSTGRES', 'db_name')
+            user=TannerConfig.get('POSTGRES', 'user')
+            timeout=TannerConfig.get('POSTGRES','timeout')
+            password=TannerConfig.get('POSTGRES', 'password')
+            if poolsize is None:
+                poolsize=TannerConfig.get('REDIS', 'poolsize')
+            db_string = 'dbname={} user={} password={} host={} port={}'.format(dbname, user, password, host, port)
+            postgres_client = await asyncio.wait_for(aiopg.create_pool(db_string,maxsize=poolsize),timeout=int(timeout))
+        except asyncio.TimeoutError as timeout_error:
+            LOGGER.exception("Failed to connect to postgres {}".format(timeout_error))
+            exit()
+        return postgres_client

tanner/server.py

@@ -7,7 +7,7 @@
 import yarl
 from aiohttp import web
 
-from tanner import dorks_manager, redis_client
+from tanner import dorks_manager, redis_client, postgres_client
 from tanner.sessions import session_manager
 from tanner.config import TannerConfig
 from tanner.emulators import base
@@ -20,7 +20,7 @@
 
 
 class TannerServer:
-    def __init__(self):
+    def __init__(self,database):
         base_dir = TannerConfig.get('EMULATORS', 'root_dir')
         db_name = TannerConfig.get('SQLI', 'db_name')
 
@@ -30,7 +30,8 @@ def __init__(self):
         self.dorks = dorks_manager.DorksManager()
         self.base_handler = base.BaseHandler(base_dir, db_name)
         self.logger = logging.getLogger(__name__)
-        self.redis_client = None
+        self.db_client = None
+        self.database=database
 
         if TannerConfig.get('HPFEEDS', 'enabled') is True:
             self.hpf = hpfeeds_report()
@@ -61,10 +62,10 @@ async def handle_event(self, request):
             response_msg = self._make_response(msg=type(error).__name__)
         else:
             session, _ = await self.session_manager.add_or_update_session(
-                data, self.redis_client
+                data, self.db_client
             )
             self.logger.info('Requested path %s', path)
-            await self.dorks.extract_path(path, self.redis_client)
+            await self.dorks.extract_path(path, self.db_client)
             detection = await self.base_handler.handle(data, session)
             session.set_attack_type(path, detection["name"])
 
@@ -92,7 +93,7 @@ async def handle_event(self, request):
         return web.json_response(response_msg)
 
     async def handle_dorks(self, request):
-        dorks = await self.dorks.choose_dorks(self.redis_client)
+        dorks = await self.dorks.choose_dorks(self.db_client)
         response_msg = dict(version=tanner_version, response=dict(dorks=dorks))
         return web.json_response(response_msg)
 
@@ -101,14 +102,14 @@ async def handle_version(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()
-        await self.redis_client.wait_closed()
+        await self.session_manager.delete_sessions_on_shutdown(self.db_client)
+        self.db_client.close()
+        await self.db_client.wait_closed()
 
     async def delete_sessions(self):
         try:
             while True:
-                await self.session_manager.delete_old_sessions(self.redis_client)
+                await self.session_manager.delete_old_sessions(self.db_client)
                 await asyncio.sleep(self.delete_timeout)
         except asyncio.CancelledError:
             pass
@@ -134,8 +135,10 @@ async def cleanup_background_tasks(self, app):
 
     def start(self):
         loop = asyncio.get_event_loop()
-        self.redis_client = loop.run_until_complete(redis_client.RedisClient.get_redis_client())
-
+        if self.database == 'redis':
+            self.db_client = loop.run_until_complete(redis_client.RedisClient.get_redis_client())
+        elif self.database == 'postgres':
+            self.db_client = loop.run_until_complete(postgres_client.PostgresClient.get_postgres_client())
         app = self.create_app(loop)
         app.on_startup.append(self.start_background_delete)
         app.on_cleanup.append(self.cleanup_background_tasks)