From ed16826b879faef7af344b9d000d3a530acfefb9 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Thu, 16 Dec 2021 09:08:39 -0500 Subject: [PATCH] Agent: Sort scan targets by IP --- .../network/scan_target_generator.py | 3 ++- .../network/test_scan_target_generator.py | 15 +++++++++++++++ 2 files changed, 17 insertions(+), 1 deletion(-) diff --git a/monkey/infection_monkey/network/scan_target_generator.py b/monkey/infection_monkey/network/scan_target_generator.py index 734cc90c6..6cec82223 100644 --- a/monkey/infection_monkey/network/scan_target_generator.py +++ b/monkey/infection_monkey/network/scan_target_generator.py @@ -1,5 +1,6 @@ import itertools import logging +import socket from collections import namedtuple from typing import List @@ -32,7 +33,7 @@ def compile_scan_target_list( scan_targets = _remove_interface_ips(scan_targets, local_network_interfaces) scan_targets = _remove_blocklisted_ips(scan_targets, blocklisted_ips) scan_targets = _remove_redundant_targets(scan_targets) - scan_targets.sort() + scan_targets.sort(key=lambda network_address: socket.inet_aton(network_address.ip)) return scan_targets 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 4f3e49b64..03febe44c 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 @@ -466,3 +466,18 @@ def test_invalid_blocklisted_ip(): blocklisted_ips=blocklisted, enable_local_network_scan=False, ) + + +def test_sorted_scan_targets(): + expected_results = [f"10.1.0.{i}" for i in range(0, 255)] + expected_results.extend([f"10.2.0.{i}" for i in range(0, 255)]) + expected_results.extend([f"10.10.0.{i}" for i in range(0, 255)]) + expected_results.extend([f"10.20.0.{i}" for i in range(0, 255)]) + + scan_targets = compile_scan_target_list( + [], ["10.1.0.0/24", "10.10.0.0/24", "10.20.0.0/24", "10.2.0.0/24"], [], [], False + ) + + actual_results = [network_address.ip for network_address in scan_targets] + + assert expected_results == actual_results