From 1da0efed4edc07c8bdef81f4803a6c194b159a96 Mon Sep 17 00:00:00 2001
From: sagiezero <61778471+sagiezero@users.noreply.github.com>
Date: Mon, 8 Feb 2021 12:58:13 +0200
Subject: [PATCH] fix(CornerShot ): total shots count improved

---
 README.md                |  1 -
 cornershot/__main__.py   |  1 -
 cornershot/cornershot.py | 38 ++++++++++++++++++++------------------
 setup.py                 |  2 +-
 4 files changed, 21 insertions(+), 21 deletions(-)

diff --git a/README.md b/README.md
index a82bdd8..48ba5cc 100644
--- a/README.md
+++ b/README.md
@@ -81,7 +81,6 @@ Adding carriers, target and ports is achieved via the *add_shots* method. Once r
 from cornershot import CornerShot
 cs = CornerShot("username", "password", "fqdn")
 cs.add_shots(carriers=["192.168.1.1"],targets=["192.168.1.2","192.168.1.3"])
-cs.lock_and_load()
 results = cs.open_fire()
 ```
 
diff --git a/cornershot/__main__.py b/cornershot/__main__.py
index 8b01719..24035f1 100644
--- a/cornershot/__main__.py
+++ b/cornershot/__main__.py
@@ -110,7 +110,6 @@ def parse_port_ranges(ranges):
 
         cs = CornerShot(args.user, args.password, args.domain, workers=args.threads)
         cs.add_shots(parse_ip_ranges(args.carrier), parse_ip_ranges(args.target), target_ports=parse_port_ranges(args.tports))
-        cs.lock_and_load()
         cs.open_fire()
 
     except KeyboardInterrupt:
diff --git a/cornershot/cornershot.py b/cornershot/cornershot.py
index 963faba..f3c3d22 100644
--- a/cornershot/cornershot.py
+++ b/cornershot/cornershot.py
@@ -11,13 +11,13 @@
 
 from . import logger
 
-MAX_QUEUE_SIZE = 1000
+MAX_QUEUE_SIZE = 3000
 TARGET_PORTS = [135, 445, 3389, 5985, 5986]
 
 DEFAULT_SHOTS = [EVENShot, RPRNShot, RRPShot, EVEN6Shot]
 
 class CornerShot(object):
-    def __init__(self, username, password, domain, skip_scanned=False, workers=250, shots=None):
+    def __init__(self, username, password, domain, workers=250, shots=None):
 
         logger.debug(f'CS created with username: {username},domain:{domain},workers:{workers}')
         if shots is None:
@@ -34,7 +34,8 @@ def __init__(self, username, password, domain, skip_scanned=False, workers=250,
         self.results = {}
         self.shot_list = []
         self.total_shots = 0
-        self.skip_scanned = skip_scanned
+        self.current_tasks = []
+        self.skip_scanned = False
         self.already_scanned = []
 
     def _takeashot(self):
@@ -50,10 +51,12 @@ def _takeashot(self):
                     except Exception:
                         logger.debug(f'Unexpected exception during shot', exc_info=True)
                     finally:
+                        logger.info("Putting bullter to result queue...")
                         self.bulletQ.task_done()
                         self.resultQ.put(res)
+                logger.debug(f'Bullet is none!')
             except (queue.Empty,TimeoutError):
-                pass
+                logger.debug(f'Exception during bullet load', exc_info=True)
             except Exception:
                 logger.debug(f'Unexpected exception during bullet load', exc_info=True)
 
@@ -121,14 +124,16 @@ def _get_next_tasks(self,remaining):
                     trgt = bullet.target + ":" + str(bullet.trgt_port)
                     if trgt not in self.already_scanned:
                         new_tasks.append(bullet)
+                    else:
+                        self.total_shots -= 1
                     iterated_shots += 1
                     remaining -= 1
                 else:
                     break
-            self.shot_list = self.shot_list[iterated_shots + 1:]
+            self.shot_list = self.shot_list[iterated_shots:]
 
         else:
-            new_tasks = self.shot_list[0:remaining]
+            new_tasks = self.shot_list[0:remaining ]
             self.shot_list = self.shot_list[remaining + 1:]
 
         return new_tasks
@@ -136,16 +141,15 @@ def _get_next_tasks(self,remaining):
 
     def _shots_manager(self):
         remaining = MAX_QUEUE_SIZE
+        self.total_shots = len(self.shot_list)
         while self.runthreads:
-            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)
+            self.current_tasks = self._get_next_tasks(remaining)
+
+            shuffle(self.current_tasks)
 
-            remaining = remaining - len(tasks)
+            remaining = remaining - len(self.current_tasks)
 
-            for bt in tasks:
+            for bt in self.current_tasks:
                 time.sleep(uniform(0, 0.026))
                 self.bulletQ.put(bt)
 
@@ -167,8 +171,9 @@ def _shots_manager(self):
 
         self.total_shots = 0
 
-    def open_fire(self,blocking=True):
-        self.lock_and_load()
+    def open_fire(self,blocking=True,skip_scanned=False):
+        self.skip_scanned = skip_scanned
+
         num_threads = min(self.total_shots,self.workers)
 
         if self.total_shots > 0:
@@ -185,9 +190,6 @@ def open_fire(self,blocking=True):
     def read_results(self):
         return self.results
 
-    def lock_and_load(self):
-        self.total_shots = len(self.shot_list)
-
     def remaining_shots(self):
         return self.total_shots
 
diff --git a/setup.py b/setup.py
index a279d99..4af0d46 100644
--- a/setup.py
+++ b/setup.py
@@ -15,7 +15,7 @@
 setup(
     name='cornershot',
     python_requires='>=3',
-    version='0.2.4',
+    version='0.2.5',
     description='Library to test network connectivity',
     long_description_content_type='text/markdown',
     long_description=long_description,