Merge pull request #437 from guardicore/feature/scan_hosts_fast

Feature/scan hosts fast 
Yay, done with my longest waiting branch. 
Next up, OS sniffing.

monkey/infection_monkey/config.py

@@ -141,10 +141,10 @@ class Configuration(object):
     exploiter_classes = []
 
     # how many victims to look for in a single scan iteration
-    victims_max_find = 30
+    victims_max_find = 100
 
     # how many victims to exploit before stopping
-    victims_max_exploit = 7
+    victims_max_exploit = 15
 
     # depth of propagation
     depth = 2
@@ -199,7 +199,7 @@ class Configuration(object):
                         9200]
     tcp_target_ports.extend(HTTP_PORTS)
     tcp_scan_timeout = 3000  # 3000 Milliseconds
-    tcp_scan_interval = 0
+    tcp_scan_interval = 0  # in milliseconds
     tcp_scan_get_banner = True
 
     # Ping Scanner

monkey/infection_monkey/example.conf

@@ -97,8 +97,8 @@
     ],
     "timeout_between_iterations": 10,
     "use_file_logging": true,
-    "victims_max_exploit": 7,
+    "victims_max_exploit": 15,
-    "victims_max_find": 30,
+    "victims_max_find": 100,
     "post_breach_actions" : []
     custom_PBA_linux_cmd = ""
     custom_PBA_windows_cmd = ""

monkey/infection_monkey/main.py

@@ -7,6 +7,7 @@ import logging.config
 import os
 import sys
 import traceback
+from multiprocessing import freeze_support
 
 from infection_monkey.utils.monkey_log_path import get_dropper_log_path, get_monkey_log_path
 from infection_monkey.config import WormConfiguration, EXTERNAL_CONFIG_FILE
@@ -43,7 +44,7 @@ def main():
 
     if 2 > len(sys.argv):
         return True
-
+    freeze_support()  # required for multiprocessing + pyinstaller on windows
     monkey_mode = sys.argv[1]
 
     if not (monkey_mode in [MONKEY_ARG, DROPPER_ARG]):

monkey/infection_monkey/model/victim_host_generator.py

@@ -0,0 +1,45 @@
+from infection_monkey.model.host import VictimHost
+
+
+class VictimHostGenerator(object):
+    def __init__(self, network_ranges, blocked_ips, same_machine_ips):
+        self.blocked_ips = blocked_ips
+        self.ranges = network_ranges
+        self.local_addresses = same_machine_ips
+
+    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.local_addresses:
+            return False
+        if ip_address in self.blocked_ips:
+            return False
+        return True

monkey/infection_monkey/model/victim_host_generator_test.py

@@ -0,0 +1,46 @@
+from unittest import TestCase
+from infection_monkey.model.victim_host_generator import VictimHostGenerator
+from common.network.network_range import CidrRange, SingleIpRange
+
+
+class VictimHostGeneratorTester(TestCase):
+
+    def setUp(self):
+        self.cidr_range = CidrRange("10.0.0.0/28", False)  # this gives us 15 hosts
+        self.local_host_range = SingleIpRange('localhost')
+        self.random_single_ip_range = SingleIpRange('41.50.13.37')
+
+    def test_chunking(self):
+        chunk_size = 3
+        # current test setup is 15+1+1-1 hosts
+        test_ranges = [self.cidr_range, self.local_host_range, self.random_single_ip_range]
+        generator = VictimHostGenerator(test_ranges, '10.0.0.1', [])
+        victims = generator.generate_victims(chunk_size)
+        for i in range(5):  # quickly check the equally sided chunks
+            self.assertEqual(len(victims.next()), chunk_size)
+        victim_chunk_last = victims.next()
+        self.assertEqual(len(victim_chunk_last), 1)
+
+    def test_remove_blocked_ip(self):
+        generator = VictimHostGenerator(self.cidr_range, ['10.0.0.1'], [])
+
+        victims = list(generator.generate_victims_from_range(self.cidr_range))
+        self.assertEqual(len(victims), 14)  # 15 minus the 1 we blocked
+
+    def test_remove_local_ips(self):
+        generator = VictimHostGenerator([], [], [])
+        generator.local_addresses = ['127.0.0.1']
+        victims = list(generator.generate_victims_from_range(self.local_host_range))
+        self.assertEqual(len(victims), 0)  # block the local IP
+
+    def test_generate_domain_victim(self):
+        # domain name victim
+        generator = VictimHostGenerator([], [], [])  # dummy object
+        victims = list(generator.generate_victims_from_range(self.local_host_range))
+        self.assertEqual(len(victims), 1)
+        self.assertEqual(victims[0].domain_name, 'localhost')
+
+        # don't generate for other victims
+        victims = list(generator.generate_victims_from_range(self.random_single_ip_range))
+        self.assertEqual(len(victims), 1)
+        self.assertEqual(victims[0].domain_name, '')

monkey/infection_monkey/network/network_scanner.py

@@ -1,28 +1,33 @@
 import time
+import logging
 
-from common.network.network_range import *
+from common.network.network_range import NetworkRange
 from infection_monkey.config import WormConfiguration
+from infection_monkey.model.victim_host_generator import VictimHostGenerator
 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
 
-__author__ = 'itamar'
+if is_windows_os():
+    from multiprocessing.dummy import Pool
+else:
+    from multiprocessing import Pool
 
 LOG = logging.getLogger(__name__)
 
-SCAN_DELAY = 0
+ITERATION_BLOCK_SIZE = 5
 
 
 class NetworkScanner(object):
     def __init__(self):
         self._ip_addresses = None
         self._ranges = None
+        self.scanners = [TcpScanner(), PingScanner()]
 
     def initialize(self):
         """
         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()
@@ -68,49 +73,35 @@ 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
         """
+        # 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)
+        victim_generator = VictimHostGenerator(self._ranges, WormConfiguration.blocked_ips, local_ips())
 
-        TCPscan = TcpScanner()
-        Pinger = PingScanner()
         victims_count = 0
+        for victim_chunk in victim_generator.generate_victims(ITERATION_BLOCK_SIZE):
+            LOG.debug("Scanning for potential victims in chunk %r", victim_chunk)
 
-        for net_range in self._ranges:
+            # check before running scans
-            LOG.debug("Scanning for potential victims in the network %r", net_range)
+            if stop_callback and stop_callback():
-            for ip_addr in net_range:
+                LOG.debug("Got stop signal")
-                if hasattr(net_range, 'domain_name'):
+                return
-                    victim = VictimHost(ip_addr, net_range.domain_name)
-                else:
-                    victim = VictimHost(ip_addr)
-                if stop_callback and stop_callback():
-                    LOG.debug("Got stop signal")
-                    break
 
-                # skip self IP address
+            results = pool.map(self.scan_machine, victim_chunk)
-                if victim.ip_addr in self._ip_addresses:
+            resulting_victims = filter(lambda x: x is not None, results)
-                    continue
+            for victim in resulting_victims:
+                LOG.debug("Found potential victim: %r", victim)
+                victims_count += 1
+                yield victim
 
-                # skip IPs marked as blocked
+                if victims_count >= max_find:
-                if victim.ip_addr in WormConfiguration.blocked_ips:
+                    LOG.debug("Found max needed victims (%d), stopping scan", max_find)
-                    LOG.info("Skipping %s due to blacklist" % victim)
+                    return
-                    continue
+            if WormConfiguration.tcp_scan_interval:
-
+                # time.sleep uses seconds, while config is in milliseconds
-                LOG.debug("Scanning %r...", victim)
+                time.sleep(WormConfiguration.tcp_scan_interval / float(1000))
-                pingAlive = Pinger.is_host_alive(victim)
-                tcpAlive = TCPscan.is_host_alive(victim)
-
-                # if scanner detect machine is up, add it to victims list
-                if pingAlive or tcpAlive:
-                    LOG.debug("Found potential victim: %r", victim)
-                    victims_count += 1
-                    yield victim
-
-                    if victims_count >= max_find:
-                        LOG.debug("Found max needed victims (%d), stopping scan", max_find)
-
-                        break
-
-                if WormConfiguration.tcp_scan_interval:
-                    # 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):
@@ -119,5 +110,18 @@ class NetworkScanner(object):
                 return True
         return False
 
+    def scan_machine(self, victim):
+        """
+        Scans specific machine using instance scanners
+        :param victim: VictimHost machine
+        :return: Victim or None if victim isn't alive
+        """
+        LOG.debug("Scanning target address: %r", victim)
+        if any([scanner.is_host_alive(victim) for scanner in self.scanners]):
+            LOG.debug("Found potential target_ip: %r", victim)
+            return victim
+        else:
+            return None
+
     def on_island(self, server):
         return bool([x for x in self._ip_addresses if x in server])

monkey/monkey_island/cc/services/config_schema.py

@@ -448,13 +448,13 @@ SCHEMA = {
                         "victims_max_find": {
                             "title": "Max victims to find",
                             "type": "integer",
-                            "default": 30,
+                            "default": 100,
                             "description": "Determines the maximum number of machines the monkey is allowed to scan"
                         },
                         "victims_max_exploit": {
                             "title": "Max victims to exploit",
                             "type": "integer",
-                            "default": 7,
+                            "default": 15,
                             "description":
                                 "Determines the maximum number of machines the monkey"
                                 " is allowed to successfully exploit. " + WARNING_SIGN