From 9697e80013d7336472e45ba59c4b0a6fa9f874e8 Mon Sep 17 00:00:00 2001
From: Miroslav Stampar <miroslav.stampar@gmail.com>
Date: Tue, 22 Nov 2011 10:54:29 +0000
Subject: [PATCH] some more optimizations

---
 _sqlmap.py                     |  3 +++
 lib/core/settings.py           |  5 ++++-
 lib/core/threads.py            |  2 +-
 lib/utils/hashdb.py            | 26 ++++++++++++++++----------
 plugins/generic/enumeration.py |  2 --
 5 files changed, 24 insertions(+), 14 deletions(-)

diff --git a/_sqlmap.py b/_sqlmap.py
index 8926b3ead..2299e2295 100755
--- a/_sqlmap.py
+++ b/_sqlmap.py
@@ -129,6 +129,9 @@ def main():
         kb.threadContinue = False
         kb.threadException = True
 
+        if conf.hashDB:
+            conf.hashDB.flush(True)
+
         # Reference: http://stackoverflow.com/questions/1635080/terminate-a-multi-thread-python-program
         if hasattr(conf, "threads") and conf.threads > 1:
             os._exit(0)
diff --git a/lib/core/settings.py b/lib/core/settings.py
index 3a47ce057..7338767bd 100644
--- a/lib/core/settings.py
+++ b/lib/core/settings.py
@@ -414,4 +414,7 @@ COMMON_USER_COLUMNS = ('user', 'username', 'user_name', 'benutzername', 'benutze
 DEFAULT_GET_POST_DELIMITER = '&'
 
 # Default delimiter in cookie values
-DEFAULT_COOKIE_DELIMITER = ';'
\ No newline at end of file
+DEFAULT_COOKIE_DELIMITER = ';'
+
+# Skip unforced HashDB flush requests below the threshold number of cached items
+HASHDB_FLUSH_THRESHOLD = 10
diff --git a/lib/core/threads.py b/lib/core/threads.py
index 2fc5db15a..bb33d2bd4 100644
--- a/lib/core/threads.py
+++ b/lib/core/threads.py
@@ -181,7 +181,7 @@ def runThreads(numThreads, threadFunction, cleanupFunction=None, forwardExceptio
         kb.threadContinue = True
         kb.threadException = False
 
-        conf.hashDB.flush()
+        conf.hashDB.flush(True)
 
         if cleanupFunction:
             cleanupFunction()
diff --git a/lib/utils/hashdb.py b/lib/utils/hashdb.py
index 51ec96652..0e3788fb4 100644
--- a/lib/utils/hashdb.py
+++ b/lib/utils/hashdb.py
@@ -12,6 +12,7 @@ import sqlite3
 import threading
 
 from lib.core.data import conf
+from lib.core.settings import HASHDB_FLUSH_THRESHOLD
 from lib.core.settings import UNICODE_ENCODING
 from lib.core.threads import getCurrentThreadData
 from lib.core.threads import getCurrentThreadName
@@ -54,15 +55,17 @@ class HashDB(object):
         retVal = None
         if key:
             hash_ = HashDB.hashKey(key)
-            while True:
-                try:
-                    for row in self.cursor.execute("SELECT value FROM storage WHERE id=?", (hash_,)):
-                        retVal = row[0]
-                except sqlite3.OperationalError, ex:
-                    if not 'locked' in ex.message:
-                        raise
-                else:
-                    break
+            retVal = self._write_cache.get(hash_, None)
+            if not retVal:
+                while True:
+                    try:
+                        for row in self.cursor.execute("SELECT value FROM storage WHERE id=?", (hash_,)):
+                            retVal = row[0]
+                    except sqlite3.OperationalError, ex:
+                        if not 'locked' in ex.message:
+                            raise
+                    else:
+                        break
         return retVal
 
     def write(self, key, value):
@@ -75,10 +78,13 @@ class HashDB(object):
         if getCurrentThreadName() in ('0', 'MainThread'):
             self.flush()
 
-    def flush(self):
+    def flush(self, forced=False):
         if not self._write_cache:
             return
 
+        if not forced and len(self._write_cache) < HASHDB_FLUSH_THRESHOLD:
+            return
+
         self._cache_lock.acquire()
         items = self._write_cache.items()
         self._write_cache.clear()
diff --git a/plugins/generic/enumeration.py b/plugins/generic/enumeration.py
index 8f4bc6f41..411c3a9d1 100644
--- a/plugins/generic/enumeration.py
+++ b/plugins/generic/enumeration.py
@@ -1564,7 +1564,6 @@ class Enumeration:
 
             try:
                 kb.dumpMode = True
-                #conf.hashDB.beginTransaction()
 
                 if not safeSQLIdentificatorNaming(conf.db) in kb.data.cachedColumns \
                    or safeSQLIdentificatorNaming(tbl, True) not in \
@@ -1789,7 +1788,6 @@ class Enumeration:
 
             finally:
                 kb.dumpMode = False
-                #conf.hashDB.endTransaction()
 
     def dumpAll(self):
         if conf.db is not None and conf.tbl is None: