forked from p15670423/monkey
Agent: Remove disused VictimHostGenerator
This commit is contained in:
parent
f0fed888cb
commit
048817d60a
|
@ -1,45 +0,0 @@
|
||||||
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
|
|
|
@ -1,41 +0,0 @@
|
||||||
from unittest import TestCase
|
|
||||||
|
|
||||||
from common.network.network_range import CidrRange, SingleIpRange
|
|
||||||
from infection_monkey.model.victim_host_generator import VictimHostGenerator
|
|
||||||
|
|
||||||
|
|
||||||
class TestVictimHostGenerator(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(next(victims)), chunk_size)
|
|
||||||
victim_chunk_last = next(victims)
|
|
||||||
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")
|
|
Loading…
Reference in New Issue