From 2420a4b626e80ea115c5f8b8dc7bf4877fb1e6e8 Mon Sep 17 00:00:00 2001
From: Miroslav Stampar <miroslav.stampar@gmail.com>
Date: Thu, 31 Jan 2013 10:01:52 +0100
Subject: [PATCH] Update for an Issue #342 and #372

---
 lib/core/agent.py     | 10 ++--------
 lib/request/inject.py |  8 ++++++++
 2 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/lib/core/agent.py b/lib/core/agent.py
index fe7bc335c..9cd727eca 100644
--- a/lib/core/agent.py
+++ b/lib/core/agent.py
@@ -888,23 +888,17 @@ class Agent(object):
         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
index 962e633da..944aaa40e 100644
--- a/lib/request/inject.py
+++ b/lib/request/inject.py
@@ -25,6 +25,7 @@ from lib.core.common import isTechniqueAvailable
 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)