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