ThreadPoolExecutor as context manager

routersploit/modules/scanners/autopwn.py

@@ -33,15 +33,17 @@ class Exploit(exploits.Exploit):
 
     def run(self):
         self.vulnerabilities = []
-        executor = threads.ThreadPoolExecutor(self.threads)
-        executor.feed(utils.iter_modules(utils.EXPLOITS_DIR))
-        executor.run(self.target_function)
 
+        with threads.ThreadPoolExecutor(self.threads) as executor:
+            for exploit in utils.iter_modules(utils.EXPLOITS_DIR):
+                executor.submit(self.target_function, exploit)
+
+        print_info()
         if self.vulnerabilities:
-            print_info()
             print_success("Device is vulnerable!")
             for v in self.vulnerabilities:
                 print_info(" - {}".format(v))
+            print_info()
         else:
             print_error("Device is not vulnerable to any exploits!\n")
 

routersploit/threads.py

@@ -15,33 +15,18 @@ from . import utils
 data_queue = queue.Queue()
 
 
-class DataProducerThread(threading.Thread):
-    def __init__(self, data):
-        super(DataProducerThread, self).__init__(name=self.__class__.__name__)
-        self.data = data
-
-    def run(self):
-        for record in self.data:
-            data_queue.put(record)
-
-    def stop(self):
-        data_queue.queue.clear()
-
-    def join_queue(self):
-        data_queue.join()
-
-
 class WorkerThread(threading.Thread):
-    def __init__(self, target, name):
-        super(WorkerThread, self).__init__(target=target, name=name)
-        self.target = target
+    def __init__(self, name):
+        super(WorkerThread, self).__init__(name=name)
         self.name = name
 
     def run(self):
         while not data_queue.empty():
             record = data_queue.get()
+            target = record[0]
+            args = record[1:]
             try:
-                self.target(record)
+                target(*args)
             finally:
                 data_queue.task_done()
 
@@ -49,21 +34,19 @@ class WorkerThread(threading.Thread):
 class ThreadPoolExecutor(object):
     def __init__(self, threads):
         self.threads = threads
-        self.data_producer = None
-
-    def feed(self, dataset):
-        self.data_producer = DataProducerThread(dataset)
-        self.data_producer.start()
-        time.sleep(0.1)
+        self.workers = []
 
-    def run(self, target):
-        workers = []
+    def __enter__(self):
+        self.workers = []
         for worker_id in xrange(int(self.threads)):
             worker = WorkerThread(
-                target=target,
                 name='worker-{}'.format(worker_id),
             )
-            workers.append(worker)
+            self.workers.append(worker)
+        return self
+
+    def __exit__(self, *args):
+        for worker in self.workers:
             worker.start()
 
         start = time.time()
@@ -73,10 +56,13 @@ class ThreadPoolExecutor(object):
         except KeyboardInterrupt:
             utils.print_info()
             utils.print_status("Waiting for already scheduled jobs to finish...")
-            self.data_producer.stop()
-            for worker in workers:
+            data_queue.queue.clear()
+            for worker in self.workers:
                 worker.join()
         else:
-            self.data_producer.join_queue()
+            data_queue.join()
 
         utils.print_status('Elapsed time: ', time.time() - start, 'seconds')
+
+    def submit(self, *args):
+        data_queue.put(args)