UT: Add SMB plugin tests

monkey/tests/unit_tests/agent_plugins/exploiters/smb/test_smb_plugin.py

@@ -0,0 +1,109 @@
+from ipaddress import IPv4Address
+from threading import Event
+from unittest.mock import MagicMock
+from uuid import UUID
+
+import pytest
+from agent_plugins.exploiters.smb.src.plugin import Plugin
+from agent_plugins.exploiters.smb.src.smb_exploiter import SMBExploiter
+
+from common import OperatingSystem
+from infection_monkey.i_puppet import ExploiterResultData, TargetHost
+from infection_monkey.propagation_credentials_repository import IPropagationCredentialsRepository
+
+AGENT_ID = UUID("5c145d4e-ec61-44f7-998e-17477112f50f")
+BAD_SMB_OPTIONS_DICT = {"blah": "blah"}
+TARGET_IP = IPv4Address("1.1.1.1")
+TARGET_HOST = TargetHost(ip=TARGET_IP, operating_system=OperatingSystem.WINDOWS)
+SERVERS = ["10.10.10.10"]
+EXPLOITER_RESULT_DATA = ExploiterResultData(True, False, error_message="Test error")
+
+
+@pytest.fixture
+def propagation_credentials_repository():
+    return MagicMock(spec=IPropagationCredentialsRepository)
+
+
+class MockSMBExploiter(SMBExploiter):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def exploit_host(self, *args, **kwargs) -> ExploiterResultData:
+        return EXPLOITER_RESULT_DATA
+
+
+class ErrorRaisingMockSMBExploiter(SMBExploiter):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def exploit_host(self, *args, **kwargs) -> ExploiterResultData:
+        raise Exception("Test error")
+
+
+@pytest.fixture
+def plugin(
+    monkeypatch, propagation_credentials_repository: IPropagationCredentialsRepository
+) -> Plugin:
+    monkeypatch.setattr("agent_plugins.exploiters.smb.src.plugin.SMBExploiter", MockSMBExploiter)
+
+    return Plugin(
+        plugin_name="SMB",
+        agent_id=AGENT_ID,
+        agent_event_publisher=MagicMock(),
+        agent_binary_repository=MagicMock(),
+        propagation_credentials_repository=propagation_credentials_repository,
+    )
+
+
+def test_run__fails_on_bad_options(plugin: Plugin):
+    result = plugin.run(
+        host=TARGET_HOST,
+        servers=SERVERS,
+        current_depth=1,
+        options=BAD_SMB_OPTIONS_DICT,
+        interrupt=Event(),
+    )
+
+    assert not result.exploitation_success
+    assert not result.propagation_success
+
+
+def test_run__returns_exploiter_result_data(plugin: Plugin):
+    result = plugin.run(
+        host=TARGET_HOST,
+        servers=SERVERS,
+        current_depth=1,
+        options={},
+        interrupt=Event(),
+    )
+
+    assert result == EXPLOITER_RESULT_DATA
+
+
+def test_run__exploit_host_raises_exception(
+    monkeypatch,
+    plugin: Plugin,
+    propagation_credentials_repository: IPropagationCredentialsRepository,
+):
+    monkeypatch.setattr(
+        "agent_plugins.exploiters.smb.src.plugin.SMBExploiter",
+        ErrorRaisingMockSMBExploiter,
+    )
+
+    plugin = Plugin(
+        plugin_name="SMB",
+        agent_id=AGENT_ID,
+        agent_event_publisher=MagicMock(),
+        agent_binary_repository=MagicMock(),
+        propagation_credentials_repository=propagation_credentials_repository,
+    )
+    result = plugin.run(
+        host=TARGET_HOST,
+        servers=SERVERS,
+        current_depth=1,
+        options={},
+        interrupt=Event(),
+    )
+
+    assert not result.exploitation_success
+    assert not result.propagation_success