add service models manager to keep track of db tables of service in use

mod_config/controllers.py

@@ -1,13 +1,16 @@
 import os
+import sys
 
 import subprocess
 import threading
 import shutil
 import zipfile
+import importlib
 
 from flask import Blueprint, g, request, send_file, jsonify, abort, \
     url_for, redirect
 from werkzeug.utils import secure_filename
+from sqlalchemy import create_engine
 
 from database import create_session
 from decorators import get_menu_entries, template_renderer
@@ -18,7 +21,7 @@
     RuleForm, DeleteRuleForm
 from mod_config.models import Service, Notification, Rule, Actions, Conditions
 from pipot.notifications import NotificationLoader
-from pipot.services import ServiceLoader
+from pipot.services import ServiceLoader, ServiceModelsManager
 
 mod_config = Blueprint('config', __name__)
 
@@ -280,20 +283,20 @@ def data_processing_ajax(action):
     return jsonify(result)
 
 
-def verify_and_import_module(temp_path, final_path, form, is_container=False):
+def verify_and_import_module(final_path, form, is_container=False):
     if is_container:
-        instance = ServiceLoader.load_from_container(temp_path)
+        instance = ServiceLoader.load_from_container(final_path, temp_folder=False)
     else:
-        instance = ServiceLoader.load_from_file(temp_path)
+        instance = ServiceLoader.load_from_file(final_path, temp_folder=False)
     # Auto-generate tables
     instance.get_used_table_names()
     # Move
-    os.rename(temp_path, final_path)
     service = Service(instance.__class__.__name__,
                       form.description.data)
     g.db.add(service)
     g.db.commit()
 
+
 @mod_config.route('/services', methods=['GET', 'POST'])
 @login_required
 @check_access_rights()
@@ -306,7 +309,6 @@ def services():
         if file:
             filename = secure_filename(file.filename)
             basename, extname = os.path.splitext(filename)
-            temp_dir = os.path.join('./pipot/services/temp', basename)
             final_dir = os.path.join('./pipot/services', basename)
             if not os.path.isdir(final_dir):
                 if extname == '.zip':
@@ -315,32 +317,32 @@ def services():
                     if ret:
                         form.errors['container'] = ['Corrupt container']
                     else:
-                        zip_file.extractall('./pipot/services/temp')
+                        zip_file.extractall('./pipot/services')
                         try:
-                            verify_and_import_module(temp_dir, final_dir, form, is_container=True)
+                            verify_and_import_module(final_dir, form, is_container=True)
                             # Reset form, all ok
                             form = NewServiceForm(None)
                         except ServiceLoader.ServiceLoaderException as e:
-                            shutil.rmtree(temp_dir)
+                            shutil.rmtree(final_dir)
                             form.errors['container'] = [e.value]
                 else:
-                    if os.path.exists(temp_dir):
-                        shutil.rmtree(temp_dir)
-                    os.mkdir(temp_dir)
-                    temp_file = os.path.join(temp_dir, filename)
+                    os.mkdir(final_dir)
+                    final_file = os.path.join(final_dir, filename)
                     # create the __init__.py for module import
-                    file.save(temp_file)
-                    open(os.path.join(temp_dir, '__init__.py'), 'w')
+                    file.save(final_file)
+                    open(os.path.join(final_dir, '__init__.py'), 'w')
                     # Import and verify module
                     try:
-                        verify_and_import_module(temp_dir, final_dir, form, is_container=False)
+                        verify_and_import_module(final_dir, form, is_container=False)
                         # Reset form, all ok
                         form = NewServiceForm(None)
                     except ServiceLoader.ServiceLoaderException as e:
                         # Remove file
-                        shutil.rmtree(temp_dir)
+                        shutil.rmtree(final_dir)
                         # Pass error to user
                         form.errors['file'] = [e.value]
+                # add service name to services.txt
+                ServiceModelsManager.add_models(basename)
             else:
                 form.errors['file'] = ['Service already exists.']
     return {
@@ -363,8 +365,16 @@ def services_ajax(action):
         if form.validate_on_submit():
             service = Service.query.filter(
                 Service.id == form.id.data).first()
-            # Delete service
+            # Delete service in db
             g.db.delete(service)
+            # Delete service model
+            removed_models = ServiceModelsManager.rm_models(service.name)
+            for model_name in removed_models:
+                module = importlib.import_module('pipot.services' + '.' + service.name + '.' + service.name)
+                model = getattr(module, model_name.lstrip('.' + service.name))
+                from database import Base, db_engine
+                Base.metadata.drop_all(bind=db_engine, tables=[model.__table__])
+                Base.metadata.remove(model.__table__)
             # Delete file
             try:
                 shutil.rmtree(service.get_file())

pipot/services/ServiceLoader.py

@@ -1,5 +1,6 @@
 import importlib
 import os
+import sys
 
 from pipot.services.IService import IService
 import pipot.services as main
@@ -17,7 +18,7 @@ def __str__(self):
         return repr(self.value)
 
 
-def load_from_container(container_dir):
+def load_from_container(container_dir, temp_folder=False):
     """Attempts to load the service from a folder with the same name
     Required container format:
     myService.zip
@@ -40,7 +41,7 @@ def load_from_container(container_dir):
             pass
         if not os.path.isfile(os.path.join(container_dir, '__init__.py')):
             open(os.path.join(container_dir, '__init__.py'), 'w')
-        instance = load_from_file(mod_file)
+        instance = load_from_file(mod_file, temp_folder=temp_folder)
         return instance
 
 

pipot/services/ServiceModelsManager.py

@@ -0,0 +1,38 @@
+from __future__ import print_function
+import os
+import sys
+import importlib
+import inspect
+from database import Base
+
+models_storage = './pipot/services/models.txt'
+
+
+def add_models(service):
+    models = get_models()
+    cls_members = inspect.getmembers(importlib.import_module('pipot.services' + '.' + service + '.' + service),
+                                     inspect.isclass)
+    cls_info = list(filter(lambda x: Base in inspect.getmro(x[1]) and x[0] not in ('IModel', 'IModelIP'), cls_members))
+    models.extend([service + '.' + name for name, _ in cls_info])
+    save_models(models)
+
+
+def rm_models(service):
+    models = get_models()
+    removed_models = list(filter(lambda x: x.startswith(service), models))
+    models = list(filter(lambda x: not x.startswith(service), models))
+    with open(models_storage, 'w') as f:
+        for model in models:
+            print(model, file=f)
+    return removed_models
+
+
+def get_models():
+    with open(models_storage, 'r') as f:
+        return [line.strip('\n') for line in f.readlines()]
+
+
+def save_models(models):
+    with open(models_storage, 'w') as f:
+        for model in models:
+            print(model, file=f)

tests/testAppBase.py

@@ -16,7 +16,7 @@
 import flask
 from flask import g, current_app, session
 from collections import namedtuple
-from database import create_session, Base
+from database import create_session
 from mod_auth.models import User, Role, Page, PageAccess
 from mod_config.models import Service, Notification, Actions, Conditions, Rule
 from mod_honeypot.models import Profile, PiModels, PiPotReport, ProfileService, \
@@ -88,6 +88,7 @@ def setUp(self):
         self.client = self.app.test_client(self)
 
     def tearDown(self):
+        from database import Base
         db = create_session(self.app.config['DATABASE_URI'], drop_tables=False)
         db_engine = create_engine(self.app.config['DATABASE_URI'], convert_unicode=True)
         Base.metadata.drop_all(bind=db_engine)

tests/testServiceManagement.py

@@ -14,18 +14,26 @@
 from database import create_session
 from mod_config.models import Service
 from tests.testAppBase import TestAppBase
+from pipot.services import ServiceModelsManager
 
 test_dir = os.path.dirname(os.path.abspath(__file__))
 service_dir = os.path.join(test_dir, '../pipot/services/')
-
+temp_dir = os.path.join(test_dir, 'temp')
+if not os.path.exists(temp_dir):
+    os.makedirs(temp_dir)
+ServiceModelsManager.models_storage = os.path.join(temp_dir, 'models.txt')
 
 class TestServiceManagement(TestAppBase):
 
     def setUp(self):
         super(TestServiceManagement, self).setUp()
+        if not os.path.isfile(ServiceModelsManager.models_storage):
+            with open(ServiceModelsManager.models_storage, 'w'):
+                pass
 
     def tearDown(self):
         super(TestServiceManagement, self).tearDown()
+        os.remove(ServiceModelsManager.models_storage)
 
     def add_and_remove_service(self, service_name, service_file_name):
         # upload the service file
@@ -45,6 +53,8 @@ def add_and_remove_service(self, service_name, service_file_name):
         self.assertTrue(os.path.isfile(os.path.join(service_dir, service_name, service_name + '.py')))
         # check service file and folder is removed under temp_path
         self.assertFalse(os.path.isdir(os.path.join(service_dir, 'temp', service_name)))
+        # check models.txt is updated
+        self.assertEqual(['TelnetService.ReportTelnet'], ServiceModelsManager.get_models())
         # check database
         try:
             db = create_session(self.app.config['DATABASE_URI'], drop_tables=False)
@@ -53,6 +63,8 @@ def add_and_remove_service(self, service_name, service_file_name):
             name = service_row.name
         finally:
             db.remove()
+        from database import db_engine
+        self.assertTrue(db_engine.has_table('report_telnet'))
         self.assertEqual(service_name, name)
         # delete service file
         with self.app.test_client() as client:
@@ -64,6 +76,8 @@ def add_and_remove_service(self, service_name, service_file_name):
             self.assertEqual(response.get_json()['status'], 'success')
         # check service file and folder is removed unser final_path
         self.assertFalse(os.path.isfile(os.path.join(service_dir, service_name, service_file_name)))
+        # check models.txt is updated
+        self.assertEqual([], ServiceModelsManager.get_models())
         # check database
         try:
             db = create_session(self.app.config['DATABASE_URI'], drop_tables=False)
@@ -72,6 +86,10 @@ def add_and_remove_service(self, service_name, service_file_name):
         finally:
             db.remove()
         self.assertTrue(result)
+        from database import db_engine
+        self.assertFalse(db_engine.has_table('report_telnet'))
+        # clean the service module, otherwise table won't be added to Base.metadata in next test
+        del sys.modules['pipot.services.' + service_name + '.' + service_name]
 
     def test_add_and_delete_service_file(self):
         service_name = 'TelnetService'