verify modules

Author: uruwhy

app/objects/c_obfuscator.py

@@ -5,6 +5,7 @@
 
 from app.objects.interfaces.i_object import FirstClassObjectInterface
 from app.utility.base_object import BaseObject
+from app.utility.base_world import BaseWorld
 
 
 class ObfuscatorSchema(ma.Schema):
@@ -28,6 +29,7 @@ def unique(self):
 
     def __init__(self, name, description, module):
         super().__init__()
+        BaseWorld.verify_module(module, 'obfuscators')
         self.name = name
         self.description = description
         self.module = module

app/objects/c_planner.py

@@ -4,6 +4,7 @@
 
 from app.objects.interfaces.i_object import FirstClassObjectInterface
 from app.utility.base_object import BaseObject
+from app.utility.base_world import BaseWorld
 from app.objects.secondclass.c_fact import Fact, FactSchema
 
 
@@ -43,6 +44,7 @@ def __init__(self, name='', planner_id='', module='', params=None, stopping_cond
         super().__init__()
         self.name = name
         self.planner_id = planner_id if planner_id else str(uuid.uuid4())
+        BaseWorld.verify_module(module, 'planners')
         self.module = module
         self.params = params if params else {}
         self.description = description

app/objects/secondclass/c_parser.py

@@ -2,6 +2,7 @@
 
 from app.objects.secondclass.c_parserconfig import ParserConfig, ParserConfigSchema
 from app.utility.base_object import BaseObject
+from app.utility.base_world import BaseWorld
 
 
 class ParserSchema(ma.Schema):
@@ -31,5 +32,6 @@ def unique(self):
 
     def __init__(self, module, parserconfigs):
         super().__init__()
+        BaseWorld.verify_module(module, 'parsers', ['app/learning'])
         self.module = module
         self.parserconfigs = parserconfigs

app/objects/secondclass/c_requirement.py

@@ -1,6 +1,7 @@
 import marshmallow as ma
 
 from app.utility.base_object import BaseObject
+from app.utility.base_world import BaseWorld
 
 
 class RequirementSchema(ma.Schema):
@@ -23,5 +24,6 @@ def unique(self):
 
     def __init__(self, module, relationship_match):
         super().__init__()
+        BaseWorld.verify_module(module, 'requirements')
         self.module = module
         self.relationship_match = relationship_match

app/utility/base_world.py

@@ -1,4 +1,5 @@
 import binascii
+import os
 import string
 import re
 import yaml
@@ -164,6 +165,35 @@ def parse_version(version_string, pattern=r'([0-9]+(?:\.[0-9]+)+)'):
             logging.getLogger('check_requirement').error(repr(e))
             return False
 
+    @staticmethod
+    def verify_module(module, module_type, additional_paths=[]):
+        '''Checks that the specified module matches the expected module path
+        given its type and that the corresponding Python file exists.
+        Raises a ModuleNotFoundError if the requested module does not meet the
+        required conditions.'''
+
+        module_path = module.replace('.', '/') + '.py'
+        module_filename = os.path.basename(module_path)
+        allowed_paths = [os.path.join('app/', module_type, module_filename)]
+        if os.path.isdir('plugins'):
+            for plugin_dir in os.listdir('plugins'):
+                if os.path.isdir(os.path.join('plugins', plugin_dir)):
+                    allowed_paths.append(
+                        os.path.join('plugins', plugin_dir, 'app', module_type, module_filename)
+                    )
+        if additional_paths:
+            allowed_paths += [os.path.join(p, module_filename) for p in additional_paths]
+        if module_path not in allowed_paths:
+            raise ModuleNotFoundError(
+                f'Module {module} does not align with allowed paths for this module type. Allowed paths for this module: {str(allowed_paths)}',
+                name=module
+            )
+        if not os.path.isfile(module_path):
+            raise ModuleNotFoundError(
+                f'Module {module} with path {module_path} was not found on disk.',
+                name=module
+            )
+
     class Access(Enum):
         APP = 0
         RED = 1

tests/utility/test_base_world.py

@@ -1,7 +1,9 @@
 import pytest
+import os
 import yaml
 
 from datetime import datetime, timezone
+from unittest import mock
 from app.utility.base_world import BaseWorld
 
 
@@ -91,3 +93,45 @@ def test_is_not_base64(self):
     def test_is_base64(self):
         b64str = 'aGVsbG8gd29ybGQgZnJvbSB1bml0IHRlc3QgbGFuZAo='
         assert BaseWorld.is_base64(b64str)
+
+    @mock.patch.object(os, 'listdir', return_value=['stockpile', 'testplugin', 'dummy'])
+    @mock.patch.object(os.path, 'isdir', return_value=True)
+    def test_verify_module(self, mock_isdir, mock_listdir):
+        def _mock_isfile(path):
+            return path in [
+                'plugins/stockpile/app/requirements/test_req.py',
+                'plugins/testplugin/app/obfuscators/test_obf.py',
+                'app/planners/test_planner.py',
+                'app/parsers/test_parser.py',
+                'app/learning/learning_parser.py',
+            ]
+
+        with mock.patch.object(os.path, 'isfile', side_effect=_mock_isfile):
+            BaseWorld.verify_module('plugins.stockpile.app.requirements.test_req', 'requirements')
+            BaseWorld.verify_module('plugins.testplugin.app.obfuscators.test_obf', 'obfuscators')
+            BaseWorld.verify_module('app.planners.test_planner', 'planners')
+            BaseWorld.verify_module('app.parsers.test_parser', 'parsers')
+            BaseWorld.verify_module('app.learning.learning_parser', 'parsers', ['app/learning'])
+
+            allowed_paths_str = str([
+                'app/parsers/myparser.py',
+                'plugins/stockpile/parsers/myparser.py',
+                'plugins/testplugin/parsers/myparser.py',
+                'plugins/dummy/parsers/myparser.py'
+            ])
+            expected_err = f'Module data.payloads.myparser does not align with allowed paths for this module type. Allowed paths for this module: {allowed_paths_str}'
+            with pytest.raises(ModuleNotFoundError, match=expected_err):
+                BaseWorld.verify_module('data.payloads.myparser', 'parsers')
+            allowed_paths_str = str([
+                'otherdir/myparser.py',
+                'app/parsers/myparser.py',
+                'plugins/stockpile/parsers/myparser.py',
+                'plugins/testplugin/parsers/myparser.py',
+                'plugins/dummy/parsers/myparser.py'
+            ])
+            expected_err = f'Module plugins.dne.myparser does not align with allowed paths for this module type. Allowed paths for this module: {allowed_paths_str}'
+            with pytest.raises(ModuleNotFoundError, match=expected_err):
+                BaseWorld.verify_module('data.payloads.myparser', 'parsers', ['otherdir'])
+            expected_err = 'Module plugins.stockpile.app.obfuscators.dne with path plugins/stockpile/app/obfuscators/dne was not found on disk.'
+            with pytest.raises(ModuleNotFoundError, match=expected_err):
+                BaseWorld.verify_module('plugins.stockpile.app.obfuscators.dne', 'obfuscators', ['otherdir'])