diff --git a/monkey/infection_monkey/model/victim_host_generator.py b/monkey/infection_monkey/model/victim_host_generator.py deleted file mode 100644 index 444c4a5ee..000000000 --- a/monkey/infection_monkey/model/victim_host_generator.py +++ /dev/null @@ -1,45 +0,0 @@ -from infection_monkey.model.host import VictimHost - - -class VictimHostGenerator(object): - def __init__(self, network_ranges, blocked_ips, same_machine_ips): - self.blocked_ips = blocked_ips - self.ranges = network_ranges - self.local_addresses = same_machine_ips - - def generate_victims(self, chunk_size): - """ - Generates VictimHosts in chunks from all the instances network ranges - :param chunk_size: Maximum size of each chunk - """ - chunk = [] - for net_range in self.ranges: - for victim in self.generate_victims_from_range(net_range): - chunk.append(victim) - if len(chunk) == chunk_size: - yield chunk - chunk = [] - if chunk: # finished with number of victims < chunk_size - yield chunk - - def generate_victims_from_range(self, net_range): - """ - Generates VictimHosts from a given netrange - :param net_range: Network range object - :return: Generator of VictimHost objects - """ - for address in net_range: - if not self.is_ip_scannable(address): # check if the IP should be skipped - continue - if hasattr(net_range, "domain_name"): - victim = VictimHost(address, net_range.domain_name) - else: - victim = VictimHost(address) - yield victim - - def is_ip_scannable(self, ip_address): - if ip_address in self.local_addresses: - return False - if ip_address in self.blocked_ips: - return False - return True diff --git a/monkey/tests/unit_tests/infection_monkey/model/test_victim_host_generator.py b/monkey/tests/unit_tests/infection_monkey/model/test_victim_host_generator.py deleted file mode 100644 index 0133102eb..000000000 --- a/monkey/tests/unit_tests/infection_monkey/model/test_victim_host_generator.py +++ /dev/null @@ -1,41 +0,0 @@ -from unittest import TestCase - -from common.network.network_range import CidrRange, SingleIpRange -from infection_monkey.model.victim_host_generator import VictimHostGenerator - - -class TestVictimHostGenerator(TestCase): - def setUp(self): - self.cidr_range = CidrRange("10.0.0.0/28", False) # this gives us 15 hosts - self.local_host_range = SingleIpRange("localhost") - self.random_single_ip_range = SingleIpRange("41.50.13.37") - - def test_chunking(self): - chunk_size = 3 - # current test setup is 15+1+1-1 hosts - test_ranges = [self.cidr_range, self.local_host_range, self.random_single_ip_range] - generator = VictimHostGenerator(test_ranges, "10.0.0.1", []) - victims = generator.generate_victims(chunk_size) - for i in range(5): # quickly check the equally sided chunks - self.assertEqual(len(next(victims)), chunk_size) - victim_chunk_last = next(victims) - self.assertEqual(len(victim_chunk_last), 1) - - def test_remove_blocked_ip(self): - generator = VictimHostGenerator(self.cidr_range, ["10.0.0.1"], []) - - victims = list(generator.generate_victims_from_range(self.cidr_range)) - self.assertEqual(len(victims), 14) # 15 minus the 1 we blocked - - def test_remove_local_ips(self): - generator = VictimHostGenerator([], [], []) - generator.local_addresses = ["127.0.0.1"] - victims = list(generator.generate_victims_from_range(self.local_host_range)) - self.assertEqual(len(victims), 0) # block the local IP - - def test_generate_domain_victim(self): - # domain name victim - generator = VictimHostGenerator([], [], []) # dummy object - victims = list(generator.generate_victims_from_range(self.local_host_range)) - self.assertEqual(len(victims), 1) - self.assertEqual(victims[0].domain_name, "localhost")