From 0466eb7239ac60e418b5701b8296e97d914dfc3f Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 3 Oct 2022 13:02:00 -0400 Subject: [PATCH 1/6] Agent: Add notify_relay_on_propagation agent event handler --- .../agent_event_handlers/__init__.py | 1 + .../notify_relay_on_propagation.py | 19 +++++++ .../agent_event_handlers/__init__.py | 0 .../test_notify_relay_on_propagation.py | 54 +++++++++++++++++++ 4 files changed, 74 insertions(+) create mode 100644 monkey/infection_monkey/agent_event_handlers/__init__.py create mode 100644 monkey/infection_monkey/agent_event_handlers/notify_relay_on_propagation.py create mode 100644 monkey/tests/unit_tests/infection_monkey/agent_event_handlers/__init__.py create mode 100644 monkey/tests/unit_tests/infection_monkey/agent_event_handlers/test_notify_relay_on_propagation.py diff --git a/monkey/infection_monkey/agent_event_handlers/__init__.py b/monkey/infection_monkey/agent_event_handlers/__init__.py new file mode 100644 index 000000000..3e0f9b03c --- /dev/null +++ b/monkey/infection_monkey/agent_event_handlers/__init__.py @@ -0,0 +1 @@ +from .notify_relay_on_propagation import notify_relay_on_propagation diff --git a/monkey/infection_monkey/agent_event_handlers/notify_relay_on_propagation.py b/monkey/infection_monkey/agent_event_handlers/notify_relay_on_propagation.py new file mode 100644 index 000000000..7cd85ec4a --- /dev/null +++ b/monkey/infection_monkey/agent_event_handlers/notify_relay_on_propagation.py @@ -0,0 +1,19 @@ +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: + def __init__(self, tcp_relay: Optional[TCPRelay]): + self._tcp_relay = tcp_relay + + def __call__(self, event: PropagationEvent): + if self._tcp_relay is None: + return + + if event.success: + self._tcp_relay.add_potential_user(event.target) diff --git a/monkey/tests/unit_tests/infection_monkey/agent_event_handlers/__init__.py b/monkey/tests/unit_tests/infection_monkey/agent_event_handlers/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/monkey/tests/unit_tests/infection_monkey/agent_event_handlers/test_notify_relay_on_propagation.py b/monkey/tests/unit_tests/infection_monkey/agent_event_handlers/test_notify_relay_on_propagation.py new file mode 100644 index 000000000..be03bada9 --- /dev/null +++ b/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_relay_not_notified_if_none(): + handler = notify_relay_on_propagation(None) + + # Raises AttributeError on failure + handler(SUCCESSFUL_PROPAGATION_EVENT) From fb7d62e31868b74942479635a1bb8a6f65a6c8d9 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 3 Oct 2022 13:05:51 -0400 Subject: [PATCH 2/6] Agent: Subscribe notify_relay_on_propagation to PropagationEvent events --- monkey/infection_monkey/monkey.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 33ae633d3..b298179c0 100644 --- a/monkey/infection_monkey/monkey.py +++ b/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 ( @@ -313,6 +314,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, From 2ad972548b0f27b108a0cfa7428cfb711c18d8a0 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 3 Oct 2022 13:07:22 -0400 Subject: [PATCH 3/6] Agent: Remove ExploitInterceptingTelemetryMessenger decoration --- monkey/infection_monkey/monkey.py | 9 +-------- 1 file changed, 1 insertion(+), 8 deletions(-) diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index b298179c0..e38cbc96e 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -81,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, ) @@ -279,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, From d3ff56138fe57ba40260ba5e79fc76fcdc991a32 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 3 Oct 2022 13:08:52 -0400 Subject: [PATCH 4/6] Agent: Remove disused ExploitInterceptingTelemetryMessenger --- ...xploit_intercepting_telemetry_messenger.py | 41 ------------ ...xploit_intercepting_telemetry_messenger.py | 62 ------------------- 2 files changed, 103 deletions(-) delete mode 100644 monkey/infection_monkey/telemetry/messengers/exploit_intercepting_telemetry_messenger.py delete mode 100644 monkey/tests/unit_tests/infection_monkey/telemetry/messengers/test_exploit_intercepting_telemetry_messenger.py diff --git a/monkey/infection_monkey/telemetry/messengers/exploit_intercepting_telemetry_messenger.py b/monkey/infection_monkey/telemetry/messengers/exploit_intercepting_telemetry_messenger.py deleted file mode 100644 index cb59aa0ae..000000000 --- a/monkey/infection_monkey/telemetry/messengers/exploit_intercepting_telemetry_messenger.py +++ /dev/null @@ -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) diff --git a/monkey/tests/unit_tests/infection_monkey/telemetry/messengers/test_exploit_intercepting_telemetry_messenger.py b/monkey/tests/unit_tests/infection_monkey/telemetry/messengers/test_exploit_intercepting_telemetry_messenger.py deleted file mode 100644 index 61ca1b971..000000000 --- a/monkey/tests/unit_tests/infection_monkey/telemetry/messengers/test_exploit_intercepting_telemetry_messenger.py +++ /dev/null @@ -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 From a8383f4a79167442f6200b09e6813231a0f93076 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 3 Oct 2022 13:24:09 -0400 Subject: [PATCH 5/6] Agent: Add docstrings to notify_relay_on_propagation --- .../notify_relay_on_propagation.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/monkey/infection_monkey/agent_event_handlers/notify_relay_on_propagation.py b/monkey/infection_monkey/agent_event_handlers/notify_relay_on_propagation.py index 7cd85ec4a..0d1991931 100644 --- a/monkey/infection_monkey/agent_event_handlers/notify_relay_on_propagation.py +++ b/monkey/infection_monkey/agent_event_handlers/notify_relay_on_propagation.py @@ -8,10 +8,22 @@ 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 From 399fedfba5ebaead9b1b80da4ad44d43a6d5116e Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 3 Oct 2022 14:31:27 -0400 Subject: [PATCH 6/6] UT: Rename test_relay_not_notified_if_none --- .../agent_event_handlers/test_notify_relay_on_propagation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/monkey/tests/unit_tests/infection_monkey/agent_event_handlers/test_notify_relay_on_propagation.py b/monkey/tests/unit_tests/infection_monkey/agent_event_handlers/test_notify_relay_on_propagation.py index be03bada9..74b896bb7 100644 --- a/monkey/tests/unit_tests/infection_monkey/agent_event_handlers/test_notify_relay_on_propagation.py +++ b/monkey/tests/unit_tests/infection_monkey/agent_event_handlers/test_notify_relay_on_propagation.py @@ -47,7 +47,7 @@ def test_relay_not_notified_on_successful_propation(mock_tcp_relay): mock_tcp_relay.add_potential_user.assert_not_called() -def test_relay_not_notified_if_none(): +def test_handler_doesnt_raise_if_relay_is_none(): handler = notify_relay_on_propagation(None) # Raises AttributeError on failure