mvt/tests/test_custom_modules.py

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"],
    }