feat: added new unpacking plugin base class

Author: jstucke

fact_extractor/helperFunctions/plugin.py

@@ -1,30 +1,51 @@
+from __future__ import annotations
+
+import importlib.util
 import logging
+import sys
+from importlib.machinery import SourceFileLoader
+from inspect import isclass
 from pathlib import Path
+from typing import Iterable, Type
 
-from common_helper_files import get_dirs_in_dir
-from pluginbase import PluginBase
-
-from helperFunctions.file_system import get_src_dir
-
+from plugins.base_class import UnpackingPlugin
 
-def import_plugins(plugin_mount, plugin_base_dir):
-    plugin_base = PluginBase(package=plugin_mount)
-    plugin_src_dirs = _get_plugin_src_dirs(plugin_base_dir)
-    return plugin_base.make_plugin_source(searchpath=plugin_src_dirs)
+SRC_DIR = Path(__file__).parent.parent
+PLUGIN_DIR = SRC_DIR / 'plugins'
 
 
-def _get_plugin_src_dirs(base_dir):
-    plug_in_base_path = Path(get_src_dir(), base_dir)
-    plugin_dirs = get_dirs_in_dir(str(plug_in_base_path))
+def import_plugins(path=PLUGIN_DIR) -> list:
+    """Returns a list of modules where each module is an unpacking plugin."""
     plugins = []
-    for plugin_path in plugin_dirs:
-        # Ignore cache directories
-        if Path(plugin_path).name == '__pycache__':
+    for plugin_file in path.glob('**/code/*.py'):  # type: Path
+        if plugin_file.name == '__init__.py':
             continue
 
-        plugin_code_dir = Path(plugin_path, 'code')
-        if plugin_code_dir.is_dir():
-            plugins.append(str(plugin_code_dir))
-        else:
-            logging.warning(f'Plugin has no code directory: {plugin_path}')
+        module_name = str(plugin_file.relative_to(SRC_DIR)).replace('/', '.')[: -len('.py')]
+        loader = SourceFileLoader(module_name, str(plugin_file))
+        spec = importlib.util.spec_from_loader(loader.name, loader)
+        plugin_module = importlib.util.module_from_spec(spec)
+
+        sys.modules[spec.name] = plugin_module
+        try:
+            loader.exec_module(plugin_module)
+            plugins.append(plugin_module)
+        except Exception as error:  # probably missing dependencies
+            sys.modules.pop(spec.name)
+            logging.exception(f'Could not import plugin {module_name} due to exception: {error}')
+
     return plugins
+
+
+def find_plugin_classes(module) -> Iterable[Type[UnpackingPlugin]]:
+    """get all subclasses of UnpackingPlugin from the module"""
+    for attr_name in dir(module):
+        if attr_name.startswith('_'):
+            continue
+        attr = getattr(module, attr_name)
+        if _is_plugin_class(attr):
+            yield attr
+
+
+def _is_plugin_class(attr) -> bool:
+    return attr != UnpackingPlugin and isclass(attr) and issubclass(attr, UnpackingPlugin)

fact_extractor/plugins/base_class.py

@@ -0,0 +1,48 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Type
+
+
+class PluginError(Exception):
+    pass
+
+
+class UnpackingPlugin(ABC):
+    MIME_PATTERNS = ()  # must be overwritten by the subclass
+    NAME = 'base'
+    VERSION = '0.0.0'
+
+    @abstractmethod
+    def unpack_file(self, file_path: str, tmp_dir: str) -> dict:
+        """
+        Unpack a file `file_path` to `tmp_dir`. Must be implemented by the concrete unpacker subclass!
+        The function returns a dictionary with metadata from unpacking the file. The key 'output' is expected and
+        its value should contain the output of the used tool or relevant logging messages.
+
+        :param file_path: Path to the file to be unpacked.
+        :param tmp_dir: Path to the temporary directory where the file will be unpacked.
+        :return: The metadata.
+        """
+
+    def validate(self):
+        if len(self.MIME_PATTERNS) == 0:
+            raise PluginError(f'{self.NAME} is not a valid plugin (no MIME patterns defined)')
+        if self.NAME == 'base':
+            raise PluginError(f'{self.NAME} is not a valid plugin (no NAME defined)')
+        if self.VERSION == '0.0.0':
+            raise PluginError(f'{self.NAME} is not a valid plugin (no VERSION defined)')
+
+    @classmethod
+    def from_old_module(cls, old_module) -> Type[UnpackingPlugin]:
+        """For backwards compatibility with old plugins, create a subclass dynamically."""
+        return type(
+            f'{old_module.__name__}.Unpacker',
+            (cls,),
+            {
+                'NAME': old_module.NAME,
+                'VERSION': old_module.VERSION,
+                'MIME_PATTERNS': old_module.MIME_PATTERNS,
+                'unpack_file': staticmethod(old_module.unpack_function),
+            },
+        )

fact_extractor/plugins/unpacking/generic_carver/test/test_plugin_generic_carver.py

@@ -1,4 +1,3 @@
-import re
 import zlib
 from pathlib import Path
 
@@ -18,7 +17,7 @@ def test_unpacker_selection_generic(self):
     def test_extraction(self):
         in_file = f'{get_test_data_dir()}/generic_carver_test'
         files, meta_data = self.unpacker._extract_files_from_file_using_specific_unpacker(
-            in_file, self.tmp_dir.name, self.unpacker.unpacker_plugins['generic/carver']
+            in_file, self.tmp_dir.name, self.unpacker.unpacking_plugins['generic_carver']
         )
         files = set(files)
         assert len(files) == 3, 'file number incorrect'
@@ -28,7 +27,7 @@ def test_extraction(self):
     def test_zlib_carving(self):
         in_file = TEST_DATA_DIR / 'zlib_carving_test'
         files, meta_data = self.unpacker._extract_files_from_file_using_specific_unpacker(
-            in_file, self.tmp_dir.name, self.unpacker.unpacker_plugins['generic/carver']
+            in_file, self.tmp_dir.name, self.unpacker.unpacking_plugins['generic_carver']
         )
         assert len(files) == 9, 'file number incorrect'
         assert sum(1 if f.endswith('.zlib_carver') else 0 for f in files), 'wrong number of carved zlib streams'
@@ -40,7 +39,7 @@ def test_filter(self):
         in_file = TEST_DATA_DIR / 'carving_test_file'
         assert Path(in_file).is_file()
         files, meta_data = self.unpacker._extract_files_from_file_using_specific_unpacker(
-            str(in_file), self.tmp_dir.name, self.unpacker.unpacker_plugins['generic/carver']
+            str(in_file), self.tmp_dir.name, self.unpacker.unpacking_plugins['generic_carver']
         )
         files = set(files)
         assert len(files) == 4, 'file number incorrect'
@@ -53,12 +52,11 @@ def test_fake_archives(self, file_format):
         in_file = TEST_DATA_DIR / f'fake_{file_format}.{file_format}'
         assert Path(in_file).is_file()
         meta = unpack_function(str(in_file), self.tmp_dir.name)
-        assert meta == {'output': 'No valid chunks found.'}
+        assert 'No valid chunks found.' in meta['output']
 
     @pytest.mark.parametrize(('file_format', 'expected_size'), [('bz2', 52), ('zip', 170)])
     def test_trailing_data(self, file_format, expected_size):
         in_file = Path(get_test_data_dir()) / f'trailing_data.{file_format}'
         assert Path(in_file).is_file()
         meta = unpack_function(str(in_file), self.tmp_dir.name)
-        carved_size = int(re.search(r'size: (\d+)', meta['output']).group(1))
-        assert carved_size == expected_size
+        assert f'chunk: 0x0-{hex(expected_size)}' in meta['output']

fact_extractor/plugins/unpacking/raw/test/test_raw.py

@@ -22,7 +22,7 @@ def test_extraction(self):
     def test_extraction_encoded(self):
         input_file = Path(TEST_DATA_DIR, 'encoded.bin')
         unpacked_files, meta_data = self.unpacker._extract_files_from_file_using_specific_unpacker(
-            str(input_file), self.tmp_dir.name, self.unpacker.unpacker_plugins['data/raw']
+            str(input_file), self.tmp_dir.name, self.unpacker.unpacking_plugins['RAW']
         )
         assert meta_data['Intel Hex'] == 1
         assert meta_data['Motorola S-Record'] == 1

fact_extractor/plugins/unpacking/senao/test/test_plugin_senao.py

@@ -15,7 +15,7 @@ def test_extraction(self):
         files, meta_data = self.unpacker._extract_files_from_file_using_specific_unpacker(
             str(in_file),
             self.tmp_dir.name,
-            self.unpacker.unpacker_plugins['firmware/senao-v2b'],
+            self.unpacker.unpacking_plugins['senao'],
         )
         assert len(files) == 1, 'unpacked file number incorrect'
         file = Path(files[0])

fact_extractor/test/data/plugin_system/plugin_one/code/plugin_one.py

@@ -1,2 +1,7 @@
-def plugin_init():
-    return 'plugin_loaded'
+NAME = 'plugin_one'
+VERSION = '0.0.1'
+MIME_PATTERNS = ('foo/bar',)
+
+
+def unpack_function(file_path, tmp_dir):  # noqa: ARG001
+    pass

fact_extractor/test/data/plugin_system/plugin_two/code/plugin_two.py

@@ -0,0 +1,10 @@
+from plugins.base_class import UnpackingPlugin
+
+
+class UnpackerTwo(UnpackingPlugin):
+    NAME = 'plugin_two'
+    VERSION = '0.0.1'
+    MIME_PATTERNS = ('test/123',)
+
+    def unpack_file(self, file_path, tmp_dir):
+        pass

fact_extractor/test/unit/helperFunctions/test_plugin.py

@@ -1,17 +1,15 @@
-from helperFunctions.plugin import _get_plugin_src_dirs, import_plugins
+from pathlib import Path
 
-TEST_PLUGINS_BASE_PATH = 'test/data/plugin_system'
+from helperFunctions.plugin import find_plugin_classes, import_plugins
 
+TEST_PLUGINS_BASE_PATH = Path(__file__).parent.parent.parent / 'data/plugin_system'
 
-class TestHelperFunctionsPlugin:
-    def test_get_plugin_src_dirs(self):
-        result = _get_plugin_src_dirs(TEST_PLUGINS_BASE_PATH)
-        assert isinstance(result, list), 'result is not a list'
-        assert 'plugin_one' in sorted(result)[0], 'plugin not found'
-        assert len(result) == 2, 'number of found plugin directories not correct'
 
-    def test_load_plugins(self):
-        result = import_plugins('plugins.test', TEST_PLUGINS_BASE_PATH)
-        imported_plugins = result.list_plugins()
-        assert len(imported_plugins) == 1, 'wrong number of plugins imported'
-        assert imported_plugins[0] == 'plugin_one', 'plugin name not correct'
+def test_load_plugins():
+    imported_modules = import_plugins(TEST_PLUGINS_BASE_PATH)
+    assert len(imported_modules) == 2, 'wrong number of plugins imported'
+    modules = {m.__name__.split('.')[-1]: m for m in imported_modules}
+    assert 'plugin_one' in modules
+    assert 'plugin_two' in modules
+    assert len(list(find_plugin_classes(modules['plugin_one']))) == 0
+    assert len(list(find_plugin_classes(modules['plugin_two']))) == 1

fact_extractor/test/unit/unpacker/test_unpacker.py

@@ -2,51 +2,59 @@
 
 from __future__ import annotations
 
-import gc
 import json
-import os
 import shutil
 from configparser import ConfigParser
 from pathlib import Path
 from tempfile import TemporaryDirectory
 from unittest.mock import Mock, patch
 
 from helperFunctions.file_system import get_test_data_dir
+from plugins.base_class import UnpackingPlugin
 from unpacker.unpack import Unpacker
 
 
 class TestUnpackerBase:
-    def setup_method(self):
-        self.config = ConfigParser()
-        self.ds_tmp_dir = TemporaryDirectory(prefix='fact_tests_')
-        self.tmp_dir = TemporaryDirectory(prefix='fact_tests_')
+    @classmethod
+    def setup_class(cls):
+        cls.config = ConfigParser()
+        cls.ds_tmp_dir = TemporaryDirectory(prefix='fact_tests_')
+        cls.files_dir = Path(cls.ds_tmp_dir.name) / 'files'
+        cls.reports_dir = Path(cls.ds_tmp_dir.name) / 'reports'
+
+        cls.config.add_section('unpack')
+        cls.config.set('unpack', 'data_folder', cls.ds_tmp_dir.name)
+        cls.config.set('unpack', 'blacklist', 'text/plain, image/png')
+        cls.config.add_section('ExpertSettings')
+        cls.config.set('ExpertSettings', 'header_overhead', '256')
+        cls.config.set('ExpertSettings', 'unpack_threshold', '0.8')
 
-        self.config.add_section('unpack')
-        self.config.set('unpack', 'data_folder', self.ds_tmp_dir.name)
-        self.config.set('unpack', 'blacklist', 'text/plain, image/png')
-        self.config.add_section('ExpertSettings')
-        self.config.set('ExpertSettings', 'header_overhead', '256')
-        self.config.set('ExpertSettings', 'unpack_threshold', '0.8')
+        cls.unpacker = Unpacker(config=cls.config)
 
-        self.unpacker = Unpacker(config=self.config)
-        os.makedirs(str(self.unpacker._report_folder), exist_ok=True)  # pylint: disable=protected-access
-        os.makedirs(str(self.unpacker._file_folder), exist_ok=True)  # pylint: disable=protected-access
+        cls.test_file_path = Path(get_test_data_dir(), 'get_files_test/testfile1')
 
-        self.test_file_path = Path(get_test_data_dir(), 'get_files_test/testfile1')
+    def setup_method(self):
+        self.tmp_dir = TemporaryDirectory(prefix='fact_tests_')
+        self.unpacker._report_folder.mkdir(parents=True, exist_ok=True)
+        self.unpacker._file_folder.mkdir(parents=True, exist_ok=True)
 
     def teardown_method(self):
-        self.ds_tmp_dir.cleanup()
         self.tmp_dir.cleanup()
-        gc.collect()
+        shutil.rmtree(self.unpacker._report_folder)
+        shutil.rmtree(self.unpacker._file_folder)
+
+    @classmethod
+    def teardown_class(cls):
+        cls.ds_tmp_dir.cleanup()
 
     def get_unpacker_meta(self):
         return json.loads(
             Path(self.unpacker._report_folder, 'meta.json').read_text()  # pylint: disable=protected-access
         )
 
     def check_unpacker_selection(self, mime_type, plugin_name):
-        name = self.unpacker.get_unpacker(mime_type)[1]
-        assert name == plugin_name, 'wrong unpacker plugin selected'
+        unpacker = self.unpacker.get_unpacker(mime_type)
+        assert unpacker.NAME == plugin_name, 'wrong unpacker plugin selected'  # noqa: SIM300
 
     def check_unpacking_of_standard_unpack_set(
         self,
@@ -59,9 +67,9 @@ def check_unpacking_of_standard_unpack_set(
         files = {f for f in files if not any(rule in f for rule in ignore or set())}
         assert len(files) == 3, f'file number incorrect: {meta_data}'
         assert files == {
-            os.path.join(self.tmp_dir.name, additional_prefix_folder, 'testfile1'),
-            os.path.join(self.tmp_dir.name, additional_prefix_folder, 'testfile2'),
-            os.path.join(self.tmp_dir.name, additional_prefix_folder, 'generic folder/test file 3_.txt'),
+            str(Path(self.tmp_dir.name, additional_prefix_folder, 'testfile1')),
+            str(Path(self.tmp_dir.name, additional_prefix_folder, 'testfile2')),
+            str(Path(self.tmp_dir.name, additional_prefix_folder, 'generic folder/test file 3_.txt')),
         }, f'not all files found: {meta_data}'
         if output:
             assert 'output' in meta_data
@@ -70,9 +78,11 @@ def check_unpacking_of_standard_unpack_set(
 
 class TestUnpackerCore(TestUnpackerBase):
     def test_generic_carver_found(self):
-        assert 'generic/carver' in list(self.unpacker.unpacker_plugins), 'generic carver plugin not found'
-        name = self.unpacker.unpacker_plugins['generic/carver'][1]
-        assert name == 'generic_carver', 'generic_carver plugin not found'
+        assert 'generic_carver' in self.unpacker.unpacking_plugins, 'generic carver plugin not found'
+        assert 'generic/carver' in self.unpacker.plugin_by_mime, 'generic carver MIME type not found'
+        plugin = self.unpacker.get_unpacker('generic/carver')
+        assert isinstance(plugin, UnpackingPlugin)
+        assert plugin.NAME == 'generic_carver', 'generic_carver plugin not found'
 
     def test_unpacker_selection_unknown(self):
         self.check_unpacker_selection('unknown/blah', 'generic_carver')
@@ -167,17 +177,17 @@ def test_main_unpack_function(self):
         test_file_path = Path(get_test_data_dir(), 'container/test.zip')
         self.main_unpack_check(test_file_path, 3, 0, '7z')
 
-    def test_main_unpack_exclude_archive(self):
+    def test_main_unpack_exclude_archive(self, monkeypatch):
         test_file_path = Path(get_test_data_dir(), 'container/test.zip')
-        self.unpacker.exclude = ['*test.zip']
+        monkeypatch.setattr(self.unpacker, 'exclude', ['*test.zip'])
         self.main_unpack_check(test_file_path, 0, 1, None)
 
-    def test_main_unpack_exclude_subdirectory(self):
+    def test_main_unpack_exclude_subdirectory(self, monkeypatch):
         test_file_path = Path(get_test_data_dir(), 'container/test.zip')
-        self.unpacker.exclude = ['*/generic folder/*']
+        monkeypatch.setattr(self.unpacker, 'exclude', ['*/generic folder/*'])
         self.main_unpack_check(test_file_path, 2, 1, '7z')
 
-    def test_main_unpack_exclude_files(self):
+    def test_main_unpack_exclude_files(self, monkeypatch):
         test_file_path = Path(get_test_data_dir(), 'container/test.zip')
-        self.unpacker.exclude = ['*/get_files_test/*test*']
+        monkeypatch.setattr(self.unpacker, 'exclude', ['*/get_files_test/*test*'])
         self.main_unpack_check(test_file_path, 0, 3, '7z')