diff --git a/monkey/infection_monkey/model/host.py b/monkey/infection_monkey/model/host.py index dcc6e7455..ebafb09f5 100644 --- a/monkey/infection_monkey/model/host.py +++ b/monkey/infection_monkey/model/host.py @@ -46,3 +46,17 @@ 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/network/network_scanner.py b/monkey/infection_monkey/network/network_scanner.py index a0b57c07b..fb2b5269e 100644 --- a/monkey/infection_monkey/network/network_scanner.py +++ b/monkey/infection_monkey/network/network_scanner.py @@ -3,8 +3,8 @@ 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.network.info import local_ips, get_interfaces_ranges -from infection_monkey.model import VictimHost from infection_monkey.network import TcpScanner, PingScanner from infection_monkey.utils import is_windows_os @@ -27,15 +27,13 @@ def generate_victims(net_ranges, chunk_size): """ chunk = [] for net_range in net_ranges: - for address in net_range: - if hasattr(net_range, 'domain_name'): - victim = VictimHost(address, net_range.domain_name) - else: - victim = VictimHost(address) + 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):