Fixes sqlmapproject#1559

aizaimenghuangu · Nov 23, 2015 · b2dc443 · b2dc443
1 parent 4d57692
commit b2dc443
Showing 1 changed file with 29 additions and 24 deletions.
diff --git a/plugins/generic/custom.py b/plugins/generic/custom.py
@@ -17,6 +17,7 @@
 from lib.core.data import logger
 from lib.core.dicts import SQL_STATEMENTS
 from lib.core.enums import AUTOCOMPLETE_TYPE
+from lib.core.exception import SqlmapNoneDataException
 from lib.core.settings import NULL
 from lib.core.settings import PARAMETER_SPLITTING_REGEX
 from lib.core.shell import autoCompletion
@@ -35,38 +36,42 @@ def sqlQuery(self, query):
         sqlType = None
         query = query.rstrip(';')
 
-        for sqlTitle, sqlStatements in SQL_STATEMENTS.items():
-            for sqlStatement in sqlStatements:
-                if query.lower().startswith(sqlStatement):
-                    sqlType = sqlTitle
-                    break
+        try:
+            for sqlTitle, sqlStatements in SQL_STATEMENTS.items():
+                for sqlStatement in sqlStatements:
+                    if query.lower().startswith(sqlStatement):
+                        sqlType = sqlTitle
+                        break
 
-        if not any(_ in query.upper() for _ in ("OPENROWSET", "INTO")) and (not sqlType or "SELECT" in sqlType):
-            infoMsg = "fetching %s query output: '%s'" % (sqlType if sqlType is not None else "SQL", query)
-            logger.info(infoMsg)
+            if not any(_ in query.upper() for _ in ("OPENROWSET", "INTO")) and (not sqlType or "SELECT" in sqlType):
+                infoMsg = "fetching %s query output: '%s'" % (sqlType if sqlType is not None else "SQL", query)
+                logger.info(infoMsg)
 
-            output = inject.getValue(query, fromUser=True)
+                output = inject.getValue(query, fromUser=True)
 
-            return output
-        elif not isStackingAvailable() and not conf.direct:
-                warnMsg = "execution of custom SQL queries is only "
-                warnMsg += "available when stacked queries are supported"
-                logger.warn(warnMsg)
+                return output
+            elif not isStackingAvailable() and not conf.direct:
+                    warnMsg = "execution of custom SQL queries is only "
+                    warnMsg += "available when stacked queries are supported"
+                    logger.warn(warnMsg)
 
-                return None
-        else:
-            if sqlType:
-                debugMsg = "executing %s query: '%s'" % (sqlType if sqlType is not None else "SQL", query)
+                    return None
             else:
-                debugMsg = "executing unknown SQL type query: '%s'" % query
-            logger.debug(debugMsg)
+                if sqlType:
+                    debugMsg = "executing %s query: '%s'" % (sqlType if sqlType is not None else "SQL", query)
+                else:
+                    debugMsg = "executing unknown SQL type query: '%s'" % query
+                logger.debug(debugMsg)
 
-            inject.goStacked(query)
+                inject.goStacked(query)
 
-            debugMsg = "done"
-            logger.debug(debugMsg)
+                debugMsg = "done"
+                logger.debug(debugMsg)
 
-            output = NULL
+                output = NULL
+
+        except SqlmapNoneDataException, ex:
+            logger.warn(ex)
 
         return output