From 90dcb0a91e1875fc0829ad8c982a1438e738ffe3 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Wed, 31 Aug 2022 14:38:35 +0000 Subject: [PATCH 1/3] Agent: Wait for relay users to disconnect --- .../network/relay/relay_user_handler.py | 19 ++++++++++++++++++- .../network/relay/tcp_relay.py | 2 +- .../network/relay/test_relay_user_handler.py | 15 +++++++++++++++ 3 files changed, 34 insertions(+), 2 deletions(-) diff --git a/monkey/infection_monkey/network/relay/relay_user_handler.py b/monkey/infection_monkey/network/relay/relay_user_handler.py index e9ee00d04..2c83f0b13 100644 --- a/monkey/infection_monkey/network/relay/relay_user_handler.py +++ b/monkey/infection_monkey/network/relay/relay_user_handler.py @@ -10,6 +10,7 @@ from common.utils.code_utils import del_key # Wait for potential new clients to connect DEFAULT_NEW_CLIENT_TIMEOUT = 2.5 * MEDIUM_REQUEST_TIMEOUT +DEFAULT_DISCONNECT_TIMEOUT = 60 * 10 # Wait up to 10 minutes for clients to disconnect @dataclass @@ -21,8 +22,13 @@ class RelayUser: class RelayUserHandler: """Manages membership to a network relay.""" - def __init__(self, new_client_timeout: float = DEFAULT_NEW_CLIENT_TIMEOUT): + def __init__( + self, + new_client_timeout: float = DEFAULT_NEW_CLIENT_TIMEOUT, + client_disconnect_timeout: float = DEFAULT_DISCONNECT_TIMEOUT, + ): self._new_client_timeout = new_client_timeout + self._client_disconnect_timeout = client_disconnect_timeout self._relay_users: Dict[IPv4Address, RelayUser] = {} self._potential_users: Dict[IPv4Address, RelayUser] = {} @@ -41,6 +47,7 @@ class RelayUserHandler: del_key(self._potential_users, user_address) timer = EggTimer() + timer.set(self._client_disconnect_timeout) self._relay_users[user_address] = RelayUser(user_address, timer) def add_potential_user(self, user_address: IPv4Address): @@ -74,3 +81,13 @@ class RelayUserHandler: ) return len(self._potential_users) > 0 + + def has_connected_users(self) -> bool: + """ + Return whether or not we have any relay users. + """ + self._relay_users = dict( + filter(lambda ru: not ru[1].timer.is_expired(), self._relay_users.items()) + ) + + return len(self._relay_users) > 0 diff --git a/monkey/infection_monkey/network/relay/tcp_relay.py b/monkey/infection_monkey/network/relay/tcp_relay.py index a12d8a0d9..f605cdc9d 100644 --- a/monkey/infection_monkey/network/relay/tcp_relay.py +++ b/monkey/infection_monkey/network/relay/tcp_relay.py @@ -36,7 +36,7 @@ class TCPRelay(Thread, InterruptableThreadMixin): """ Blocks until the users disconnect or the timeout has elapsed. """ - while self._user_handler.has_potential_users(): + while self._user_handler.has_potential_users() or self._user_handler.has_connected_users(): sleep(0.5) def _wait_for_pipes_to_close(self): diff --git a/monkey/tests/unit_tests/infection_monkey/network/relay/test_relay_user_handler.py b/monkey/tests/unit_tests/infection_monkey/network/relay/test_relay_user_handler.py index ca0eb4103..6f3ecb8fa 100644 --- a/monkey/tests/unit_tests/infection_monkey/network/relay/test_relay_user_handler.py +++ b/monkey/tests/unit_tests/infection_monkey/network/relay/test_relay_user_handler.py @@ -33,3 +33,18 @@ def test_potential_users_time_out(): sleep(0.003) assert not handler.has_potential_users() + + +def test_relay_users_added(handler): + assert not handler.has_connected_users() + handler.add_relay_user(USER_ADDRESS) + assert handler.has_connected_users() + + +def test_relay_users_time_out(): + handler = RelayUserHandler(client_disconnect_timeout=0.001) + + handler.add_relay_user(USER_ADDRESS) + sleep(0.003) + + assert not handler.has_connected_users() From e2736a9273ac04161dee3e1cb4d35349b9408430 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 7 Sep 2022 00:02:23 -0400 Subject: [PATCH 2/3] Agent: Lock RelayUserHandler.has_*_users() --- .../network/relay/relay_user_handler.py | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/monkey/infection_monkey/network/relay/relay_user_handler.py b/monkey/infection_monkey/network/relay/relay_user_handler.py index 2c83f0b13..741045a29 100644 --- a/monkey/infection_monkey/network/relay/relay_user_handler.py +++ b/monkey/infection_monkey/network/relay/relay_user_handler.py @@ -76,18 +76,20 @@ class RelayUserHandler: """ Return whether or not we have any potential users. """ - self._potential_users = dict( - filter(lambda ru: not ru[1].timer.is_expired(), self._potential_users.items()) - ) + with self._lock: + self._potential_users = dict( + filter(lambda ru: not ru[1].timer.is_expired(), self._potential_users.items()) + ) - return len(self._potential_users) > 0 + return len(self._potential_users) > 0 def has_connected_users(self) -> bool: """ Return whether or not we have any relay users. """ - self._relay_users = dict( - filter(lambda ru: not ru[1].timer.is_expired(), self._relay_users.items()) - ) + with self._lock: + self._relay_users = dict( + filter(lambda ru: not ru[1].timer.is_expired(), self._relay_users.items()) + ) - return len(self._relay_users) > 0 + return len(self._relay_users) > 0 From 5a83401a0267703a793c7221e2a86249dc0950e7 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 7 Sep 2022 00:16:35 -0400 Subject: [PATCH 3/3] Agent: Add RelayUserHandler._remove_expired_users() --- .../network/relay/relay_user_handler.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/monkey/infection_monkey/network/relay/relay_user_handler.py b/monkey/infection_monkey/network/relay/relay_user_handler.py index 741045a29..5755d2794 100644 --- a/monkey/infection_monkey/network/relay/relay_user_handler.py +++ b/monkey/infection_monkey/network/relay/relay_user_handler.py @@ -77,9 +77,7 @@ class RelayUserHandler: Return whether or not we have any potential users. """ with self._lock: - self._potential_users = dict( - filter(lambda ru: not ru[1].timer.is_expired(), self._potential_users.items()) - ) + self._potential_users = RelayUserHandler._remove_expired_users(self._potential_users) return len(self._potential_users) > 0 @@ -88,8 +86,12 @@ class RelayUserHandler: Return whether or not we have any relay users. """ with self._lock: - self._relay_users = dict( - filter(lambda ru: not ru[1].timer.is_expired(), self._relay_users.items()) - ) + self._relay_users = RelayUserHandler._remove_expired_users(self._relay_users) return len(self._relay_users) > 0 + + @staticmethod + def _remove_expired_users( + user_list: Dict[IPv4Address, RelayUser] + ) -> Dict[IPv4Address, RelayUser]: + return dict(filter(lambda ru: not ru[1].timer.is_expired(), user_list.items()))