Update for an Issue #342 and #372

stamparm · stamparm · commit 2420a4b626e8 · 2013-01-31T10:01:52.000+01:00
diff --git a/lib/core/agent.py b/lib/core/agent.py
@@ -888,23 +888,17 @@ def forgeQueryOutputLength(self, expression):
         lengthQuery = queries[Backend.getIdentifiedDbms()].length.query
         select = re.search("\ASELECT\s+", expression, re.I)
         selectTopExpr = re.search("\ASELECT\s+TOP\s+[\d]+\s+(.+?)\s+FROM", expression, re.I)
-        selectDistinctExpr = re.search("\ASELECT\s+DISTINCT\((.+?)\)\s+FROM", expression, re.I)
         selectFromExpr = re.search("\ASELECT\s+(.+?)\s+FROM", expression, re.I)
         selectExpr = re.search("\ASELECT\s+(.+)$", expression, re.I)
 
         _, _, _, _, _, _, fieldsStr, _ = self.getFields(expression)
 
-        if any((selectTopExpr, selectDistinctExpr, selectFromExpr, selectExpr)):
+        if any((selectTopExpr, selectFromExpr, selectExpr)):
             query = fieldsStr
         else:
             query = expression
 
-        if selectDistinctExpr:
-            lengthExpr = "SELECT %s FROM (%s)" % (lengthQuery % query, expression)
-
-            if Backend.getIdentifiedDbms() in (DBMS.MYSQL, DBMS.PGSQL):
-                lengthExpr += " AS %s" % randomStr(lowercase=True)
-        elif select:
+        if select:
             lengthExpr = expression.replace(query, lengthQuery % query, 1)
         else:
             lengthExpr = lengthQuery % expression
diff --git a/lib/request/inject.py b/lib/request/inject.py
@@ -25,6 +25,7 @@
 from lib.core.common import parseUnionPage
 from lib.core.common import popValue
 from lib.core.common import pushValue
+from lib.core.common import randomStr
 from lib.core.common import readInput
 from lib.core.common import singleTimeWarnMessage
 from lib.core.data import conf
@@ -76,6 +77,13 @@ def _goInference(payload, expression, charsetType=None, firstChar=None, lastChar
 
     if not (timeBasedCompare and kb.dnsTest):
         if (conf.eta or conf.threads > 1) and Backend.getIdentifiedDbms() and not re.search("(COUNT|LTRIM)\(", expression, re.I) and not timeBasedCompare:
+
+            if field and re.search("\ASELECT\s+DISTINCT\((.+?)\)\s+FROM", expression, re.I):
+                expression = "SELECT %s FROM (%s)" % (field, expression)
+
+                if Backend.getIdentifiedDbms() in (DBMS.MYSQL, DBMS.PGSQL):
+                    expression += " AS %s" % randomStr(lowercase=True)
+
             if field and conf.hexConvert:
                 nulledCastedField = agent.nullAndCastField(field)
                 injExpression = expression.replace(field, nulledCastedField, 1)
