From cd0b3077cfa339d76916e94207b2772ead1720e7 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Tue, 30 Aug 2022 21:31:06 +0000 Subject: [PATCH] Agent: Notify TCPRelay of exploit --- monkey/infection_monkey/tcp_relay.py | 9 +++++++ ...xploit_intercepting_telemetry_messenger.py | 22 +++++++++++++--- ...xploit_intercepting_telemetry_messenger.py | 26 ++++++++++++------- 3 files changed, 43 insertions(+), 14 deletions(-) diff --git a/monkey/infection_monkey/tcp_relay.py b/monkey/infection_monkey/tcp_relay.py index ed8340dd7..eaa564c3b 100644 --- a/monkey/infection_monkey/tcp_relay.py +++ b/monkey/infection_monkey/tcp_relay.py @@ -22,6 +22,7 @@ class TCPRelay(Thread): super(TCPRelay, self).__init__(name="MonkeyTcpRelayThread") self.daemon = True self._relay_users: List[RelayUser] = [] + self._potential_users: List[RelayUser] = [] self._lock = Lock() def run(self): @@ -44,13 +45,21 @@ class TCPRelay(Thread): self._stopped.set() def on_user_connected(self, user: str): + """Handle new user connection.""" with self._lock: self._relay_users.append(RelayUser(user)) def on_user_disconnected(self, user: str): + """Handle user disconnection.""" with self._lock: self._relay_users = [u for u in self._relay_users if u.address != user] def relay_users(self) -> List[RelayUser]: + """Get the list of users connected to the relay.""" with self._lock: return self._relay_users.copy() + + def on_potential_new_user(self, user: str): + """Notify TCPRelay that a new user may try and connect.""" + with self._lock: + self._potential_users.append(RelayUser(user)) diff --git a/monkey/infection_monkey/telemetry/messengers/exploit_intercepting_telemetry_messenger.py b/monkey/infection_monkey/telemetry/messengers/exploit_intercepting_telemetry_messenger.py index b2a254061..5eb8f0310 100644 --- a/monkey/infection_monkey/telemetry/messengers/exploit_intercepting_telemetry_messenger.py +++ b/monkey/infection_monkey/telemetry/messengers/exploit_intercepting_telemetry_messenger.py @@ -1,5 +1,6 @@ from functools import singledispatch +from infection_monkey.tcp_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 @@ -7,26 +8,39 @@ from infection_monkey.tunnel import MonkeyTunnel class ExploitInterceptingTelemetryMessenger(ITelemetryMessenger): - def __init__(self, telemetry_messenger: ITelemetryMessenger, tunnel: MonkeyTunnel): + def __init__( + self, telemetry_messenger: ITelemetryMessenger, tunnel: MonkeyTunnel, relay: TCPRelay = None + ): self._telemetry_messenger = telemetry_messenger self._tunnel = tunnel + self._relay = relay def send_telemetry(self, telemetry: ITelem): - _send_telemetry(telemetry, self._telemetry_messenger, self._tunnel) + _send_telemetry(telemetry, self._telemetry_messenger, self._tunnel, 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, tunnel: MonkeyTunnel + telemetry: ITelem, + telemetry_messenger: ITelemetryMessenger, + tunnel: MonkeyTunnel, + relay: TCPRelay = None, ): telemetry_messenger.send_telemetry(telemetry) @_send_telemetry.register -def _(telemetry: ExploitTelem, telemetry_messenger: ITelemetryMessenger, tunnel: MonkeyTunnel): +def _( + telemetry: ExploitTelem, + telemetry_messenger: ITelemetryMessenger, + tunnel: MonkeyTunnel, + relay: TCPRelay = None, +): if telemetry.propagation_result is True: tunnel.set_wait_for_exploited_machines() + if relay: + relay.on_potential_new_user(str(telemetry.host["ip_addr"])) 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 index 969489107..7e9e1e7fd 100644 --- 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 @@ -8,7 +8,7 @@ from infection_monkey.telemetry.messengers.exploit_intercepting_telemetry_messen ) -class MockExpliotTelem(ExploitTelem): +class MockExploitTelem(ExploitTelem): def __init__(self, propagation_success): erd = ExploiterResultData() erd.propagation_success = propagation_success @@ -21,42 +21,48 @@ class MockExpliotTelem(ExploitTelem): def test_generic_telemetry(TestTelem): mock_telemetry_messenger = MagicMock() mock_tunnel = MagicMock() + mock_relay = MagicMock() telemetry_messenger = ExploitInterceptingTelemetryMessenger( - mock_telemetry_messenger, mock_tunnel + mock_telemetry_messenger, mock_tunnel, mock_relay ) telemetry_messenger.send_telemetry(TestTelem()) assert mock_telemetry_messenger.send_telemetry.called assert not mock_tunnel.set_wait_for_exploited_machines.called + assert not mock_relay.on_potential_new_user.called -def test_propagation_successful_expliot_telemetry(): +def test_propagation_successful_exploit_telemetry(): mock_telemetry_messenger = MagicMock() mock_tunnel = MagicMock() - mock_expliot_telem = MockExpliotTelem(True) + mock_relay = MagicMock() + mock_exploit_telem = MockExploitTelem(True) telemetry_messenger = ExploitInterceptingTelemetryMessenger( - mock_telemetry_messenger, mock_tunnel + mock_telemetry_messenger, mock_tunnel, mock_relay ) - telemetry_messenger.send_telemetry(mock_expliot_telem) + telemetry_messenger.send_telemetry(mock_exploit_telem) assert mock_telemetry_messenger.send_telemetry.called assert mock_tunnel.set_wait_for_exploited_machines.called + assert mock_relay.on_potential_new_user.called -def test_propagation_failed_expliot_telemetry(): +def test_propagation_failed_exploit_telemetry(): mock_telemetry_messenger = MagicMock() mock_tunnel = MagicMock() - mock_expliot_telem = MockExpliotTelem(False) + mock_relay = MagicMock() + mock_exploit_telem = MockExploitTelem(False) telemetry_messenger = ExploitInterceptingTelemetryMessenger( - mock_telemetry_messenger, mock_tunnel + mock_telemetry_messenger, mock_tunnel, mock_relay ) - telemetry_messenger.send_telemetry(mock_expliot_telem) + telemetry_messenger.send_telemetry(mock_exploit_telem) assert mock_telemetry_messenger.send_telemetry.called assert not mock_tunnel.set_wait_for_exploited_machines.called + assert not mock_relay.on_potential_new_user.called