Add custom module loading (#816)

Author: besendorf

README.md

@@ -60,6 +60,11 @@ For alternative installation options and known issues, please refer to the [docu
 
 MVT provides two commands `mvt-ios` and `mvt-android`. [Check out the documentation to learn how to use them!](https://docs.mvt.re/)
 
+Module-running `check-*` commands can load custom Python modules with
+`--load-module PATH` or from a folder set in `MVT_CUSTOM_MODULES`. See the
+[development documentation](https://docs.mvt.re/en/latest/development/) for
+details.
+
 
 ## License
 

docs/development.md

@@ -39,6 +39,82 @@ Selecting a single module also runs its transitive dependencies. If a dependency
 is unavailable or the dependency graph contains a cycle, the command logs a
 warning and does not run any modules.
 
+## Custom modules
+
+Module-running `check-*` commands can load custom modules from Python files that
+are not installed as part of MVT. Load one file with:
+
+```bash
+mvt-ios check-backup --load-module ./example_module.py --output ./out ./backup
+```
+
+You can also load a folder. MVT loads non-hidden top-level `*.py` files in
+sorted order and skips `__init__.py`:
+
+```bash
+mvt-ios check-fs --load-module ./custom_modules ./filesystem-dump
+```
+
+Set `MVT_CUSTOM_MODULES` to load a folder for every module-running command. This
+folder is loaded before any `--load-module` path:
+
+```bash
+MVT_CUSTOM_MODULES=./custom_modules mvt-android check-bugreport ./bugreport.zip
+```
+
+Custom modules are normal `MVTModule` subclasses:
+
+```python
+from mvt.common.module import MVTModule
+
+
+class ExampleCustomModule(MVTModule):
+    supported_commands = (("ios", "check-backup"), ("ios", "check-fs"))
+    slug = "example_custom_module"
+
+    def run(self):
+        self.results = [{"message": "custom module ran"}]
+
+    def check_indicators(self):
+        pass
+
+    def serialize(self, result):
+        return None
+```
+
+Use `supported_commands` to restrict a module to specific platform/command
+pairs. Missing or empty `supported_commands` means the module is available to
+all commands, which keeps older modules compatible. Supported pairs are:
+
+```python
+("ios", "check-backup")
+("ios", "check-fs")
+("ios", "check-iocs")
+("android", "check-backup")
+("android", "check-bugreport")
+("android", "check-androidqf")
+("android", "check-intrusion-logs")
+("android", "check-iocs")
+```
+
+Custom modules can depend on existing MVT module classes. Dependencies are
+resolved with the same ordering logic as built-in modules, and custom modules
+are appended after built-ins before ordering:
+
+```python
+from mvt.common.module import MVTModule
+from mvt.ios.modules.backup.manifest import Manifest
+
+
+class DependentCustomModule(MVTModule):
+    supported_commands = (("ios", "check-backup"),)
+    dependencies = (Manifest,)
+
+    def run(self):
+        manifest_results = self.get_dependency_results(Manifest)
+        self.results = [{"manifest_entries": len(manifest_results)}]
+```
+
 ## Profiling
 
 Some MVT modules extract and process significant amounts of data during the analysis process or while checking results against known indicators. Care must be

src/mvt/android/cli.py

@@ -22,6 +22,7 @@
     HELP_MSG_HASHES,
     HELP_MSG_IOC,
     HELP_MSG_LIST_MODULES,
+    HELP_MSG_LOAD_MODULE,
     HELP_MSG_MODULE,
     HELP_MSG_NONINTERACTIVE,
     HELP_MSG_OUTPUT,
@@ -30,6 +31,7 @@
     HELP_MSG_VERSION,
 )
 from mvt.common.logo import logo
+from mvt.common.module_loader import CustomModuleLoadError, load_custom_modules
 from mvt.common.updates import IndicatorsUpdates
 from mvt.common.utils import init_logging, set_verbose_logging
 
@@ -59,6 +61,13 @@ def _get_disable_flags(ctx):
     )
 
 
+def _load_custom_modules(load_module):
+    try:
+        return load_custom_modules(load_module)
+    except CustomModuleLoadError as exc:
+        raise click.ClickException(str(exc)) from exc
+
+
 # ==============================================================================
 # Main
 # ==============================================================================
@@ -119,11 +128,28 @@ def check_adb(ctx):
 @click.option("--output", "-o", type=click.Path(exists=False), help=HELP_MSG_OUTPUT)
 @click.option("--list-modules", "-l", is_flag=True, help=HELP_MSG_LIST_MODULES)
 @click.option("--module", "-m", help=HELP_MSG_MODULE)
+@click.option(
+    "--load-module",
+    type=click.Path(exists=True),
+    multiple=True,
+    default=[],
+    help=HELP_MSG_LOAD_MODULE,
+)
 @click.option("--verbose", "-v", is_flag=True, help=HELP_MSG_VERBOSE)
 @click.argument("BUGREPORT_PATH", type=click.Path(exists=True))
 @click.pass_context
-def check_bugreport(ctx, iocs, output, list_modules, module, verbose, bugreport_path):
+def check_bugreport(
+    ctx,
+    iocs,
+    output,
+    list_modules,
+    module,
+    load_module,
+    verbose,
+    bugreport_path,
+):
     set_verbose_logging(verbose)
+    custom_modules = _load_custom_modules(load_module)
     # Always generate hashes as bug reports are small.
     cmd = CmdAndroidCheckBugreport(
         target_path=bugreport_path,
@@ -133,6 +159,7 @@ def check_bugreport(ctx, iocs, output, list_modules, module, verbose, bugreport_
         hashes=True,
         disable_version_check=_get_disable_flags(ctx)[0],
         disable_indicator_check=_get_disable_flags(ctx)[1],
+        custom_modules=custom_modules,
     )
 
     if list_modules:
@@ -164,6 +191,13 @@ def check_bugreport(ctx, iocs, output, list_modules, module, verbose, bugreport_
 )
 @click.option("--output", "-o", type=click.Path(exists=False), help=HELP_MSG_OUTPUT)
 @click.option("--list-modules", "-l", is_flag=True, help=HELP_MSG_LIST_MODULES)
+@click.option(
+    "--load-module",
+    type=click.Path(exists=True),
+    multiple=True,
+    default=[],
+    help=HELP_MSG_LOAD_MODULE,
+)
 @click.option("--non-interactive", "-n", is_flag=True, help=HELP_MSG_NONINTERACTIVE)
 @click.option("--backup-password", "-p", help=HELP_MSG_ANDROID_BACKUP_PASSWORD)
 @click.option("--verbose", "-v", is_flag=True, help=HELP_MSG_VERBOSE)
@@ -174,12 +208,14 @@ def check_backup(
     iocs,
     output,
     list_modules,
+    load_module,
     non_interactive,
     backup_password,
     verbose,
     backup_path,
 ):
     set_verbose_logging(verbose)
+    custom_modules = _load_custom_modules(load_module)
 
     # Always generate hashes as backups are generally small.
     cmd = CmdAndroidCheckBackup(
@@ -193,6 +229,7 @@ def check_backup(
         },
         disable_version_check=_get_disable_flags(ctx)[0],
         disable_indicator_check=_get_disable_flags(ctx)[1],
+        custom_modules=custom_modules,
     )
 
     if list_modules:
@@ -223,6 +260,13 @@ def check_backup(
 @click.option("--output", "-o", type=click.Path(exists=False), help=HELP_MSG_OUTPUT)
 @click.option("--list-modules", "-l", is_flag=True, help=HELP_MSG_LIST_MODULES)
 @click.option("--module", "-m", help=HELP_MSG_MODULE)
+@click.option(
+    "--load-module",
+    type=click.Path(exists=True),
+    multiple=True,
+    default=[],
+    help=HELP_MSG_LOAD_MODULE,
+)
 @click.option("--hashes", "-H", is_flag=True, help=HELP_MSG_HASHES)
 @click.option("--non-interactive", "-n", is_flag=True, help=HELP_MSG_NONINTERACTIVE)
 @click.option("--backup-password", "-p", help=HELP_MSG_ANDROID_BACKUP_PASSWORD)
@@ -235,13 +279,15 @@ def check_androidqf(
     output,
     list_modules,
     module,
+    load_module,
     hashes,
     non_interactive,
     backup_password,
     verbose,
     androidqf_path,
 ):
     set_verbose_logging(verbose)
+    custom_modules = _load_custom_modules(load_module)
 
     cmd = CmdAndroidCheckAndroidQF(
         target_path=androidqf_path,
@@ -255,6 +301,7 @@ def check_androidqf(
         },
         disable_version_check=_get_disable_flags(ctx)[0],
         disable_indicator_check=_get_disable_flags(ctx)[1],
+        custom_modules=custom_modules,
     )
 
     if list_modules:
@@ -288,6 +335,13 @@ def check_androidqf(
 @click.option("--output", "-o", type=click.Path(exists=False), help=HELP_MSG_OUTPUT)
 @click.option("--list-modules", "-l", is_flag=True, help=HELP_MSG_LIST_MODULES)
 @click.option("--module", "-m", help=HELP_MSG_MODULE)
+@click.option(
+    "--load-module",
+    type=click.Path(exists=True),
+    multiple=True,
+    default=[],
+    help=HELP_MSG_LOAD_MODULE,
+)
 @click.option(
     "--timezone",
     "-t",
@@ -307,11 +361,13 @@ def check_intrusion_logs(
     output,
     list_modules,
     module,
+    load_module,
     timezone,
     verbose,
     logs_path,
 ):
     set_verbose_logging(verbose)
+    custom_modules = _load_custom_modules(load_module)
 
     module_options = {}
     if timezone:
@@ -325,6 +381,7 @@ def check_intrusion_logs(
         module_options=module_options if module_options else None,
         disable_version_check=_get_disable_flags(ctx)[0],
         disable_indicator_check=_get_disable_flags(ctx)[1],
+        custom_modules=custom_modules,
     )
 
     if list_modules:
@@ -352,15 +409,25 @@ def check_intrusion_logs(
 )
 @click.option("--list-modules", "-l", is_flag=True, help=HELP_MSG_LIST_MODULES)
 @click.option("--module", "-m", help=HELP_MSG_MODULE)
+@click.option(
+    "--load-module",
+    type=click.Path(exists=True),
+    multiple=True,
+    default=[],
+    help=HELP_MSG_LOAD_MODULE,
+)
 @click.argument("FOLDER", type=click.Path(exists=True))
 @click.pass_context
-def check_iocs(ctx, iocs, list_modules, module, folder):
+def check_iocs(ctx, iocs, list_modules, module, load_module, folder):
+    custom_modules = _load_custom_modules(load_module)
     cmd = CmdCheckIOCS(
         target_path=folder,
         ioc_files=iocs,
         module_name=module,
         disable_version_check=_get_disable_flags(ctx)[0],
         disable_indicator_check=_get_disable_flags(ctx)[1],
+        custom_modules=custom_modules,
+        platform="android",
     )
     cmd.modules = (
         BACKUP_MODULES + BUGREPORT_MODULES + ANDROIDQF_MODULES + INTRUSION_LOGS_MODULES

src/mvt/android/cmd_check_androidqf.py

@@ -17,6 +17,7 @@
 from mvt.android.cmd_check_bugreport import CmdAndroidCheckBugreport
 from mvt.common.command import Command
 from mvt.common.indicators import Indicators
+from mvt.common.module import MVTModule
 
 from .modules.androidqf import ANDROIDQF_MODULES
 from .modules.androidqf.base import AndroidQFModule
@@ -50,6 +51,7 @@ def __init__(
         sub_command: Optional[bool] = False,
         disable_version_check: bool = False,
         disable_indicator_check: bool = False,
+        custom_modules: Optional[list[type[MVTModule]]] = None,
     ) -> None:
         super().__init__(
             target_path=target_path,
@@ -64,8 +66,10 @@ def __init__(
             log=log,
             disable_version_check=disable_version_check,
             disable_indicator_check=disable_indicator_check,
+            custom_modules=custom_modules,
         )
 
+        self.platform = "android"
         self.name = "check-androidqf"
         self.modules = ANDROIDQF_MODULES
 
@@ -210,6 +214,7 @@ def run_bugreport_cmd(self) -> bool:
                 module_options=self.module_options,
                 hashes=self.hashes,
                 sub_command=True,
+                custom_modules=self.custom_modules,
             )
             cmd.from_zip(bugreport)
             cmd.run()
@@ -239,6 +244,7 @@ def run_backup_cmd(self) -> bool:
             module_options=self.module_options,
             hashes=self.hashes,
             sub_command=True,
+            custom_modules=self.custom_modules,
         )
         try:
             cmd.from_ab(backup)
@@ -318,6 +324,7 @@ def run_intrusion_logs_cmd(self) -> bool:
                 module_options=adv_module_options,
                 hashes=self.hashes,
                 sub_command=True,
+                custom_modules=self.custom_modules,
             )
             cmd.run()
 

src/mvt/android/cmd_check_backup.py

@@ -21,6 +21,7 @@
 )
 from mvt.common.command import Command
 from mvt.common.indicators import Indicators
+from mvt.common.module import MVTModule
 
 from .modules.backup import BACKUP_MODULES
 
@@ -45,6 +46,7 @@ def __init__(
         sub_command: Optional[bool] = False,
         disable_version_check: bool = False,
         disable_indicator_check: bool = False,
+        custom_modules: Optional[list[type[MVTModule]]] = None,
     ) -> None:
         super().__init__(
             target_path=target_path,
@@ -59,8 +61,10 @@ def __init__(
             log=log,
             disable_version_check=disable_version_check,
             disable_indicator_check=disable_indicator_check,
+            custom_modules=custom_modules,
         )
 
+        self.platform = "android"
         self.name = "check-backup"
         self.modules = BACKUP_MODULES