Merge pull request #2389 from guardicore/2269-notify-relay-on-propagation

2269 notify relay on propagation

monkey/infection_monkey/agent_event_handlers/__init__.py

@@ -0,0 +1 @@
+from .notify_relay_on_propagation import notify_relay_on_propagation

monkey/infection_monkey/agent_event_handlers/notify_relay_on_propagation.py

@@ -0,0 +1,31 @@
+import logging
+from typing import Optional
+
+from common.agent_events import PropagationEvent
+from infection_monkey.network.relay import TCPRelay
+
+logger = logging.getLogger(__name__)
+
+
+class notify_relay_on_propagation:
+    """
+    Notifies a TCPRelay of potential relay users if propagation is successful
+    """
+
+    def __init__(self, tcp_relay: Optional[TCPRelay]):
+        """
+        :param tcp_relay: A TCPRelay to notify on successful propagation
+        """
+        self._tcp_relay = tcp_relay
+
+    def __call__(self, event: PropagationEvent):
+        """
+        Notify a TCPRelay of potential relay users if propagation is successful
+
+        :param event: A `PropagationEvent`
+        """
+        if self._tcp_relay is None:
+            return
+
+        if event.success:
+            self._tcp_relay.add_potential_user(event.target)

monkey/infection_monkey/monkey.py

@@ -13,7 +13,7 @@ from common.agent_event_serializers import (
     AgentEventSerializerRegistry,
     register_common_agent_event_serializers,
 )
-from common.agent_events import CredentialsStolenEvent
+from common.agent_events import CredentialsStolenEvent, PropagationEvent
 from common.agent_registration_data import AgentRegistrationData
 from common.event_queue import IAgentEventQueue, PyPubSubAgentEventQueue
 from common.network.network_utils import get_my_ip_addresses, get_network_interfaces
@@ -22,6 +22,7 @@ from common.utils.argparse_types import positive_int
 from common.utils.attack_utils import ScanStatus, UsageEnum
 from common.version import get_version
 from infection_monkey.agent_event_forwarder import AgentEventForwarder
+from infection_monkey.agent_event_handlers import notify_relay_on_propagation
 from infection_monkey.config import GUID
 from infection_monkey.control import ControlClient
 from infection_monkey.credential_collectors import (
@@ -80,9 +81,6 @@ from infection_monkey.puppet.puppet import Puppet
 from infection_monkey.system_singleton import SystemSingleton
 from infection_monkey.telemetry.attack.t1106_telem import T1106Telem
 from infection_monkey.telemetry.attack.t1107_telem import T1107Telem
-from infection_monkey.telemetry.messengers.exploit_intercepting_telemetry_messenger import (
-    ExploitInterceptingTelemetryMessenger,
-)
 from infection_monkey.telemetry.messengers.legacy_telemetry_messenger_adapter import (
     LegacyTelemetryMessengerAdapter,
 )
@@ -278,15 +276,11 @@ class InfectionMonkey:
 
         victim_host_factory = self._build_victim_host_factory(local_network_interfaces)
 
-        telemetry_messenger = ExploitInterceptingTelemetryMessenger(
-            self._telemetry_messenger, self._relay
-        )
-
         self._master = AutomatedMaster(
             self._current_depth,
             servers,
             puppet,
-            telemetry_messenger,
+            self._telemetry_messenger,
             victim_host_factory,
             self._control_channel,
             local_network_interfaces,
@@ -313,6 +307,7 @@ class InfectionMonkey:
         agent_event_queue.subscribe_all_events(
             AgentEventForwarder(self._island_api_client, agent_event_serializer_registry).send_event
         )
+        agent_event_queue.subscribe_type(PropagationEvent, notify_relay_on_propagation(self._relay))
 
     def _build_puppet(
         self,

monkey/infection_monkey/telemetry/messengers/exploit_intercepting_telemetry_messenger.py

@@ -1,41 +0,0 @@
-from functools import singledispatch
-from ipaddress import IPv4Address
-
-from infection_monkey.network.relay import TCPRelay
-from infection_monkey.telemetry.exploit_telem import ExploitTelem
-from infection_monkey.telemetry.i_telem import ITelem
-from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger
-
-
-class ExploitInterceptingTelemetryMessenger(ITelemetryMessenger):
-    def __init__(self, telemetry_messenger: ITelemetryMessenger, relay: TCPRelay):
-        self._telemetry_messenger = telemetry_messenger
-        self._relay = relay
-
-    def send_telemetry(self, telemetry: ITelem):
-        _send_telemetry(telemetry, self._telemetry_messenger, self._relay)
-
-
-# Note: We can use @singledispatchmethod instead of @singledispatch if we migrate to Python 3.8 or
-# later.
-@singledispatch
-def _send_telemetry(
-    telemetry: ITelem,
-    telemetry_messenger: ITelemetryMessenger,
-    relay: TCPRelay,
-):
-    telemetry_messenger.send_telemetry(telemetry)
-
-
-@_send_telemetry.register
-def _(
-    telemetry: ExploitTelem,
-    telemetry_messenger: ITelemetryMessenger,
-    relay: TCPRelay,
-):
-    if telemetry.propagation_result is True:
-        if relay:
-            address = IPv4Address(str(telemetry.host["ip_addr"]))
-            relay.add_potential_user(address)
-
-    telemetry_messenger.send_telemetry(telemetry)

monkey/tests/unit_tests/infection_monkey/agent_event_handlers/test_notify_relay_on_propagation.py

@@ -0,0 +1,54 @@
+from ipaddress import IPv4Address
+from unittest.mock import MagicMock
+from uuid import UUID
+
+import pytest
+
+from common.agent_events import PropagationEvent
+from infection_monkey.agent_event_handlers import notify_relay_on_propagation
+from infection_monkey.network.relay import TCPRelay
+
+TARGET_ADDRESS = IPv4Address("192.168.1.10")
+
+SUCCESSFUL_PROPAGATION_EVENT = PropagationEvent(
+    source=UUID("f811ad00-5a68-4437-bd51-7b5cc1768ad5"),
+    target=TARGET_ADDRESS,
+    tags=frozenset({"test"}),
+    success=True,
+    exploiter_name="test_exploiter",
+)
+
+FAILED_PROPAGATION_EVENT = PropagationEvent(
+    source=UUID("f811ad00-5a68-4437-bd51-7b5cc1768ad5"),
+    target=TARGET_ADDRESS,
+    tags=frozenset({"test"}),
+    success=False,
+    exploiter_name="test_exploiter",
+    error_message="everything is broken",
+)
+
+
+@pytest.fixture
+def mock_tcp_relay():
+    return MagicMock(spec=TCPRelay)
+
+
+def test_relay_notified_on_successful_propation(mock_tcp_relay):
+    handler = notify_relay_on_propagation(mock_tcp_relay)
+    handler(SUCCESSFUL_PROPAGATION_EVENT)
+
+    mock_tcp_relay.add_potential_user.assert_called_once_with(TARGET_ADDRESS)
+
+
+def test_relay_not_notified_on_successful_propation(mock_tcp_relay):
+    handler = notify_relay_on_propagation(mock_tcp_relay)
+    handler(FAILED_PROPAGATION_EVENT)
+
+    mock_tcp_relay.add_potential_user.assert_not_called()
+
+
+def test_handler_doesnt_raise_if_relay_is_none():
+    handler = notify_relay_on_propagation(None)
+
+    # Raises AttributeError on failure
+    handler(SUCCESSFUL_PROPAGATION_EVENT)

monkey/tests/unit_tests/infection_monkey/telemetry/messengers/test_exploit_intercepting_telemetry_messenger.py

@@ -1,62 +0,0 @@
-from unittest.mock import MagicMock
-
-from infection_monkey.i_puppet.i_puppet import ExploiterResultData
-from infection_monkey.model.host import VictimHost
-from infection_monkey.telemetry.exploit_telem import ExploitTelem
-from infection_monkey.telemetry.messengers.exploit_intercepting_telemetry_messenger import (
-    ExploitInterceptingTelemetryMessenger,
-)
-
-
-class MockExploitTelem(ExploitTelem):
-    def __init__(self, propagation_success):
-        erd = ExploiterResultData()
-        erd.propagation_success = propagation_success
-        super().__init__("TestExploiter", VictimHost("127.0.0.1"), erd)
-
-    def get_data(self):
-        return {}
-
-
-def test_generic_telemetry(TestTelem):
-    mock_telemetry_messenger = MagicMock()
-    mock_relay = MagicMock()
-
-    telemetry_messenger = ExploitInterceptingTelemetryMessenger(
-        mock_telemetry_messenger, mock_relay
-    )
-
-    telemetry_messenger.send_telemetry(TestTelem())
-
-    assert mock_telemetry_messenger.send_telemetry.called
-    assert not mock_relay.add_potential_user.called
-
-
-def test_propagation_successful_exploit_telemetry():
-    mock_telemetry_messenger = MagicMock()
-    mock_relay = MagicMock()
-    mock_exploit_telem = MockExploitTelem(True)
-
-    telemetry_messenger = ExploitInterceptingTelemetryMessenger(
-        mock_telemetry_messenger, mock_relay
-    )
-
-    telemetry_messenger.send_telemetry(mock_exploit_telem)
-
-    assert mock_telemetry_messenger.send_telemetry.called
-    assert mock_relay.add_potential_user.called
-
-
-def test_propagation_failed_exploit_telemetry():
-    mock_telemetry_messenger = MagicMock()
-    mock_relay = MagicMock()
-    mock_exploit_telem = MockExploitTelem(False)
-
-    telemetry_messenger = ExploitInterceptingTelemetryMessenger(
-        mock_telemetry_messenger, mock_relay
-    )
-
-    telemetry_messenger.send_telemetry(mock_exploit_telem)
-
-    assert mock_telemetry_messenger.send_telemetry.called
-    assert not mock_relay.add_potential_user.called