diff --git a/cornershot/cornershot.py b/cornershot/cornershot.py index ec29ed7..fe6416e 100644 --- a/cornershot/cornershot.py +++ b/cornershot/cornershot.py @@ -3,7 +3,7 @@ import time from random import uniform,shuffle -from .shots import PORT_UNKNOWN,PORT_FILTERED +from .shots import PORT_UNKNOWN,PORT_FILTERED,PORT_OPEN from .shots.even import EVENShot from .shots.even6 import EVEN6Shot from .shots.rprn import RPRNShot @@ -37,6 +37,7 @@ def __init__(self, username, password, domain, workers=250, shots=None): self.current_tasks = [] self.skip_scanned = False self.already_scanned = [] + self.batch_scanned_event = None def _takeashot(self): while self.runthreads: @@ -92,7 +93,7 @@ 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: + if self.skip_scanned and PORT_OPEN in state: tp_pair = target + ":" + str(tport) if tp_pair not in self.already_scanned: self.already_scanned.append(tp_pair) @@ -140,15 +141,14 @@ def _get_next_tasks(self,remaining): return new_tasks - def _shots_manager(self): - remaining = MAX_QUEUE_SIZE + remaining = min(self.workers,MAX_QUEUE_SIZE) + shuffle(self.shot_list) self.total_shots = len(self.shot_list) + while self.runthreads: self.current_tasks = self._get_next_tasks(remaining) - shuffle(self.current_tasks) - remaining = remaining - len(self.current_tasks) for bt in self.current_tasks: @@ -163,28 +163,34 @@ def _shots_manager(self): self.resultQ.task_done() remaining += 1 self.total_shots -= 1 + if self.total_shots % 500 == 0: + self.batch_scanned_event.set() if self.total_shots < 1: self.runthreads = False except (TimeoutError,queue.Empty): break self.total_shots = 0 + self.batch_scanned_event.set() def open_fire(self,blocking=True,skip_scanned=False): self.skip_scanned = skip_scanned - num_threads = min(self.total_shots,self.workers) + self.workers = min(self.total_shots,self.workers) if self.total_shots > 0: - for _ in range(num_threads): + for _ in range(self.workers): w = threading.Thread(target=self._takeashot, daemon=True) w.start() + if blocking: self._shots_manager() return self.results else: main_thread = threading.Thread(target=self._shots_manager,daemon=True) main_thread.start() + self.batch_scanned_event = threading.Event() + return self.batch_scanned_event def read_results(self): return self.results diff --git a/cornershot/shots/__init__.py b/cornershot/shots/__init__.py index 7ba499c..487f43f 100644 --- a/cornershot/shots/__init__.py +++ b/cornershot/shots/__init__.py @@ -1,4 +1,5 @@ import time +import asyncio from abc import ABCMeta, abstractmethod from impacket.dcerpc.v5 import transport @@ -80,7 +81,7 @@ def close(self): except Exception as err: pass - def connect_and_bind(self): + async def connect_and_bind(self): try: rpctransport = transport.DCERPCTransportFactory(self.do_binding()) if hasattr(rpctransport, 'set_credentials'): @@ -96,15 +97,15 @@ def connect_and_bind(self): except Exception as err: logger.debug(f'{type(self).__name__} - Connection failed for {self.destination}->{self.target}:{self.trgt_port} - {err}') - def shoot(self): + async def shoot(self): err = None state = PORT_UNKNOWN - self.connect_and_bind() + await self.connect_and_bind() elapsed = 0 if self.dce: start = time.time() try: - self.do_rpc_logic() + await self.do_rpc_logic() except DCERPCException as e: err = e except Exception as e: diff --git a/setup.py b/setup.py index 0694812..629beca 100644 --- a/setup.py +++ b/setup.py @@ -15,7 +15,7 @@ setup( name='cornershot', python_requires='>=3', - version='0.2.9', + version='0.2.10', description='Library to test network connectivity', long_description_content_type='text/markdown', long_description=long_description,