Agent: Add telemetry messenger to report exploited machines to tunnel

monkey/infection_monkey/telemetry/messengers/exploit_intercepting_telemetry_messenger.py

@@ -0,0 +1,30 @@
+from functools import singledispatch
+
+from infection_monkey.telemetry.i_telem import ITelem
+from infection_monkey.telemetry.exploit_telem import ExploitTelem
+from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger
+from infection_monkey.tunnel import MonkeyTunnel
+
+
+class ExploitInterceptingTelemetryMessenger(ITelemetryMessenger):
+    def __init__(self, telemetry_messenger: ITelemetryMessenger, tunnel: MonkeyTunnel):
+        self._telemetry_messenger = telemetry_messenger
+        self._tunnel = tunnel
+
+    def send_telemetry(self, telemetry: ITelem):
+        _send_telemetry(telemetry, self._telemetry_messenger, self._tunnel)
+
+
+# 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_messenger.send_telemetry(telemetry)
+
+
+@_send_telemetry.register
+def _(telemetry: ExploitTelem, telemetry_messenger: ITelemetryMessenger, tunnel: MonkeyTunnel):
+    tunnel.set_wait_for_exploited_machines()
+    telemetry_messenger.send_telemetry(telemetry)

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

@@ -0,0 +1,55 @@
+from unittest.mock import MagicMock
+
+from infection_monkey.telemetry.base_telem import BaseTelem
+from infection_monkey.telemetry.exploit_telem import ExploitTelem
+from infection_monkey.telemetry.i_telem import ITelem
+from infection_monkey.telemetry.messengers.exploit_intercepting_telemetry_messenger import (
+    ExploitInterceptingTelemetryMessenger,
+)
+
+
+class TestTelem(BaseTelem):
+    telem_category = None
+
+    def __init__(self):
+        pass
+
+    def get_data(self):
+        return {}
+
+
+class MockExpliotTelem(ExploitTelem):
+    def __init__(self):
+        pass
+
+    def get_data(self):
+        return {}
+
+
+def test_generic_telemetry():
+    mock_telemetry_messenger = MagicMock()
+    mock_tunnel = MagicMock()
+
+    telemetry_messenger = ExploitInterceptingTelemetryMessenger(
+        mock_telemetry_messenger, mock_tunnel
+    )
+
+    telemetry_messenger.send_telemetry(TestTelem())
+
+    assert mock_telemetry_messenger.send_telemetry.called
+    assert not mock_tunnel.set_wait_for_exploited_machines.called
+
+
+def test_expliot_telemetry():
+    mock_telemetry_messenger = MagicMock()
+    mock_tunnel = MagicMock()
+    mock_expliot_telem = MockExpliotTelem()
+
+    telemetry_messenger = ExploitInterceptingTelemetryMessenger(
+        mock_telemetry_messenger, mock_tunnel
+    )
+
+    telemetry_messenger.send_telemetry(mock_expliot_telem)
+
+    assert mock_telemetry_messenger.send_telemetry.called
+    assert mock_tunnel.set_wait_for_exploited_machines.called