From a32a78325743b6d4d63d7d17b38d801efc711298 Mon Sep 17 00:00:00 2001
From: Daniel Goldberg <danielg@guardicore.com>
Date: Fri, 20 Sep 2019 22:58:20 +0300
Subject: [PATCH] Factor out generating VictimHosts from NetworkRange object

---
 monkey/infection_monkey/model/host.py              | 14 ++++++++++++++
 monkey/infection_monkey/network/network_scanner.py | 10 ++++------
 2 files changed, 18 insertions(+), 6 deletions(-)

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):