From 9bc735963bc3d3640ff299c861d0f2dd485f558d Mon Sep 17 00:00:00 2001
From: Miroslav Stampar <miroslav.stampar@gmail.com>
Date: Sun, 4 Dec 2011 22:42:19 +0000
Subject: [PATCH] update of redirection mechanism (now 3-state - redirected,
 original and "ignored" (containing redirection message itself))

---
 lib/core/enums.py      |  5 +++++
 lib/core/option.py     |  2 +-
 lib/request/connect.py | 23 +++++++++++++----------
 3 files changed, 19 insertions(+), 11 deletions(-)

diff --git a/lib/core/enums.py b/lib/core/enums.py
index 4f5f3b1e0..3c215f747 100644
--- a/lib/core/enums.py
+++ b/lib/core/enums.py
@@ -110,6 +110,11 @@ class EXPECTED:
     BOOL = "bool"
     INT = "int"
 
+class REDIRECTION:
+    FOLLOW = "1"
+    ORIGINAL = "2"
+    IGNORE = "3"
+
 class PAYLOAD:
     SQLINJECTION = {
                      1: "boolean-based blind",
diff --git a/lib/core/option.py b/lib/core/option.py
index 69a42094c..b8c644041 100644
--- a/lib/core/option.py
+++ b/lib/core/option.py
@@ -1383,7 +1383,6 @@ def __setKnowledgeBaseAttributes(flushAll=True):
 
     kb.absFilePaths = set()
     kb.adjustTimeDelay = False
-    kb.alwaysRedirect = None
     kb.alwaysRefresh = None
     kb.arch = None
     kb.authHeader = None
@@ -1455,6 +1454,7 @@ def __setKnowledgeBaseAttributes(flushAll=True):
     kb.processResponseCounter = 0
     kb.proxyAuthHeader = None
     kb.queryCounter = 0
+    kb.redirectChoice = None
     kb.redirectSetCookie = None
     kb.reflectiveMechanism = True
     kb.reflectiveCounters = {REFLECTIVE_COUNTER.MISS:0, REFLECTIVE_COUNTER.HIT:0}
diff --git a/lib/request/connect.py b/lib/request/connect.py
index 584c9d2d4..e50db0e96 100644
--- a/lib/request/connect.py
+++ b/lib/request/connect.py
@@ -49,6 +49,7 @@ from lib.core.enums import HTTPMETHOD
 from lib.core.enums import NULLCONNECTION
 from lib.core.enums import PAYLOAD
 from lib.core.enums import PLACE
+from lib.core.enums import REDIRECTION
 from lib.core.exception import sqlmapConnectionException
 from lib.core.exception import sqlmapSyntaxException
 from lib.core.settings import HTTP_ACCEPT_HEADER_VALUE
@@ -312,21 +313,23 @@ class Connect:
             if hasattr(conn, "redurl") and hasattr(conn, "redcode") and target\
               and not redirecting and not conf.realTest:
 
-                if kb.alwaysRedirect is None:
+                if kb.redirectChoice is None:
                     msg = "sqlmap got a %d redirect to " % conn.redcode
-                    msg += "'%s'. Do you want to follow redirects " % conn.redurl
-                    msg += "from now on (or stay on the original page otherwise)? [Y/n]"
-                    choice = readInput(msg, default="Y")
+                    msg += "'%s'. What do you want to do? " % conn.redurl
+                    msg += "\n[1] Follow the redirection (default)"
+                    msg += "\n[2] Stay on the original page"
+                    msg += "\n[3] Ignore"
+                    choice = readInput(msg, default="1")
 
-                    kb.alwaysRedirect = choice not in ("n", "N")
+                    kb.redirectChoice = choice
 
-                if kb.alwaysRedirect:
-                    kwargs['url'] = conn.redurl
-                    kwargs['redirecting'] = conn.redcode
-                    return Connect.__getPageProxy(**kwargs)
-                else:
+                if kb.redirectChoice == REDIRECTION.IGNORE:
                     redirecting = conn.redcode
                     page = threadData.lastRedirectMsg[1]
+                else:
+                    kwargs['url'] = conf.url if kb.redirectChoice == REDIRECTION.ORIGINAL else conn.redurl
+                    kwargs['redirecting'] = conn.redcode
+                    return Connect.__getPageProxy(**kwargs)
 
             # Return response object
             if response: