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_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")
# don't generate for other victims
victims = list(generator.generate_victims_from_range(self.random_single_ip_range))
self.assertEqual(len(victims), 1)
self.assertEqual(victims[0].domain_name, "")
def get_victim_machines(self, max_find=5, stop_callback=None):
"""
Finds machines according to the ranges specified in the object
:param max_find: Max number of victims to find regardless of ranges
:param stop_callback: A callback to check at any point if we should stop scanning
:return: yields a sequence of VictimHost instances
"""
# 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)
victim_generator = VictimHostGenerator(self._ranges,
WormConfiguration.blocked_ips,
local_ips())
victims_count = 0
for victim_chunk in victim_generator.generate_victims(
ITERATION_BLOCK_SIZE):
LOG.debug("Scanning for potential victims in chunk %r",
victim_chunk)
# check before running scans
if stop_callback and stop_callback():
LOG.debug("Got stop signal")
return
results = pool.map(self.scan_machine, victim_chunk)
resulting_victims = filter(lambda x: x is not None, results)
for victim in resulting_victims:
LOG.debug("Found potential victim: %r", victim)
victims_count += 1
yield victim
if victims_count >= max_find:
LOG.debug("Found max needed victims (%d), stopping scan",
max_find)
return
if WormConfiguration.tcp_scan_interval:
# time.sleep uses seconds, while config is in milliseconds
time.sleep(WormConfiguration.tcp_scan_interval / float(1000))
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_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