From 744636a8c1f7c7c2ac4acd6a509bf8a2da2f4f06 Mon Sep 17 00:00:00 2001
From: Miroslav Stampar <miroslav.stampar@gmail.com>
Date: Sun, 25 Sep 2011 20:36:32 +0000
Subject: [PATCH] switching to SQLite resume support (on error and union
 techniques this moment)

---
 lib/core/option.py          |   2 +
 lib/core/target.py          |  10 +++
 lib/techniques/error/use.py | 135 ++++++++++++++++++------------------
 lib/techniques/union/use.py | 101 ++++++++++++++-------------
 lib/utils/hashdb.py         |  63 +++++++++++++++++
 5 files changed, 197 insertions(+), 114 deletions(-)
 create mode 100644 lib/utils/hashdb.py

diff --git a/lib/core/option.py b/lib/core/option.py
index 207fee01eee..493cddc4517 100644
--- a/lib/core/option.py
+++ b/lib/core/option.py
@@ -1360,6 +1360,8 @@ def __setConfAttributes():
     conf.dbmsConnector = None
     conf.dbmsHandler = None
     conf.dumpPath = None
+    conf.hashDB = None
+    conf.hashDBFile = None
     conf.httpHeaders = []
     conf.hostname = None
     conf.multipleTargets = False
diff --git a/lib/core/target.py b/lib/core/target.py
index c09dab81e43..4a454c7b2e5 100644
--- a/lib/core/target.py
+++ b/lib/core/target.py
@@ -40,6 +40,7 @@
 from lib.core.settings import URI_INJECTABLE_REGEX
 from lib.core.settings import URI_INJECTION_MARK_CHAR
 from lib.core.settings import USER_AGENT_ALIASES
+from lib.utils.hashdb import HashDB
 from lib.core.xmldump import dumper as xmldumper
 from lib.request.connect import Connect as Request
 
@@ -174,6 +175,9 @@ def __setOutputResume():
     if not conf.sessionFile:
         conf.sessionFile = "%s%ssession" % (conf.outputPath, os.sep)
 
+    if not conf.hashDBFile:
+        conf.hashDBFile = "%s%shashdb" % (conf.outputPath, os.sep)
+
     logger.info("using '%s' as session file" % conf.sessionFile)
 
     if os.path.exists(conf.sessionFile):
@@ -223,6 +227,7 @@ def __setOutputResume():
         else:
             try:
                 os.remove(conf.sessionFile)
+                os.remove(conf.hashDBFile)
                 logger.info("flushing session file")
             except OSError, msg:
                 errMsg = "unable to flush the session file (%s)" % msg
@@ -230,6 +235,7 @@ def __setOutputResume():
 
     try:
         conf.sessionFP = codecs.open(conf.sessionFile, "a", UNICODE_ENCODING)
+        conf.hashDB = HashDB(conf.hashDBFile)
         dataToSessionFile("\n[%s]\n" % time.strftime("%X %x"))
     except IOError:
         errMsg = "unable to write on the session file specified"
@@ -338,12 +344,16 @@ def initTargetEnv():
         if conf.sessionFP:
             conf.sessionFP.close()
 
+        if conf.hashDB:
+            conf.hashDB.close()
+
         if conf.cj:
             conf.cj.clear()
 
         conf.paramDict = {}
         conf.parameters = {}
         conf.sessionFile = None
+        conf.hashDBFile = None
 
         __setKnowledgeBaseAttributes(False)
         __restoreCmdLineOptions()
diff --git a/lib/techniques/error/use.py b/lib/techniques/error/use.py
index f43dbbc2150..0e03465736b 100644
--- a/lib/techniques/error/use.py
+++ b/lib/techniques/error/use.py
@@ -53,83 +53,86 @@ def __oneShotErrorUse(expression, field):
 
     threadData = getCurrentThreadData()
 
-    retVal = None
+    retVal = conf.hashDB.retrieve(expression) if not conf.freshQueries else None
+
     offset = 1
     chunk_length = None
 
-    while True:
-        check = "%s(?P<result>.*?)%s" % (kb.misc.start, kb.misc.stop)
-        trimcheck = "%s(?P<result>.*?)</" % (kb.misc.start)
-
-        nulledCastedField = agent.nullAndCastField(field)
-
-        if Backend.isDbms(DBMS.MYSQL):
-            chunk_length = MYSQL_ERROR_CHUNK_LENGTH
-            nulledCastedField = queries[DBMS.MYSQL].substring.query % (nulledCastedField, offset, chunk_length)
-        elif Backend.isDbms(DBMS.MSSQL):
-            chunk_length = MSSQL_ERROR_CHUNK_LENGTH
-            nulledCastedField = queries[DBMS.MSSQL].substring.query % (nulledCastedField, offset, chunk_length)
-
-        # Forge the error-based SQL injection request
-        vector = kb.injection.data[PAYLOAD.TECHNIQUE.ERROR].vector
-        query = agent.prefixQuery(vector)
-        query = agent.suffixQuery(query)
-        injExpression = expression.replace(field, nulledCastedField, 1)
-        injExpression = unescaper.unescape(injExpression)
-        injExpression = query.replace("[QUERY]", injExpression)
-        payload = agent.payload(newValue=injExpression)
-
-        # Perform the request
-        page, headers = Request.queryPage(payload, content=True)
-
-        reqCount += 1
-
-        # Parse the returned page to get the exact error-based
-        # sql injection output
-        output = reduce(lambda x, y: x if x is not None else y, [ \
-                extractRegexResult(check, page, re.DOTALL | re.IGNORECASE), \
-                extractRegexResult(check, listToStrValue(headers.headers \
-                if headers else None), re.DOTALL | re.IGNORECASE), \
-                extractRegexResult(check, threadData.lastRedirectMsg[1] \
-                if threadData.lastRedirectMsg and threadData.lastRedirectMsg[0] == \
-                threadData.lastRequestUID else None, re.DOTALL | re.IGNORECASE)], \
-                None)
-
-        if output is not None:
-            output = getUnicode(output, kb.pageEncoding)
-        else:
-            trimmed = extractRegexResult(trimcheck, page, re.DOTALL | re.IGNORECASE) \
-                or extractRegexResult(trimcheck, listToStrValue(headers.headers \
-                if headers else None), re.DOTALL | re.IGNORECASE) \
-                or extractRegexResult(trimcheck, threadData.lastRedirectMsg[1] \
-                if threadData.lastRedirectMsg and threadData.lastRedirectMsg[0] == \
-                threadData.lastRequestUID else None, re.DOTALL | re.IGNORECASE)
-
-            if trimmed:
-                warnMsg = "possible server trimmed output detected (due to its length): "
-                warnMsg += trimmed
-                logger.warn(warnMsg)
+    if not retVal:
+        while True:
+            check = "%s(?P<result>.*?)%s" % (kb.misc.start, kb.misc.stop)
+            trimcheck = "%s(?P<result>.*?)</" % (kb.misc.start)
+
+            nulledCastedField = agent.nullAndCastField(field)
+
+            if Backend.isDbms(DBMS.MYSQL):
+                chunk_length = MYSQL_ERROR_CHUNK_LENGTH
+                nulledCastedField = queries[DBMS.MYSQL].substring.query % (nulledCastedField, offset, chunk_length)
+            elif Backend.isDbms(DBMS.MSSQL):
+                chunk_length = MSSQL_ERROR_CHUNK_LENGTH
+                nulledCastedField = queries[DBMS.MSSQL].substring.query % (nulledCastedField, offset, chunk_length)
+
+            # Forge the error-based SQL injection request
+            vector = kb.injection.data[PAYLOAD.TECHNIQUE.ERROR].vector
+            query = agent.prefixQuery(vector)
+            query = agent.suffixQuery(query)
+            injExpression = expression.replace(field, nulledCastedField, 1)
+            injExpression = unescaper.unescape(injExpression)
+            injExpression = query.replace("[QUERY]", injExpression)
+            payload = agent.payload(newValue=injExpression)
+
+            # Perform the request
+            page, headers = Request.queryPage(payload, content=True)
+
+            reqCount += 1
+
+            # Parse the returned page to get the exact error-based
+            # sql injection output
+            output = reduce(lambda x, y: x if x is not None else y, [ \
+                    extractRegexResult(check, page, re.DOTALL | re.IGNORECASE), \
+                    extractRegexResult(check, listToStrValue(headers.headers \
+                    if headers else None), re.DOTALL | re.IGNORECASE), \
+                    extractRegexResult(check, threadData.lastRedirectMsg[1] \
+                    if threadData.lastRedirectMsg and threadData.lastRedirectMsg[0] == \
+                    threadData.lastRequestUID else None, re.DOTALL | re.IGNORECASE)], \
+                    None)
 
-        if any(map(lambda dbms: Backend.isDbms(dbms), [DBMS.MYSQL, DBMS.MSSQL])):
-            if offset == 1:
-                retVal = output
+            if output is not None:
+                output = getUnicode(output, kb.pageEncoding)
             else:
-                retVal += output if output else ''
+                trimmed = extractRegexResult(trimcheck, page, re.DOTALL | re.IGNORECASE) \
+                    or extractRegexResult(trimcheck, listToStrValue(headers.headers \
+                    if headers else None), re.DOTALL | re.IGNORECASE) \
+                    or extractRegexResult(trimcheck, threadData.lastRedirectMsg[1] \
+                    if threadData.lastRedirectMsg and threadData.lastRedirectMsg[0] == \
+                    threadData.lastRequestUID else None, re.DOTALL | re.IGNORECASE)
+
+                if trimmed:
+                    warnMsg = "possible server trimmed output detected (due to its length): "
+                    warnMsg += trimmed
+                    logger.warn(warnMsg)
+
+            if any(map(lambda dbms: Backend.isDbms(dbms), [DBMS.MYSQL, DBMS.MSSQL])):
+                if offset == 1:
+                    retVal = output
+                else:
+                    retVal += output if output else ''
 
-            if output and len(output) >= chunk_length:
-                offset += chunk_length
+                if output and len(output) >= chunk_length:
+                    offset += chunk_length
+                else:
+                    break
             else:
+                retVal = output
                 break
-        else:
-            retVal = output
-            break
 
-    if isinstance(retVal, basestring):
-        retVal = htmlunescape(retVal).replace("<br>", "\n")
+        if isinstance(retVal, basestring):
+            retVal = htmlunescape(retVal).replace("<br>", "\n")
 
-    retVal = __errorReplaceChars(retVal)
+        retVal = __errorReplaceChars(retVal)
 
-    dataToSessionFile("[%s][%s][%s][%s][%s]\n" % (conf.url, kb.injection.place, conf.parameters[kb.injection.place], expression, replaceNewlineTabs(retVal)))
+        #dataToSessionFile("[%s][%s][%s][%s][%s]\n" % (conf.url, kb.injection.place, conf.parameters[kb.injection.place], expression, replaceNewlineTabs(retVal)))
+        conf.hashDB.write(expression, retVal)
 
     return safecharencode(retVal) if kb.safeCharEncode else retVal
 
diff --git a/lib/techniques/union/use.py b/lib/techniques/union/use.py
index 1eb3fba8f9d..6dd72afe0b1 100644
--- a/lib/techniques/union/use.py
+++ b/lib/techniques/union/use.py
@@ -50,54 +50,59 @@
 def __oneShotUnionUse(expression, unpack=True, limited=False):
     global reqCount
 
-    check = "(?P<result>%s.*%s)" % (kb.misc.start, kb.misc.stop)
-    trimcheck = "%s(?P<result>.*?)</" % (kb.misc.start)
-
-    # Prepare expression with delimiters
-    expression = agent.concatQuery(expression, unpack)
-    expression = unescaper.unescape(expression)
-
-    if conf.limitStart or conf.limitStop:
-        where = PAYLOAD.WHERE.NEGATIVE
-    else:
-        where = None
-
-    # Forge the inband SQL injection request
-    vector = kb.injection.data[PAYLOAD.TECHNIQUE.UNION].vector
-    query = agent.forgeInbandQuery(expression, vector[0], vector[1], vector[2], vector[3], vector[4], vector[5], None, limited)
-    payload = agent.payload(newValue=query, where=where)
-
-    # Perform the request
-    page, headers = Request.queryPage(payload, content=True, raise404=False)
-
-    reqCount += 1
-
-    # Parse the returned page to get the exact union-based
-    # sql injection output
-    output = reduce(lambda x, y: x if x is not None else y, [ \
-            extractRegexResult(check, removeReflectiveValues(page, payload), re.DOTALL | re.IGNORECASE), \
-            extractRegexResult(check, removeReflectiveValues(listToStrValue(headers.headers \
-            if headers else None), payload, True), re.DOTALL | re.IGNORECASE)], \
-            None)
-
-    if output is not None:
-        output = getUnicode(output, kb.pageEncoding)
-    else:
-        trimmed = extractRegexResult(trimcheck, removeReflectiveValues(page, payload), re.DOTALL | re.IGNORECASE) \
-                or extractRegexResult(trimcheck, removeReflectiveValues(listToStrValue(headers.headers \
-                if headers else None), payload, True), re.DOTALL | re.IGNORECASE)
-
-        if trimmed:
-            warnMsg = "possible server trimmed output detected (due to its length): "
-            warnMsg += trimmed
-            logger.warn(warnMsg)
-        elif Backend.isDbms(DBMS.MYSQL) and not kb.multiThreadMode:
-            warnMsg = "if the problem persists with 'None' values please try to use "
-            warnMsg += "hidden switch --no-cast (fixing problems with some collation "
-            warnMsg += "issues)"
-            singleTimeWarnMessage(warnMsg)
-
-    return output
+    retVal = conf.hashDB.retrieve(expression) if not conf.freshQueries else None
+
+    if not retVal:
+        check = "(?P<result>%s.*%s)" % (kb.misc.start, kb.misc.stop)
+        trimcheck = "%s(?P<result>.*?)</" % (kb.misc.start)
+
+        # Prepare expression with delimiters
+        expression = agent.concatQuery(expression, unpack)
+        expression = unescaper.unescape(expression)
+
+        if conf.limitStart or conf.limitStop:
+            where = PAYLOAD.WHERE.NEGATIVE
+        else:
+            where = None
+
+        # Forge the inband SQL injection request
+        vector = kb.injection.data[PAYLOAD.TECHNIQUE.UNION].vector
+        query = agent.forgeInbandQuery(expression, vector[0], vector[1], vector[2], vector[3], vector[4], vector[5], None, limited)
+        payload = agent.payload(newValue=query, where=where)
+
+        # Perform the request
+        page, headers = Request.queryPage(payload, content=True, raise404=False)
+
+        reqCount += 1
+
+        # Parse the returned page to get the exact union-based
+        # sql injection output
+        retVal = reduce(lambda x, y: x if x is not None else y, [ \
+                extractRegexResult(check, removeReflectiveValues(page, payload), re.DOTALL | re.IGNORECASE), \
+                extractRegexResult(check, removeReflectiveValues(listToStrValue(headers.headers \
+                if headers else None), payload, True), re.DOTALL | re.IGNORECASE)], \
+                None)
+
+        if retVal is not None:
+            retVal = getUnicode(retVal, kb.pageEncoding)
+        else:
+            trimmed = extractRegexResult(trimcheck, removeReflectiveValues(page, payload), re.DOTALL | re.IGNORECASE) \
+                    or extractRegexResult(trimcheck, removeReflectiveValues(listToStrValue(headers.headers \
+                    if headers else None), payload, True), re.DOTALL | re.IGNORECASE)
+
+            if trimmed:
+                warnMsg = "possible server trimmed output detected (due to its length): "
+                warnMsg += trimmed
+                logger.warn(warnMsg)
+            elif Backend.isDbms(DBMS.MYSQL) and not kb.multiThreadMode:
+                warnMsg = "if the problem persists with 'None' values please try to use "
+                warnMsg += "hidden switch --no-cast (fixing problems with some collation "
+                warnMsg += "issues)"
+                singleTimeWarnMessage(warnMsg)
+
+        conf.hashDB.write(expression, retVal)
+
+    return retVal
 
 def configUnion(char=None, columns=None):
     def __configUnionChar(char):
diff --git a/lib/utils/hashdb.py b/lib/utils/hashdb.py
new file mode 100644
index 00000000000..c72691dcd63
--- /dev/null
+++ b/lib/utils/hashdb.py
@@ -0,0 +1,63 @@
+#!/usr/bin/env python
+
+"""
+$Id$
+
+Copyright (c) 2006-2011 sqlmap developers (http://www.sqlmap.org/)
+See the file 'doc/COPYING' for copying permission
+"""
+
+import hashlib
+import sqlite3
+
+from lib.core.settings import UNICODE_ENCODING
+
+class HashDB:
+    def __init__(self, filepath):
+        self.connection = sqlite3.connect(filepath)
+        self.cursor = self.connection.cursor()
+        self.cursor.execute("CREATE TABLE IF NOT EXISTS storage (id INTEGER PRIMARY KEY, value TEXT)")
+
+    def __del__(self):
+        self.close()
+
+    def close(self):
+        try:
+            self.endTransaction()
+            self.connection.close()
+        except:
+            pass
+
+    def hashKey(self, key):
+        key = key.encode(UNICODE_ENCODING) if isinstance(key, unicode) else repr(key)
+        retVal = int(hashlib.md5(key).hexdigest()[:8], 16)
+        return retVal
+
+    def beginTransaction(self):
+        """
+        Great speed improvement can be gained by using explicit transactions around multiple inserts.
+        Reference: http://stackoverflow.com/questions/4719836/python-and-sqlite3-adding-thousands-of-rows
+        """
+        self.cursor.execute('BEGIN TRANSACTION')
+
+    def endTransaction(self):
+        try:
+            self.cursor.execute('END TRANSACTION')
+        except sqlite3.OperationalError:
+            pass
+
+    def retrieve(self, key):
+        retVal = None
+        if key:
+            hash_ = self.hashKey(key)
+            for row in self.cursor.execute("SELECT value FROM storage WHERE id=?", (hash_,)):
+                retVal = row[0]
+        return retVal
+
+    def write(self, key, value):
+        if key:
+            hash_ = self.hashKey(key)
+            try:
+                self.cursor.execute("INSERT INTO storage VALUES (?, ?)", (hash_, value,))
+            except sqlite3.IntegrityError:
+                self.cursor.execute("UPDATE storage SET value=? WHERE id=?", (value, hash_,))