PEP8 changes

Documentation improvements
This commit is contained in:
Daniel Goldberg 2019-09-16 16:32:13 +03:00
parent d8bac57eb5
commit d8d55cb546
2 changed files with 14 additions and 13 deletions

View File

@ -1,17 +1,18 @@
import sys import sys
import itertools
import time import time
import logging 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 common.network.network_range import NetworkRange
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
from infection_monkey.model import VictimHost from infection_monkey.model import VictimHost
from infection_monkey.network import TcpScanner, PingScanner 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' __author__ = 'itamar'
@ -26,7 +27,6 @@ def generate_victims(net_ranges, chunk_size):
Generates VictimHosts in chunks from all the netranges Generates VictimHosts in chunks from all the netranges
:param net_ranges: Iterable of network ranges :param net_ranges: Iterable of network ranges
:param chunk_size: Maximum size of each chunk :param chunk_size: Maximum size of each chunk
:return:
""" """
chunk = [] chunk = []
for net_range in net_ranges: for net_range in net_ranges:
@ -51,7 +51,6 @@ class NetworkScanner(object):
""" """
Set up scanning. Set up scanning.
based on configuration: scans local network and/or scans fixed list of IPs/subnets. based on configuration: scans local network and/or scans fixed list of IPs/subnets.
:return:
""" """
# get local ip addresses # get local ip addresses
self._ip_addresses = local_ips() 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 :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() # 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 victims_count = 0
for victim_chunk in generate_victims(self._ranges, ITERATION_BLOCK_SIZE): for victim_chunk in generate_victims(self._ranges, ITERATION_BLOCK_SIZE):
LOG.debug("Scanning for potential victims in chunk %r", victim_chunk) 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 uses seconds, while config is in milliseconds
time.sleep(WormConfiguration.tcp_scan_interval / float(1000)) time.sleep(WormConfiguration.tcp_scan_interval / float(1000))
@staticmethod @staticmethod
def _is_any_ip_in_subnet(ip_addresses, subnet_str): def _is_any_ip_in_subnet(ip_addresses, subnet_str):
for ip_address in ip_addresses: for ip_address in ip_addresses:
@ -139,10 +142,9 @@ class NetworkScanner(object):
return True return True
return False return False
def scan_machine(self, victim): def scan_machine(self, victim):
""" """
Scans specific machine using given scanner Scans specific machine using instance scanners
:param victim: VictimHost machine :param victim: VictimHost machine
:return: Victim or None if victim isn't alive :return: Victim or None if victim isn't alive
""" """
@ -153,6 +155,5 @@ class NetworkScanner(object):
else: else:
return None 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])