From aba0446e61e17b1c52d81febd5dba938aab29f38 Mon Sep 17 00:00:00 2001
From: Mike Salvatore <mike.s.salvatore@gmail.com>
Date: Wed, 2 Mar 2022 09:43:14 -0500
Subject: [PATCH] Agent: Add telemetry messenger to report exploited machines
 to tunnel

---
 ...xploit_intercepting_telemetry_messenger.py | 30 ++++++++++
 ...xploit_intercepting_telemetry_messenger.py | 55 +++++++++++++++++++
 2 files changed, 85 insertions(+)
 create mode 100644 monkey/infection_monkey/telemetry/messengers/exploit_intercepting_telemetry_messenger.py
 create 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
new file mode 100644
index 000000000..3b92235fb
--- /dev/null
+++ b/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)
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
new file mode 100644
index 000000000..c6b85df3e
--- /dev/null
+++ b/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