From 996f2b3c7a369e465efffb11f41569ad877760fe Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Thu, 24 Mar 2022 10:38:41 -0400 Subject: [PATCH] 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 --- ...xploit_intercepting_telemetry_messenger.py | 6 ++-- ...xploit_intercepting_telemetry_messenger.py | 28 +++++++++++++++---- 2 files changed, 27 insertions(+), 7 deletions(-) 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 3b92235fb..b2a254061 100644 --- a/monkey/infection_monkey/telemetry/messengers/exploit_intercepting_telemetry_messenger.py +++ b/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) 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 c6b85df3e..f949738f6 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 @@ -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): - pass + def __init__(self, propagation_success): + 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