From 5a83401a0267703a793c7221e2a86249dc0950e7 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 7 Sep 2022 00:16:35 -0400 Subject: [PATCH] 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()))