diff --git a/monkey/infection_monkey/network/scan_target_generator.py b/monkey/infection_monkey/network/scan_target_generator.py index 41b570622..cdcfbdb31 100644 --- a/monkey/infection_monkey/network/scan_target_generator.py +++ b/monkey/infection_monkey/network/scan_target_generator.py @@ -12,8 +12,11 @@ def compile_scan_target_list( ) -> List[str]: scan_targets = _get_ips_from_ranges_to_scan(ranges_to_scan) + _remove_blocklisted_ips(scan_targets, blocklisted_ips) + scan_target_list = list(scan_targets) scan_target_list.sort() + return scan_target_list @@ -25,3 +28,12 @@ def _get_ips_from_ranges_to_scan(ranges_to_scan: List[str]) -> Set[str]: scan_targets.update(set(_range)) return scan_targets + + +def _remove_blocklisted_ips(scan_targets: Set[str], blocked_ips: List[str]): + for blocked_ip in blocked_ips: + try: + scan_targets.remove(blocked_ip) + except KeyError: + # We don't need to remove the blocked ip if it's already missing from the scan_targets + pass diff --git a/monkey/tests/unit_tests/infection_monkey/network/test_scan_target_generator.py b/monkey/tests/unit_tests/infection_monkey/network/test_scan_target_generator.py index 89a0775d8..9e1b5fc0b 100644 --- a/monkey/tests/unit_tests/infection_monkey/network/test_scan_target_generator.py +++ b/monkey/tests/unit_tests/infection_monkey/network/test_scan_target_generator.py @@ -63,3 +63,43 @@ def test_ip_range(ip_range): for i in range(25, 34): assert f"192.168.56.{i}" in scan_targets + + +def test_no_duplicates(): + scan_targets = compile_ranges_only(["192.168.56.0/29", "192.168.56.2", "192.168.56.4"]) + + assert len(scan_targets) == 7 + + for i in range(0, 7): + assert f"192.168.56.{i}" in scan_targets + + +def test_blocklisted_ips(): + blocklisted_ips = ["10.0.0.5", "10.0.0.32", "10.0.0.119", "192.168.1.33"] + + scan_targets = compile_scan_target_list( + local_ips=[], + ranges_to_scan=["10.0.0.0/24"], + inaccessible_subnets=[], + blocklisted_ips=blocklisted_ips, + enable_local_network_scan=False, + ) + + assert len(scan_targets) == 252 + for blocked_ip in blocklisted_ips: + assert blocked_ip not in scan_targets + + +@pytest.mark.parametrize("ranges_to_scan", [["10.0.0.5"], []]) +def test_only_ip_blocklisted(ranges_to_scan): + blocklisted_ips = ["10.0.0.5"] + + scan_targets = compile_scan_target_list( + local_ips=[], + ranges_to_scan=ranges_to_scan, + inaccessible_subnets=[], + blocklisted_ips=blocklisted_ips, + enable_local_network_scan=False, + ) + + assert len(scan_targets) == 0