Agent: Notify TCPRelay of exploit

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))

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)

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