From e57ce1099f4a830113ee31094118638c6cb7cf84 Mon Sep 17 00:00:00 2001 From: Itay Mizeretz Date: Mon, 26 Feb 2018 17:53:16 +0200 Subject: [PATCH 1/2] Remove unecessary parameters and members. Create better abstraction --- infection_monkey/network/network_scanner.py | 3 ++- infection_monkey/network/range.py | 20 ++++++-------------- 2 files changed, 8 insertions(+), 15 deletions(-) diff --git a/infection_monkey/network/network_scanner.py b/infection_monkey/network/network_scanner.py index 7bdddc904..df1b4f369 100644 --- a/infection_monkey/network/network_scanner.py +++ b/infection_monkey/network/network_scanner.py @@ -45,7 +45,8 @@ class NetworkScanner(object): for net_range in self._ranges: LOG.debug("Scanning for potential victims in the network %r", net_range) - for victim in net_range: + for ip_addr in net_range: + victim = VictimHost(ip_addr) if stop_callback and stop_callback(): LOG.debug("Got stop signal") break diff --git a/infection_monkey/network/range.py b/infection_monkey/network/range.py index 4eb4e122c..31a88e0df 100644 --- a/infection_monkey/network/range.py +++ b/infection_monkey/network/range.py @@ -5,18 +5,14 @@ from abc import ABCMeta, abstractmethod import ipaddress -from model.host import VictimHost - __author__ = 'itamar' class NetworkRange(object): __metaclass__ = ABCMeta - def __init__(self, base_address, shuffle=True): - self._base_address = base_address + def __init__(self, shuffle=True): self._shuffle = shuffle - self._config = __import__('config').WormConfiguration def get_range(self): return [x for x in self._get_range() if (x & 0xFF != 0)] # remove broadcast ips @@ -27,7 +23,7 @@ class NetworkRange(object): random.shuffle(base_range) for x in base_range: - yield VictimHost(self._number_to_ip(self._base_address + x)) + yield self._number_to_ip(x) @abstractmethod def is_in_range(self, ip_address): @@ -59,13 +55,12 @@ class NetworkRange(object): class CidrRange(NetworkRange): def __init__(self, cidr_range, shuffle=True): - base_address = 0 - super(CidrRange, self).__init__(base_address, shuffle=shuffle) + super(CidrRange, self).__init__(shuffle=shuffle) self._cidr_range = cidr_range.strip() self._ip_network = ipaddress.ip_network(unicode(self._cidr_range), strict=False) def __repr__(self): - return "" % (self._cidr_range, ) + return "" % (self._cidr_range,) def is_in_range(self, ip_address): return ipaddress.ip_address(ip_address) in self._ip_network @@ -76,8 +71,7 @@ class CidrRange(NetworkRange): class IpRange(NetworkRange): def __init__(self, ip_range=None, lower_end_ip=None, higher_end_ip=None, shuffle=True): - base_address = 0 - super(IpRange, self).__init__(base_address, shuffle=shuffle) + super(IpRange, self).__init__(shuffle=shuffle) if ip_range is not None: addresses = ip_range.split('-') if len(addresses) != 2: @@ -106,8 +100,7 @@ class IpRange(NetworkRange): class SingleIpRange(NetworkRange): def __init__(self, ip_address, shuffle=True): - base_address = 0 - super(SingleIpRange, self).__init__(base_address, shuffle=shuffle) + super(SingleIpRange, self).__init__(shuffle=shuffle) self._ip_address = ip_address def __repr__(self): @@ -118,4 +111,3 @@ class SingleIpRange(NetworkRange): def _get_range(self): return [SingleIpRange._ip_to_number(self._ip_address)] - From 1d07e5f98fc792cbf3ea513b6106570a72d3a181 Mon Sep 17 00:00:00 2001 From: Itay Mizeretz Date: Mon, 26 Feb 2018 18:39:49 +0200 Subject: [PATCH 2/2] Move range to common code folder --- common/__init__.py | 1 + common/network/__init__.py | 1 + {infection_monkey => common}/network/range.py | 0 infection_monkey/monkey-linux.spec | 2 +- infection_monkey/monkey.spec | 2 +- infection_monkey/network/info.py | 2 +- infection_monkey/network/network_scanner.py | 3 ++- 7 files changed, 7 insertions(+), 4 deletions(-) create mode 100644 common/__init__.py create mode 100644 common/network/__init__.py rename {infection_monkey => common}/network/range.py (100%) diff --git a/common/__init__.py b/common/__init__.py new file mode 100644 index 000000000..ee5b79ad0 --- /dev/null +++ b/common/__init__.py @@ -0,0 +1 @@ +__author__ = 'itay.mizeretz' diff --git a/common/network/__init__.py b/common/network/__init__.py new file mode 100644 index 000000000..ee5b79ad0 --- /dev/null +++ b/common/network/__init__.py @@ -0,0 +1 @@ +__author__ = 'itay.mizeretz' diff --git a/infection_monkey/network/range.py b/common/network/range.py similarity index 100% rename from infection_monkey/network/range.py rename to common/network/range.py diff --git a/infection_monkey/monkey-linux.spec b/infection_monkey/monkey-linux.spec index c8c4c631b..fac69536e 100644 --- a/infection_monkey/monkey-linux.spec +++ b/infection_monkey/monkey-linux.spec @@ -4,7 +4,7 @@ block_cipher = None a = Analysis(['main.py'], - pathex=['.'], + pathex=['.', '..'], binaries=None, datas=None, hiddenimports=['_cffi_backend'], diff --git a/infection_monkey/monkey.spec b/infection_monkey/monkey.spec index 8e004145b..cb9c6130e 100644 --- a/infection_monkey/monkey.spec +++ b/infection_monkey/monkey.spec @@ -2,7 +2,7 @@ import os import platform a = Analysis(['main.py'], - pathex=['.'], + pathex=['.', '..'], hiddenimports=['_cffi_backend', 'queue'], hookspath=None, runtime_hooks=None) diff --git a/infection_monkey/network/info.py b/infection_monkey/network/info.py index 30a75780e..da30e7b0f 100644 --- a/infection_monkey/network/info.py +++ b/infection_monkey/network/info.py @@ -8,7 +8,7 @@ import itertools import netifaces from subprocess import check_output from random import randint -from range import CidrRange +from common.network.range import CidrRange def get_host_subnets(): diff --git a/infection_monkey/network/network_scanner.py b/infection_monkey/network/network_scanner.py index df1b4f369..7837e36c3 100644 --- a/infection_monkey/network/network_scanner.py +++ b/infection_monkey/network/network_scanner.py @@ -3,7 +3,8 @@ import time from config import WormConfiguration from info import local_ips, get_interfaces_ranges -from range import * +from common.network.range import * +from model import VictimHost from . import HostScanner __author__ = 'itamar'