diff --git a/cornershot/cornershot.py b/cornershot/cornershot.py index 6b53eba..963faba 100644 --- a/cornershot/cornershot.py +++ b/cornershot/cornershot.py @@ -17,7 +17,7 @@ DEFAULT_SHOTS = [EVENShot, RPRNShot, RRPShot, EVEN6Shot] class CornerShot(object): - def __init__(self, username, password, domain, workers=250, shots=None): + def __init__(self, username, password, domain, skip_scanned=False, workers=250, shots=None): logger.debug(f'CS created with username: {username},domain:{domain},workers:{workers}') if shots is None: @@ -34,6 +34,8 @@ def __init__(self, username, password, domain, workers=250, shots=None): self.results = {} self.shot_list = [] self.total_shots = 0 + self.skip_scanned = skip_scanned + self.already_scanned = [] def _takeashot(self): while self.runthreads: @@ -86,6 +88,11 @@ def _shots_generator(self, destinations, targets, target_ports, destination_port self.shot_list.append(cls(self.username, self.password, self.domain, destination, target,target_port=target_port)) def _merge_result(self, dest, target, tport, state): + if self.skip_scanned and PORT_UNKNOWN not in state: + tp_pair = target + ":" + str(tport) + if tp_pair not in self.already_scanned: + self.already_scanned.append(tp_pair) + if dest not in self.results: self.results[dest] = {} if target not in self.results[dest]: @@ -103,11 +110,36 @@ def _merge_result(self, dest, target, tport, state): elif state not in self.results[dest][target][tport]: self.results[dest][target][tport] += "|" + state + def _get_next_tasks(self,remaining): + + new_tasks = [] + remaining = min(len(self.shot_list),remaining) + if self.skip_scanned: + iterated_shots = 0 + for bullet in self.shot_list: + if remaining > 0: + trgt = bullet.target + ":" + str(bullet.trgt_port) + if trgt not in self.already_scanned: + new_tasks.append(bullet) + iterated_shots += 1 + remaining -= 1 + else: + break + self.shot_list = self.shot_list[iterated_shots + 1:] + + else: + new_tasks = self.shot_list[0:remaining] + self.shot_list = self.shot_list[remaining + 1:] + + return new_tasks + + def _shots_manager(self): remaining = MAX_QUEUE_SIZE while self.runthreads: - new_tasks = self.shot_list[0:remaining] - self.shot_list = self.shot_list[remaining + 1:] + new_tasks = self._get_next_tasks(remaining) + # new_tasks = self.shot_list[0:remaining] + # self.shot_list = self.shot_list[remaining + 1:] tasks = new_tasks shuffle(tasks) diff --git a/setup.py b/setup.py index a9d9f62..b2b41f1 100644 --- a/setup.py +++ b/setup.py @@ -15,7 +15,7 @@ setup( name='cornershot', python_requires='>=3', - version='0.2.2', + version='0.2.1', description='Library to test network connectivity', long_description_content_type='text/markdown', long_description=long_description,