implementation of multithreading for UNION and ERROR techniques

lib/core/threads.py

@@ -11,6 +11,11 @@ import difflib
 import threading
 
 from lib.core.data import kb
+from lib.core.data import logger
+from lib.core.datatype import advancedDict
+from lib.core.exception import sqlmapThreadException
+
+shared = advancedDict()
 
 class ThreadData():
     """
@@ -18,6 +23,8 @@ class ThreadData():
     """
 
     def __init__(self):
+        global shared
+
         self.disableStdOut = False
         self.lastErrorPage = None
         self.lastHTTPError = None
@@ -26,6 +33,7 @@ class ThreadData():
         self.lastRequestUID = 0
         self.retriesCount = 0
         self.seqMatcher = difflib.SequenceMatcher(None)
+        self.shared = shared
         self.valueStack = []
 
 def getCurrentThreadUID():
@@ -40,3 +48,55 @@ def getCurrentThreadData():
     if threadUID not in kb.threadData:
         kb.threadData[threadUID] = ThreadData()
     return kb.threadData[threadUID]
+
+def runThreads(numThreads, threadFunction, cleanupFunction=None):
+    threads = []
+
+    kb.threadContinue = True
+    kb.threadException = False
+
+    if numThreads > 1:
+        infoMsg = "starting %d threads" % numThreads
+        logger.info(infoMsg)
+    else:
+        threadFunction()
+        return
+
+    # Start the threads
+    for numThread in range(numThreads):
+        thread = threading.Thread(target=threadFunction, name=str(numThread))
+        thread.start()
+        threads.append(thread)
+
+    # And wait for them to all finish
+    try:
+        alive = True
+
+        while alive:
+            alive = False
+
+            for thread in threads:
+                if thread.isAlive():
+                    alive = True
+                    thread.join(1)
+
+    except KeyboardInterrupt:
+        kb.threadContinue = False
+        kb.threadException = True
+
+        print
+        logger.debug("waiting for threads to finish")
+
+        try:
+            while (threading.activeCount() > 1):
+                pass
+
+        except KeyboardInterrupt:
+            raise sqlmapThreadException, "user aborted (Ctrl+C was pressed multiple times)"
+
+    finally:
+        kb.threadContinue = True
+        kb.threadException = False
+
+        if cleanupFunction:
+            cleanupFunction()