Agent: Remove blocklisted IPs from scan targets

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

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