From abbd3523922aa0d64dea345807f49db70b6dc6cc Mon Sep 17 00:00:00 2001
From: Miroslav Stampar <miroslav.stampar@gmail.com>
Date: Thu, 23 Oct 2014 14:33:22 +0200
Subject: [PATCH] Support for X-CSRF-TOKEN header (Issue #2)

---
 lib/core/target.py     |  4 ++--
 lib/request/connect.py | 23 ++++++++++++++++-------
 2 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/lib/core/target.py b/lib/core/target.py
index 50a4d0d74..cd542d928 100644
--- a/lib/core/target.py
+++ b/lib/core/target.py
@@ -346,9 +346,9 @@ def _setRequestParams():
         raise SqlmapGenericException(errMsg)
 
     if conf.csrfToken:
-        if not any(conf.csrfToken in _ for _ in (conf.paramDict.get(PLACE.GET, {}), conf.paramDict.get(PLACE.POST, {}))):
+        if not any(conf.csrfToken in _ for _ in (conf.paramDict.get(PLACE.GET, {}), conf.paramDict.get(PLACE.POST, {}))) and not conf.csrfToken in set(_[0].lower() for _ in conf.httpHeaders):
             errMsg = "CSRF protection token parameter '%s' not " % conf.csrfToken
-            errMsg += "found in provided GET and/or POST values"
+            errMsg += "found in provided GET, POST or header values"
             raise SqlmapGenericException(errMsg)
     else:
         for place in (PLACE.GET, PLACE.POST):
diff --git a/lib/request/connect.py b/lib/request/connect.py
index 157e0cce2..29550c41c 100644
--- a/lib/request/connect.py
+++ b/lib/request/connect.py
@@ -758,16 +758,21 @@ class Connect(object):
                     retVal = re.sub("%s=[^&]*" % parameter, "%s=%s" % (parameter, newValue), paramString)
                 return retVal
 
-            page, _, _ = Connect.getPage(url=conf.csrfUrl or conf.url, cookie=conf.parameters.get(PLACE.COOKIE), direct=True, silent=True, ua=conf.parameters.get(PLACE.USER_AGENT), referer=conf.parameters.get(PLACE.REFERER), host=conf.parameters.get(PLACE.HOST))
-            match = re.search(r"<input[^>]+name=[\"']?%s[\"']?\s[^>]*value=(\"([^\"]+)|'([^']+)|([^ >]+))" % conf.csrfToken, page)
+            page, headers, code = Connect.getPage(url=conf.csrfUrl or conf.url, cookie=conf.parameters.get(PLACE.COOKIE), direct=True, silent=True, ua=conf.parameters.get(PLACE.USER_AGENT), referer=conf.parameters.get(PLACE.REFERER), host=conf.parameters.get(PLACE.HOST))
+            match = re.search(r"<input[^>]+name=[\"']?%s[\"']?\s[^>]*value=(\"([^\"]+)|'([^']+)|([^ >]+))" % conf.csrfToken, page or "")
             token = (match.group(2) or match.group(3) or match.group(4)) if match else None
 
             if not token:
-                errMsg = "CSRF token value '%s' can't be found at '%s'" % (conf.csrfToken, conf.csrfUrl or conf.url)
-                if not conf.csrfUrl:
-                    errMsg += ". You can try to rerun by providing "
-                    errMsg += "a valid value for option '--csrf-url'"
-                raise SqlmapTokenException, errMsg
+                if conf.csrfUrl != conf.url and code == httplib.OK:
+                    if headers and "text/plain" in headers.get(HTTP_HEADER.CONTENT_TYPE, ""):
+                        token = page
+
+                if not token:
+                    errMsg = "CSRF token value '%s' can't be found at '%s'" % (conf.csrfToken, conf.csrfUrl or conf.url)
+                    if not conf.csrfUrl:
+                        errMsg += ". You can try to rerun by providing "
+                        errMsg += "a valid value for option '--csrf-url'"
+                    raise SqlmapTokenException, errMsg
 
             if token:
                 for item in (PLACE.GET, PLACE.POST):
@@ -777,6 +782,10 @@ class Connect(object):
                         elif item == PLACE.POST and post:
                             post = _adjustParameter(post, conf.csrfToken, token)
 
+                for i in xrange(len(conf.httpHeaders)):
+                    if conf.httpHeaders[i][0].lower() == conf.csrfToken.lower():
+                        conf.httpHeaders[i] = (conf.httpHeaders[i][0], token)
+
         if conf.rParam:
             def _randomizeParameter(paramString, randomParameter):
                 retVal = paramString