From c16b74ce1a71e9f307e433abb95b7b8c3bb21154 Mon Sep 17 00:00:00 2001
From: Miroslav Stampar <miroslav.stampar@gmail.com>
Date: Fri, 15 Apr 2011 14:21:13 +0000
Subject: [PATCH] covering __pivotDumpTable for keyboard and connection
 exceptions too

---
 plugins/generic/enumeration.py | 70 ++++++++++++++++++++--------------
 1 file changed, 41 insertions(+), 29 deletions(-)

diff --git a/plugins/generic/enumeration.py b/plugins/generic/enumeration.py
index f57819231..55964aceb 100644
--- a/plugins/generic/enumeration.py
+++ b/plugins/generic/enumeration.py
@@ -1184,40 +1184,52 @@ class Enumeration:
         pivotValue = " "
         breakRetrieval = False
 
-        for i in xrange(int(count)):
-            if breakRetrieval:
-                break
+        try:
+            for i in xrange(int(count)):
+                if breakRetrieval:
+                    break
 
-            for column in colList:
-                if column not in lengths:
-                    lengths[column] = 0
+                for column in colList:
+                    if column not in lengths:
+                        lengths[column] = 0
 
-                if column not in entries:
-                    entries[column] = []
+                    if column not in entries:
+                        entries[column] = []
 
-                if column == colList[0]:
-                    # Correction for pivotValues with unrecognized chars
-                    if pivotValue and '?' in pivotValue and pivotValue[0] != '?':
-                        pivotValue = pivotValue.split('?')[0]
-                        pivotValue = pivotValue[:-1] + chr(ord(pivotValue[-1]) + 1)
-                    query = dumpNode.query % (column, table, column, pivotValue)
-                else:
-                    query = dumpNode.query2 % (column, table, colList[0], pivotValue)
-
-                if blind:
-                    value = inject.getValue(query, inband=False, error=False)
-                else:
-                    value = inject.getValue(query, blind=False)
-
-                if column == colList[0]:
-                    if not value:
-                        breakRetrieval = True
-                        break
+                    if column == colList[0]:
+                        # Correction for pivotValues with unrecognized chars
+                        if pivotValue and '?' in pivotValue and pivotValue[0] != '?':
+                            pivotValue = pivotValue.split('?')[0]
+                            pivotValue = pivotValue[:-1] + chr(ord(pivotValue[-1]) + 1)
+                        query = dumpNode.query % (column, table, column, pivotValue)
                     else:
-                        pivotValue = safechardecode(value)
+                        query = dumpNode.query2 % (column, table, colList[0], pivotValue)
 
-                lengths[column] = max(lengths[column], len(value) if value else 0)
-                entries[column].append(value)
+                    if blind:
+                        value = inject.getValue(query, inband=False, error=False)
+                    else:
+                        value = inject.getValue(query, blind=False)
+
+                    if column == colList[0]:
+                        if not value:
+                            breakRetrieval = True
+                            break
+                        else:
+                            pivotValue = safechardecode(value)
+
+                    lengths[column] = max(lengths[column], len(value) if value else 0)
+                    entries[column].append(value)
+
+        except KeyboardInterrupt:
+            warnMsg = "user aborted during enumeration. sqlmap "
+            warnMsg += "will display partial output"
+            logger.warn(warnMsg)
+
+        except sqlmapConnectionException, e:
+            errMsg = "connection exception detected. sqlmap "
+            errMsg += "will display partial output"
+            errMsg += "'%s'" % e
+            logger.critical(errMsg)
 
         return entries, lengths