Add container-based service upload (#25)

install/install.sh

@@ -114,6 +114,9 @@ echo "We need some information for the admin account"
 read -e -p "Admin username: " -i "admin" admin_name
 read -e -p "Admin email: " admin_email
 read -e -p "Admin password: " admin_password
+while [ ${#admin_password} == 0 ]; do
+    read -e -p "Invalid Admin password(cannot be empty), pleasee retry: " admin_password
+done
 echo "Creating admin account: "
 python "${dir}/init_db.py" "${config_db_uri}" "${admin_name}" "${admin_email}" "${admin_password}"
 echo ""

mod_config/controllers.py

@@ -3,6 +3,7 @@
 import subprocess
 import threading
 import shutil
+import zipfile
 
 from flask import Blueprint, g, request, send_file, jsonify, abort, \
     url_for, redirect
@@ -278,40 +279,68 @@ def data_processing_ajax(action):
     return jsonify(result)
 
 
+def verify_and_import_module(temp_path, final_path, form, is_container=False):
+    # import pdb; pdb.set_trace()
+    if is_container:
+        instance = ServiceLoader.load_from_container(temp_path)
+    else:
+        instance = ServiceLoader.load_from_file(temp_path)
+    # 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()
 @template_renderer()
 def services():
     form = NewServiceForm()
-
     if form.validate_on_submit():
         # Process uploaded file
         file = request.files[form.file.name]
         if file:
             filename = secure_filename(file.filename)
-            temp_path = os.path.join('./pipot/services/temp', filename)
-            final_path = os.path.join('./pipot/services', filename)
-            if not os.path.isfile(final_path):
-                file.save(temp_path)
-                # Import and verify module
-                try:
-                    instance = ServiceLoader.load_from_file(temp_path)
-                    # 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()
-                    # Reset form, all ok
-                    form = NewServiceForm(None)
-                except ServiceLoader.ServiceLoaderException as e:
-                    # Remove file
-                    os.remove(temp_path)
-                    # Pass error to user
-                    form.errors['file'] = [e.value]
+            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':
+                    zip_file = zipfile.ZipFile(file)
+                    ret = zip_file.testzip()
+                    if ret:
+                        form.errors['container'] = ['Corrupt container']
+                    else:
+                        zip_file.extractall('./pipot/services/temp')
+                        try:
+                            verify_and_import_module(temp_dir, final_dir, form, is_container=True)
+                            # Reset form, all ok
+                            form = NewServiceForm(None)
+                        except ServiceLoader.ServiceLoaderException as e:
+                            shutil.rmtree(temp_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)
+                    # create the __init__.py for module import
+                    file.save(temp_file)
+                    open(os.path.join(temp_dir, '__init__.py'), 'w')
+                    # Import and verify module
+                    try:
+                        verify_and_import_module(temp_dir, final_dir, form, is_container=False)
+                        # Reset form, all ok
+                        form = NewServiceForm(None)
+                    except ServiceLoader.ServiceLoaderException as e:
+                        # Remove file
+                        shutil.rmtree(temp_dir)
+                        # Pass error to user
+                        form.errors['file'] = [e.value]
             else:
                 form.errors['file'] = ['Service already exists.']
     return {
@@ -338,7 +367,7 @@ def services_ajax(action):
             g.db.delete(service)
             # Delete file
             try:
-                os.remove(service.get_file())
+                shutil.rmtree(service.get_file())
                 # Finalize service delete
                 g.db.commit()
                 result['status'] = 'success'

mod_config/forms.py

@@ -1,4 +1,6 @@
 import re
+import os
+from enum import Enum
 from flask_wtf import Form
 from wtforms import SubmitField, FileField, TextAreaField, HiddenField, \
     SelectField, StringField, IntegerField
@@ -7,44 +9,54 @@
 from mod_config.models import Service, Notification
 
 
-def is_python(file_name):
+class FileType(Enum):
+    PYTHONFILE = 1
+    CONTAINER = 2
+
+
+def is_python_or_container(file_name):
     # Check if it ends on .py
-    is_py = re.compile("^[^/\\\]*\.py$")
-    if not is_py.match(file_name):
-        raise ValidationError('Provided file is not a python (.py) file!')
+    is_py = re.compile(r"^[^/\\]*.py$").match(file_name)
+    is_container = re.compile((r"^[^/\\]*.zip$")).match(file_name)
+    if not is_py and not is_container:
+        raise ValidationError('Provided file is not a python (.py) file or a container (.zip)!')
+    return FileType.CONTAINER if is_container else FileType.PYTHONFILE
 
 
 def simple_service_file_validation(check_service=True):
     def validate_file(form, field):
-        is_python(field.data.filename)
-        # Name cannot be one of the files we already have
-        if field.data.filename in ['__init__py', 'IService.py',
-                                   'ServiceLoader.py']:
-            raise ValidationError('Illegal file name!')
-        if check_service:
-            # Name cannot be registered already
-            service = Service.query.filter(Service.name ==
-                                           field.data.filename).first()
-            if service is not None:
-                raise ValidationError('There is already an interface with '
-                                      'this name!')
+        field.data.filename = os.path.basename(field.data.filename)
+        file_type = is_python_or_container(field.data.filename)
+        if file_type is FileType.PYTHONFILE:
+            # Name cannot be one of the files we already have
+            if field.data.filename in ['__init__py', 'IService.py',
+                                       'ServiceLoader.py']:
+                raise ValidationError('Illegal file name!')
+            if check_service:
+                # Name cannot be registered already
+                service = Service.query.filter(Service.name ==
+                                               field.data.filename).first()
+                if service is not None:
+                    raise ValidationError('There is already an interface with '
+                                          'this name!')
     return validate_file
 
 
 def simple_notification_file_validation(check_notification=True):
     def validate_file(form, field):
-        is_python(field.data.filename)
-        # Name cannot be one of the files we already have
-        if field.data.filename in ['__init__py', 'INotification.py',
-                                   'NotificationLoader.py']:
-            raise ValidationError('Illegal file name!')
-        if check_notification:
-            # Name cannot be registered already
-            notification = Notification.query.filter(
-                Notification.name == field.data.filename).first()
-            if notification is not None:
-                raise ValidationError('There is already an interface with '
-                                      'this name!')
+        file_type = is_python_or_container(field.data.filename)
+        if file_type == FileType.PYTHONFILE:
+            # Name cannot be one of the files we already have
+            if field.data.filename in ['__init__py', 'INotification.py',
+                                       'NotificationLoader.py']:
+                raise ValidationError('Illegal file name!')
+            if check_notification:
+                # Name cannot be registered already
+                notification = Notification.query.filter(
+                    Notification.name == field.data.filename).first()
+                if notification is not None:
+                    raise ValidationError('There is already an interface with '
+                                          'this name!')
     return validate_file
 
 

mod_config/models.py

@@ -26,7 +26,7 @@ def get_file(self, temp_folder=False):
         return os.path.join(
             './pipot/services',
             'temp' if temp_folder else '',
-            self.name + '.py'
+            self.name
         )
 
 

pipot/services/ServiceLoader.py

@@ -17,6 +17,33 @@ def __str__(self):
         return repr(self.value)
 
 
+def load_from_container(container_dir):
+    """Attempts to load the service from a folder with the same name
+    Required container format:
+    myService.zip
+    |-myService.py
+    |-__init__.py (will be created if doesn't exist)
+    |-requirement.txt (optional)
+    |-other file/folder
+
+    :param container_dir: The path of container
+    :type container_dir: str
+    :return: A class instance of the loaded class.
+    :rtype: pipot.services.IService.IService
+    """
+    mod_name = os.path.split(container_dir)[-1]
+    mod_file = os.path.join(container_dir, mod_name + '.py')
+    if not os.path.isfile(mod_file):
+        raise ServiceLoaderException('There is no service file %s.py found inside container' % mod_name)
+    else:
+        if os.path.isfile(os.path.join(container_dir, 'requirement.txt')):
+            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)
+        return instance
+
+
 def load_from_file(file_name, temp_folder=True):
     """
     Attempts to load a given class from a file with the same name in this
@@ -33,7 +60,7 @@ def load_from_file(file_name, temp_folder=True):
 
     try:
         py_mod = importlib.import_module(
-            '.' + mod_name,
+            '.' + mod_name + '.' + mod_name,
             temp.__name__ if temp_folder else main.__name__)
 
         if hasattr(py_mod, mod_name):

tests/testAppBase.py

@@ -13,18 +13,19 @@
 sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 
 import tests.config
+import flask
+from flask import g, current_app, session
 from collections import namedtuple
-from flask import g, current_app
 from database import create_session, Base
 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, \
     CollectorTypes, Deployment
 
 
-def generate_keys(tempdir):
-    secret_csrf_path = os.path.join(tempdir, "secret_csrf")
-    secret_key_path = os.path.join(tempdir, "secret_key")
+def generate_keys(keydir):
+    secret_csrf_path = os.path.join(keydir, "secret_csrf")
+    secret_key_path = os.path.join(keydir, "secret_key")
     if not os.path.exists(secret_csrf_path):
         secret_csrf_cmd = "head -c 24 /dev/urandom > {path}".format(path=secret_csrf_path)
         os.system(secret_csrf_cmd)
@@ -35,8 +36,8 @@ def generate_keys(tempdir):
     return {'secret_csrf_path': secret_csrf_path, 'secret_key_path': secret_key_path}
 
 
-def load_config(tempdir):
-    key_paths = generate_keys(tempdir)
+def load_config(keydir):
+    key_paths = generate_keys(keydir)
     with open(key_paths['secret_key_path'], 'rb') as secret_key_file:
         secret_key = secret_key_file.read()
     with open(key_paths['secret_csrf_path'], 'rb') as secret_csrf_file:
@@ -59,34 +60,47 @@ def load_config(tempdir):
             }
 
 
-class TestAppBaseTest(unittest.TestCase):
-    tempdir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "temp")
-    
+class TestAppBase(unittest.TestCase):
+    keydir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "keys")
+
     def create_app(self):
-        with patch('config_parser.parse_config', return_value=load_config(self.tempdir)):
+        with patch('config_parser.parse_config', return_value=load_config(self.keydir)):
             from run import app
             return app
 
     def create_admin(self):
         # test if there is admin existed
+        name, password, email = "admin", "adminpwd", "admin@email.com"
         db = create_session(self.app.config['DATABASE_URI'], drop_tables=False)
-        role = Role(name="Admin")
+        role = Role(name=name)
         db.add(role)
         db.commit()
-        admin_user = User(role_id=role.id, name="Admin", password="admin", email="admin@sample.com")
+        admin_user = User(role_id=role.id, name=name, password=password, email=email)
         db.add(admin_user)
         db.commit()
         db.remove()
-        return admin_user
+        return name, password, email
 
     def setUp(self):
-        if not os.path.exists(self.tempdir):
-            os.mkdir(self.tempdir)
+        if not os.path.exists(self.keydir):
+            os.mkdir(self.keydir)
         self.app = self.create_app()
+        self.client = self.app.test_client(self)
 
     def tearDown(self):
+        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)
+        db.remove()
+
+
+class TestApp(TestAppBase):
+
+    def setUp(self):
+        super(TestApp, self).setUp()
+
+    def tearDown(self):
+        super(TestApp, self).tearDown()
 
     def test_app_is_running(self):
         self.assertFalse(current_app is None)