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