p15670423
/
monkey
1
0
Fork 5
Code Issues Pull Requests 3 Projects Releases Wiki Activity

Changed VictimHostGenerator to accept the local addresses rather than generating them itself. 

Changed UTs to be independent.
This commit is contained in:
Daniel Goldberg 2019-09-27 18:10:59 +03:00
parent 297686dc53
commit f55a3e483b
3 changed files with 19 additions and 18 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
monkey/infection_monkey/model/victim_host_generator.py
View File

@ -1,12 +1,11 @@
from infection_monkey.model.host import VictimHost
from infection_monkey.network.info import local_ips
class VictimHostGenerator(object):
def __init__(self, network_ranges, blocked_ips):
self._ip_addresses = local_ips()
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):
"""
@ -39,7 +38,7 @@ class VictimHostGenerator(object):
yield victim
def is_ip_scannable(self, ip_address):
if ip_address in self._ip_addresses:
if ip_address in self.local_addresses:
return False
if ip_address in self.blocked_ips:
return False

28
monkey/infection_monkey/model/victim_host_generator_test.py
View File

@ -6,39 +6,41 @@ from common.network.network_range import CidrRange, SingleIpRange
class VictimHostGeneratorTester(TestCase):
def setUp(self):
self.test_ranges = [CidrRange("10.0.0.0/28", False), # this gives us 15 hosts
SingleIpRange('41.50.13.37'),
SingleIpRange('localhost')
]
self.generator = VictimHostGenerator(self.test_ranges, '10.0.0.1')
self.generator._ip_addresses = [] # test later on
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
victims = self.generator.generate_victims(chunk_size)
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):
victims = list(self.generator.generate_victims_from_range(self.test_ranges[0]))
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):
self.generator._ip_addresses = ['127.0.0.1']
victims = list(self.generator.generate_victims_from_range(self.test_ranges[-1]))
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
self.generator._ip_addresses = []
victims = list(self.generator.generate_victims_from_range(self.test_ranges[-1]))
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(self.generator.generate_victims_from_range(self.test_ranges[1]))
victims = list(generator.generate_victims_from_range(self.random_single_ip_range))
self.assertEqual(len(victims), 1)
self.assertEqual(victims[0].domain_name, '')

2
monkey/infection_monkey/network/network_scanner.py
View File

@ -78,7 +78,7 @@ class NetworkScanner(object):
# 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)
victim_generator = VictimHostGenerator(self._ranges, WormConfiguration.blocked_ips, local_ips())
victims_count = 0
for victim_chunk in victim_generator.generate_victims(ITERATION_BLOCK_SIZE):