From 8c55d2acd48dbe349c15221b83b3aceede1271b7 Mon Sep 17 00:00:00 2001 From: Daniel Goldberg Date: Mon, 23 Sep 2019 18:01:39 +0300 Subject: [PATCH] Refactor victim generation. Now we have a VictimHost generator that handles all the filtering. --- monkey/infection_monkey/model/host.py | 14 ------ .../model/victim_host_generator.py | 46 +++++++++++++++++++ .../network/network_scanner.py | 31 ++----------- 3 files changed, 49 insertions(+), 42 deletions(-) create mode 100644 monkey/infection_monkey/model/victim_host_generator.py diff --git a/monkey/infection_monkey/model/host.py b/monkey/infection_monkey/model/host.py index ebafb09f5..dcc6e7455 100644 --- a/monkey/infection_monkey/model/host.py +++ b/monkey/infection_monkey/model/host.py @@ -46,17 +46,3 @@ class VictimHost(object): def set_default_server(self, default_server): self.default_server = default_server - - -def generate_victims_from_range(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 hasattr(net_range, 'domain_name'): - victim = VictimHost(address, net_range.domain_name) - else: - victim = VictimHost(address) - yield victim diff --git a/monkey/infection_monkey/model/victim_host_generator.py b/monkey/infection_monkey/model/victim_host_generator.py new file mode 100644 index 000000000..1309278c8 --- /dev/null +++ b/monkey/infection_monkey/model/victim_host_generator.py @@ -0,0 +1,46 @@ +from infection_monkey.model.host import VictimHost +from infection_monkey.network.info import local_ips + + +class VictimHostGenerator(object): + def __init__(self, network_ranges, blocked_ips): + self._ip_addresses = local_ips() + self.blocked_ips = blocked_ips + self.ranges = network_ranges + + 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._ip_addresses: + return False + if ip_address in self.blocked_ips: + return False + return True diff --git a/monkey/infection_monkey/network/network_scanner.py b/monkey/infection_monkey/network/network_scanner.py index eb7c29572..da3d88e5f 100644 --- a/monkey/infection_monkey/network/network_scanner.py +++ b/monkey/infection_monkey/network/network_scanner.py @@ -3,7 +3,7 @@ import logging from common.network.network_range import NetworkRange from infection_monkey.config import WormConfiguration -from infection_monkey.model.host import generate_victims_from_range +from infection_monkey.model.victim_host_generator import VictimHostGenerator from infection_monkey.network.info import local_ips, get_interfaces_ranges from infection_monkey.network import TcpScanner, PingScanner from infection_monkey.utils import is_windows_os @@ -18,23 +18,6 @@ LOG = logging.getLogger(__name__) ITERATION_BLOCK_SIZE = 5 -def generate_victims(net_ranges, chunk_size): - """ - Generates VictimHosts in chunks from all the netranges - :param net_ranges: Iterable of network ranges - :param chunk_size: Maximum size of each chunk - """ - chunk = [] - for net_range in net_ranges: - for victim in 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 - - class NetworkScanner(object): def __init__(self): self._ip_addresses = None @@ -95,20 +78,12 @@ class NetworkScanner(object): # Because we are using this to spread out IO heavy tasks, we can probably go a lot higher than CPU core size # But again, balance pool = Pool(ITERATION_BLOCK_SIZE) + victim_generator = VictimHostGenerator(self._ranges, WormConfiguration.blocked_ips) victims_count = 0 - for victim_chunk in generate_victims(self._ranges, ITERATION_BLOCK_SIZE): + for victim_chunk in victim_generator.generate_victims(ITERATION_BLOCK_SIZE): LOG.debug("Scanning for potential victims in chunk %r", victim_chunk) - # skip self IP addresses - victim_chunk = [x for x in victim_chunk if x.ip_addr not in self._ip_addresses] - # skip IPs marked as blocked - - bad_victims = [x for x in victim_chunk if x.ip_addr in WormConfiguration.blocked_ips] - for victim in bad_victims: - LOG.info("Skipping %s due to blacklist" % victim) - victim_chunk = [x for x in victim_chunk if x.ip_addr not in WormConfiguration.blocked_ips] - # check before running scans if stop_callback and stop_callback(): LOG.debug("Got stop signal")