From d8d55cb546d77239289db032365a7b400c70ec32 Mon Sep 17 00:00:00 2001
From: Daniel Goldberg <danielg@guardicore.com>
Date: Mon, 16 Sep 2019 16:32:13 +0300
Subject: [PATCH] PEP8 changes Documentation improvements

---
 monkey/infection_monkey/main.py               |  2 +-
 .../network/network_scanner.py                | 25 ++++++++++---------
 2 files changed, 14 insertions(+), 13 deletions(-)

diff --git a/monkey/infection_monkey/main.py b/monkey/infection_monkey/main.py
index ca5bb9832..9c9d24d73 100644
--- a/monkey/infection_monkey/main.py
+++ b/monkey/infection_monkey/main.py
@@ -44,7 +44,7 @@ def main():
 
     if 2 > len(sys.argv):
         return True
-    freeze_support() # required for multiprocessing + pyinstaller on windows
+    freeze_support()  # required for multiprocessing + pyinstaller on windows
     monkey_mode = sys.argv[1]
 
     if not (monkey_mode in [MONKEY_ARG, DROPPER_ARG]):
diff --git a/monkey/infection_monkey/network/network_scanner.py b/monkey/infection_monkey/network/network_scanner.py
index b751e99d7..a89920c72 100644
--- a/monkey/infection_monkey/network/network_scanner.py
+++ b/monkey/infection_monkey/network/network_scanner.py
@@ -1,17 +1,18 @@
 import sys
-import itertools
 import time
 import logging
 
-if sys.platform.startswith("win"):
-    from multiprocessing.dummy import Pool
-else:
-    from multiprocessing import Pool
 from common.network.network_range import NetworkRange
 from infection_monkey.config import WormConfiguration
 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
+
+if is_windows_os():
+    from multiprocessing.dummy import Pool
+else:
+    from multiprocessing import Pool
 
 __author__ = 'itamar'
 
@@ -26,7 +27,6 @@ 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
-    :return:
     """
     chunk = []
     for net_range in net_ranges:
@@ -51,7 +51,6 @@ class NetworkScanner(object):
         """
         Set up scanning.
         based on configuration: scans local network and/or scans fixed list of IPs/subnets.
-        :return:
         """
         # get local ip addresses
         self._ip_addresses = local_ips()
@@ -97,7 +96,12 @@ class NetworkScanner(object):
         :param stop_callback: A callback to check at any point if we should stop scanning
         :return: yields a sequence of VictimHost instances
         """
-        pool = Pool()
+        # We currently use the ITERATION_BLOCK_SIZE as the pool size, however, this may not be the best decision
+        # However, the decision what ITERATION_BLOCK_SIZE also requires balancing network usage (pps and bw)
+        # 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)
+
         victims_count = 0
         for victim_chunk in generate_victims(self._ranges, ITERATION_BLOCK_SIZE):
             LOG.debug("Scanning for potential victims in chunk %r", victim_chunk)
@@ -131,7 +135,6 @@ class NetworkScanner(object):
                 # time.sleep uses seconds, while config is in milliseconds
                 time.sleep(WormConfiguration.tcp_scan_interval / float(1000))
 
-
     @staticmethod
     def _is_any_ip_in_subnet(ip_addresses, subnet_str):
         for ip_address in ip_addresses:
@@ -139,10 +142,9 @@ class NetworkScanner(object):
                 return True
         return False
 
-
     def scan_machine(self, victim):
         """
-        Scans specific machine using given scanner
+        Scans specific machine using instance scanners
         :param victim: VictimHost machine
         :return: Victim or None if victim isn't alive
         """
@@ -153,6 +155,5 @@ class NetworkScanner(object):
         else:
             return None
 
-
     def on_island(self, server):
         return bool([x for x in self._ip_addresses if x in server])