From 0be709a9eb8b9999247cc4e4c2a03d0372e2fae7 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Tue, 5 Apr 2022 13:03:55 +0530 Subject: [PATCH 1/5] Agent: Use Timer in BatchingTelemetryMessenger --- .../batching_telemetry_messenger.py | 20 +++++++------------ 1 file changed, 7 insertions(+), 13 deletions(-) diff --git a/monkey/infection_monkey/telemetry/messengers/batching_telemetry_messenger.py b/monkey/infection_monkey/telemetry/messengers/batching_telemetry_messenger.py index 4d34012d8..9dc051666 100644 --- a/monkey/infection_monkey/telemetry/messengers/batching_telemetry_messenger.py +++ b/monkey/infection_monkey/telemetry/messengers/batching_telemetry_messenger.py @@ -1,11 +1,11 @@ import queue import threading -import time from typing import Dict from infection_monkey.telemetry.i_batchable_telem import IBatchableTelem from infection_monkey.telemetry.i_telem import ITelem from infection_monkey.telemetry.messengers.i_telemetry_messenger import ITelemetryMessenger +from infection_monkey.utils.timer import Timer DEFAULT_PERIOD = 5 WAKES_PER_PERIOD = 4 @@ -40,8 +40,6 @@ class BatchingTelemetryMessenger(ITelemetryMessenger): self._period = period self._should_run_batch_thread = True - # TODO: Replace with infection_monkey.utils.timer.Timer - self._last_sent_time = time.time() self._telemetry_batches: Dict[str, IBatchableTelem] = {} self._manage_telemetry_batches_thread = None @@ -59,21 +57,20 @@ class BatchingTelemetryMessenger(ITelemetryMessenger): self._manage_telemetry_batches_thread = None def _manage_telemetry_batches(self): - self._reset() + timer = Timer() + timer.set(self._period) + self._telemetry_batches = {} while self._should_run_batch_thread: self._process_next_telemetry() - if self._period_elapsed(): + if timer.is_expired(): self._send_telemetry_batches() - self._reset() + timer.reset() + self._telemetry_batches = {} self._send_remaining_telemetry_batches() - def _reset(self): - self._last_sent_time = time.time() - self._telemetry_batches = {} - def _process_next_telemetry(self): try: telemetry = self._queue.get(block=True, timeout=self._period / WAKES_PER_PERIOD) @@ -93,9 +90,6 @@ class BatchingTelemetryMessenger(ITelemetryMessenger): else: self._telemetry_batches[telem_category] = new_telemetry - def _period_elapsed(self): - return (time.time() - self._last_sent_time) > self._period - def _send_remaining_telemetry_batches(self): while not self._queue.empty(): self._process_next_telemetry() From 60ca543117a06a59305ddf22571d316f6a3d7805 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Tue, 5 Apr 2022 18:55:08 +0530 Subject: [PATCH 2/5] Agent: Use Timer in MonkeyTunnel --- monkey/infection_monkey/tunnel.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/monkey/infection_monkey/tunnel.py b/monkey/infection_monkey/tunnel.py index b0f778534..55349dfd6 100644 --- a/monkey/infection_monkey/tunnel.py +++ b/monkey/infection_monkey/tunnel.py @@ -8,6 +8,7 @@ from infection_monkey.network.firewall import app as firewall from infection_monkey.network.info import get_free_tcp_port, local_ips from infection_monkey.network.tools import check_tcp_port, get_interface_to_target from infection_monkey.transport.base import get_last_serve_time +from infection_monkey.utils.timer import Timer logger = logging.getLogger(__name__) @@ -181,8 +182,9 @@ class MonkeyTunnel(Thread): # wait till all of the tunnel clients has been disconnected, or no one used the tunnel in # QUIT_TIMEOUT seconds - # TODO: Replace with infection_monkey.utils.timer.Timer - while self._clients and (time.time() - get_last_serve_time() < QUIT_TIMEOUT): + timer = Timer() + timer.set(self._calculate_timeout()) + while self._clients and not timer.is_expired(): try: search, address = self._broad_sock.recvfrom(BUFFER_READ) if b"-" == search: @@ -191,11 +193,16 @@ class MonkeyTunnel(Thread): except socket.timeout: continue + timer.set(self._calculate_timeout()) + logger.info("Closing tunnel") self._broad_sock.close() proxy.stop() proxy.join() + def _calculate_timeout(self): + return QUIT_TIMEOUT - (time.time() - get_last_serve_time()) + def get_tunnel_for_ip(self, ip: str): if not self.local_port: From 8a393e0b0391e2e2fbab49cefc177824432379c1 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Tue, 5 Apr 2022 19:07:38 +0530 Subject: [PATCH 3/5] Agent: Update comment in Timer --- monkey/infection_monkey/utils/timer.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/monkey/infection_monkey/utils/timer.py b/monkey/infection_monkey/utils/timer.py index 2ed17d551..6095d466e 100644 --- a/monkey/infection_monkey/utils/timer.py +++ b/monkey/infection_monkey/utils/timer.py @@ -13,7 +13,8 @@ class Timer: def set(self, timeout_sec: float): """ Set a timer - :param float timeout_sec: A fractional number of seconds to set the timeout for. + :param float timeout_sec: A nonnegative floating point number expressing the number of + seconds to set the timeout for. """ self._timeout_sec = timeout_sec self._start_time = time.time() From c9b9a27c34d4ea3f68b97fdee524c5d6f63b3429 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Wed, 6 Apr 2022 12:43:48 +0530 Subject: [PATCH 4/5] Agent: Fix logic (indenting) in MonkeyTunnel --- monkey/infection_monkey/tunnel.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/monkey/infection_monkey/tunnel.py b/monkey/infection_monkey/tunnel.py index 55349dfd6..64b204b15 100644 --- a/monkey/infection_monkey/tunnel.py +++ b/monkey/infection_monkey/tunnel.py @@ -193,7 +193,7 @@ class MonkeyTunnel(Thread): except socket.timeout: continue - timer.set(self._calculate_timeout()) + timer.set(self._calculate_timeout()) logger.info("Closing tunnel") self._broad_sock.close() From 2e087cf5214f93685e6ec1e8b5ecff9ff6b9acda Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Thu, 7 Apr 2022 19:14:24 +0530 Subject: [PATCH 5/5] Agent: Add try/except block in _calculate_timeout in MonkeyTunnel --- monkey/infection_monkey/tunnel.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/monkey/infection_monkey/tunnel.py b/monkey/infection_monkey/tunnel.py index 64b204b15..26368bff6 100644 --- a/monkey/infection_monkey/tunnel.py +++ b/monkey/infection_monkey/tunnel.py @@ -200,8 +200,11 @@ class MonkeyTunnel(Thread): proxy.stop() proxy.join() - def _calculate_timeout(self): - return QUIT_TIMEOUT - (time.time() - get_last_serve_time()) + def _calculate_timeout(self) -> float: + try: + return QUIT_TIMEOUT - (time.time() - get_last_serve_time()) + except TypeError: # get_last_serve_time() may return None + return 0.0 def get_tunnel_for_ip(self, ip: str):