Agent: Fix unnecessary waiting in MonkeyTunnel

The monkey tunnel only needs to wait before closing if propagation was
successful. Previously, it waited before closing if any exploiter was
run.

PR: #1811

monkey/infection_monkey/telemetry/messengers/exploit_intercepting_telemetry_messenger.py

@@ -1,7 +1,7 @@
 from functools import singledispatch
 
-from infection_monkey.telemetry.i_telem import ITelem
 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
 from infection_monkey.tunnel import MonkeyTunnel
 
@@ -26,5 +26,7 @@ def _send_telemetry(
 
 @_send_telemetry.register
 def _(telemetry: ExploitTelem, telemetry_messenger: ITelemetryMessenger, tunnel: MonkeyTunnel):
-    tunnel.set_wait_for_exploited_machines()
+    if telemetry.propagation_result is True:
+        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

@@ -1,8 +1,9 @@
 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.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,
 )
@@ -19,8 +20,10 @@ class TestTelem(BaseTelem):
 
 
 class MockExpliotTelem(ExploitTelem):
-    def __init__(self):
+    def __init__(self, propagation_success):
-        pass
+        erd = ExploiterResultData()
+        erd.propagation_success = propagation_success
+        super().__init__("TestExploiter", VictimHost("127.0.0.1"), erd)
 
     def get_data(self):
         return {}
@@ -40,10 +43,10 @@ def test_generic_telemetry():
     assert not mock_tunnel.set_wait_for_exploited_machines.called
 
 
-def test_expliot_telemetry():
+def test_propagation_successful_expliot_telemetry():
     mock_telemetry_messenger = MagicMock()
     mock_tunnel = MagicMock()
-    mock_expliot_telem = MockExpliotTelem()
+    mock_expliot_telem = MockExpliotTelem(True)
 
     telemetry_messenger = ExploitInterceptingTelemetryMessenger(
         mock_telemetry_messenger, mock_tunnel
@@ -53,3 +56,18 @@ def test_expliot_telemetry():
 
     assert mock_telemetry_messenger.send_telemetry.called
     assert mock_tunnel.set_wait_for_exploited_machines.called
+
+
+def test_propagation_failed_expliot_telemetry():
+    mock_telemetry_messenger = MagicMock()
+    mock_tunnel = MagicMock()
+    mock_expliot_telem = MockExpliotTelem(False)
+
+    telemetry_messenger = ExploitInterceptingTelemetryMessenger(
+        mock_telemetry_messenger, mock_tunnel
+    )
+
+    telemetry_messenger.send_telemetry(mock_expliot_telem)
+
+    assert mock_telemetry_messenger.send_telemetry.called
+    assert not mock_tunnel.set_wait_for_exploited_machines.called