Add custom module loading (#816)

tests/test_custom_modules.py

@@ -0,0 +1,193 @@
+from click.testing import CliRunner
+
+from mvt.android.cli import check_bugreport
+from mvt.android.cmd_check_androidqf import CmdAndroidCheckAndroidQF
+from mvt.android.cmd_check_backup import CmdAndroidCheckBackup
+from mvt.android.cmd_check_bugreport import CmdAndroidCheckBugreport
+from mvt.android.cmd_check_intrusion_logs import CmdAndroidCheckIntrusionLogs
+from mvt.common.module import MVTModule
+from mvt.ios.cli import check_backup, check_fs
+
+
+CUSTOM_MODULE = """
+from mvt.common.module import MVTModule
+
+
+class {name}(MVTModule):
+    supported_commands = {supported_commands!r}
+    slug = "{slug}"
+
+    def run(self):
+        self.results = [{{"message": "custom module ran"}}]
+
+    def check_indicators(self):
+        pass
+
+    def serialize(self, result):
+        return None
+"""
+
+
+def _write_custom_module(path, name, supported_commands, slug=None):
+    path.write_text(
+        CUSTOM_MODULE.format(
+            name=name,
+            supported_commands=supported_commands,
+            slug=slug or name.lower(),
+        ),
+        encoding="utf-8",
+    )
+    return path
+
+
+def test_load_module_appears_only_for_supported_cli_command(tmp_path):
+    module_path = _write_custom_module(
+        tmp_path / "custom.py",
+        "IOSBackupOnlyModule",
+        (("ios", "check-backup"),),
+    )
+
+    backup_result = CliRunner().invoke(
+        check_backup,
+        ["--list-modules", "--load-module", str(module_path), str(tmp_path)],
+    )
+    fs_result = CliRunner().invoke(
+        check_fs,
+        ["--list-modules", "--load-module", str(module_path), str(tmp_path)],
+    )
+
+    assert backup_result.exit_code == 0
+    assert "IOSBackupOnlyModule" in backup_result.output
+    assert fs_result.exit_code == 0
+    assert "IOSBackupOnlyModule" not in fs_result.output
+
+
+def test_module_option_runs_supported_custom_module(tmp_path):
+    module_path = _write_custom_module(
+        tmp_path / "custom.py",
+        "CustomRunModule",
+        (("ios", "check-backup"),),
+        slug="custom_run_module",
+    )
+    output_path = tmp_path / "out"
+
+    result = CliRunner().invoke(
+        check_backup,
+        [
+            "--module",
+            "CustomRunModule",
+            "--load-module",
+            str(module_path),
+            "--output",
+            str(output_path),
+            str(tmp_path),
+        ],
+    )
+
+    assert result.exit_code == 0
+    assert (output_path / "custom_run_module.json").exists()
+
+
+def test_custom_modules_load_from_environment_without_cli_flag(tmp_path, monkeypatch):
+    custom_modules_path = tmp_path / "custom_modules"
+    custom_modules_path.mkdir()
+    _write_custom_module(
+        custom_modules_path / "env_module.py",
+        "EnvBugreportModule",
+        (("android", "check-bugreport"),),
+    )
+    monkeypatch.setenv("MVT_CUSTOM_MODULES", str(custom_modules_path))
+
+    result = CliRunner().invoke(check_bugreport, ["--list-modules", str(tmp_path)])
+
+    assert result.exit_code == 0
+    assert "EnvBugreportModule" in result.output
+
+
+class NestedBugreportModule(MVTModule):
+    supported_commands = (("android", "check-bugreport"),)
+
+
+class NestedBackupModule(MVTModule):
+    supported_commands = (("android", "check-backup"),)
+
+
+class NestedIntrusionLogsModule(MVTModule):
+    supported_commands = (("android", "check-intrusion-logs"),)
+
+
+class NestedAndroidQFModule(MVTModule):
+    supported_commands = (("android", "check-androidqf"),)
+
+
+class DummyZip:
+    def close(self):
+        pass
+
+
+def test_androidqf_propagates_custom_modules_to_nested_commands(tmp_path, monkeypatch):
+    records = {}
+    custom_modules = [
+        NestedBugreportModule,
+        NestedBackupModule,
+        NestedIntrusionLogsModule,
+        NestedAndroidQFModule,
+    ]
+    cmd = CmdAndroidCheckAndroidQF(
+        target_path=str(tmp_path),
+        custom_modules=custom_modules,
+    )
+
+    def record_available(name):
+        def _record(command):
+            records[name] = [
+                module.__name__
+                for module in command._available_modules()
+                if module.__name__.startswith("Nested")
+            ]
+
+        return _record
+
+    monkeypatch.setattr(cmd, "load_bugreport", lambda: DummyZip())
+    monkeypatch.setattr(
+        CmdAndroidCheckBugreport,
+        "from_zip",
+        lambda self, bugreport: None,
+    )
+    monkeypatch.setattr(
+        CmdAndroidCheckBugreport,
+        "run",
+        record_available("bugreport"),
+    )
+
+    monkeypatch.setattr(cmd, "load_backup", lambda: b"")
+    monkeypatch.setattr(CmdAndroidCheckBackup, "from_ab", lambda self, backup: None)
+    monkeypatch.setattr(
+        CmdAndroidCheckBackup,
+        "run",
+        record_available("backup"),
+    )
+
+    intrusion_logs_path = tmp_path / "intrusion_logs"
+    intrusion_logs_path.mkdir()
+    setattr(cmd, "_CmdAndroidCheckAndroidQF__format", "dir")
+    setattr(
+        cmd,
+        "_CmdAndroidCheckAndroidQF__files",
+        ["androidqf/intrusion_logs/security.txt"],
+    )
+    monkeypatch.setattr(cmd, "_read_device_timezone", lambda: None)
+    monkeypatch.setattr(
+        CmdAndroidCheckIntrusionLogs,
+        "run",
+        record_available("intrusion_logs"),
+    )
+
+    assert cmd.run_bugreport_cmd()
+    assert cmd.run_backup_cmd()
+    assert cmd.run_intrusion_logs_cmd()
+    assert records == {
+        "bugreport": ["NestedBugreportModule"],
+        "backup": ["NestedBackupModule"],
+        "intrusion_logs": ["NestedIntrusionLogsModule"],
+    }