Move scanners to be instance variable.

Add MP support (threading/process) for scanning victims in chunks
This commit is contained in:
Daniel Goldberg 2019-09-09 11:37:14 +03:00
parent 7357c8c168
commit 6c5d6a5ecc
1 changed files with 31 additions and 17 deletions

View File

@ -1,6 +1,10 @@
import sys
import itertools import itertools
import time import time
if sys.platform.startswith("win"):
from multiprocessing.dummy import Pool
else:
from multiprocessing import Pool
from common.network.network_range import * from common.network.network_range import *
from infection_monkey.config import WormConfiguration from infection_monkey.config import WormConfiguration
from infection_monkey.network.info import local_ips, get_interfaces_ranges from infection_monkey.network.info import local_ips, get_interfaces_ranges
@ -34,6 +38,7 @@ class NetworkScanner(object):
def __init__(self): def __init__(self):
self._ip_addresses = None self._ip_addresses = None
self._ranges = None self._ranges = None
self.scanners = [TcpScanner(), PingScanner()]
def initialize(self): def initialize(self):
""" """
@ -85,17 +90,15 @@ class NetworkScanner(object):
:param stop_callback: A callback to check at any point if we should stop scanning :param stop_callback: A callback to check at any point if we should stop scanning
:return: yields a sequence of VictimHost instances :return: yields a sequence of VictimHost instances
""" """
pool = Pool()
tcp_scan = TcpScanner()
ping_scan = PingScanner()
victims_count = 0 victims_count = 0
for network_chunk in _grouper(self._ranges, ITERATION_BLOCK_SIZE): for network_chunk in _grouper(self._ranges, ITERATION_BLOCK_SIZE):
LOG.debug("Scanning for potential victims in chunk %r", network_chunk) LOG.debug("Scanning for potential victims in chunk %r", network_chunk)
victim_chunk = [] victim_chunk = []
for address in network_chunk: for address in network_chunk:
#if hasattr(net_range, 'domain_name'): # if hasattr(net_range, 'domain_name'):
# victim = VictimHost(address, net_range.domain_name) # victim = VictimHost(address, net_range.domain_name)
#else: # else:
victim = VictimHost(address) victim = VictimHost(address)
victim_chunk.append(victim) victim_chunk.append(victim)
@ -113,19 +116,17 @@ class NetworkScanner(object):
LOG.debug("Got stop signal") LOG.debug("Got stop signal")
break break
for victim in victim_chunk: results = pool.map(self.scan_machine, victim_chunk)
LOG.debug("Scanning %r...", victim) resulting_victims = [x for x in results if x] # filter out dead addresses
pingAlive = ping_scan.is_host_alive(victim) for victim in resulting_victims:
tcpAlive = tcp_scan.is_host_alive(victim) LOG.debug("Found potential victim: %r", victim)
if pingAlive or tcpAlive: victims_count += 1
LOG.debug("Found potential victim: %r", victim) yield victim
victims_count += 1
yield victim
if victims_count >= max_find: if victims_count >= max_find:
LOG.debug("Found max needed victims (%d), stopping scan", max_find) LOG.debug("Found max needed victims (%d), stopping scan", max_find)
break break
if WormConfiguration.tcp_scan_interval: if WormConfiguration.tcp_scan_interval:
# time.sleep uses seconds, while config is in milliseconds # time.sleep uses seconds, while config is in milliseconds
time.sleep(WormConfiguration.tcp_scan_interval / float(1000)) time.sleep(WormConfiguration.tcp_scan_interval / float(1000))
@ -137,5 +138,18 @@ class NetworkScanner(object):
return True return True
return False return False
def scan_machine(self, victim):
"""
Scans specific machine using given scanner
: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.ip_addr) for scanner in self.scanners]):
LOG.debug("Found potential target_ip: %r", victim)
return victim
else:
return None
def on_island(self, server): def on_island(self, server):
return bool([x for x in self._ip_addresses if x in server]) return bool([x for x in self._ip_addresses if x in server])