From 7c786b08831a9ccac6e24b7414d84241d66a5e36 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Thu, 16 Dec 2021 08:29:45 -0500 Subject: [PATCH] Agent: Improve performance of _remove_ips_from_scan_targets() --- monkey/infection_monkey/network/scan_target_generator.py | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/monkey/infection_monkey/network/scan_target_generator.py b/monkey/infection_monkey/network/scan_target_generator.py index aaee67a0d..d8f3339d7 100644 --- a/monkey/infection_monkey/network/scan_target_generator.py +++ b/monkey/infection_monkey/network/scan_target_generator.py @@ -99,13 +99,8 @@ def _remove_blocklisted_ips( def _remove_ips_from_scan_targets( scan_targets: List[NetworkAddress], ips_to_remove: List[str] ) -> List[NetworkAddress]: - for ip in ips_to_remove: - try: - scan_targets = [address for address in scan_targets if address.ip != ip] - except KeyError: - # We don't need to remove the ip if it's already missing from the scan_targets - pass - return scan_targets + ips_to_remove_set = set(ips_to_remove) + return [address for address in scan_targets if address.ip not in ips_to_remove_set] def _get_segmentation_check_targets(