From e53fef546e8457df5af377f317b7613ea812b64e Mon Sep 17 00:00:00 2001
From: Miroslav Stampar <miroslav.stampar@gmail.com>
Date: Tue, 7 Dec 2010 14:35:31 +0000
Subject: [PATCH] update regarding session page templates

---
 lib/controller/checks.py | 10 +++++++---
 lib/request/inject.py    |  2 +-
 lib/request/templates.py |  9 +++++----
 3 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/lib/controller/checks.py b/lib/controller/checks.py
index df722d8c4..794dc3361 100644
--- a/lib/controller/checks.py
+++ b/lib/controller/checks.py
@@ -47,6 +47,7 @@ from lib.core.session import setString
 from lib.core.session import setRegexp
 from lib.core.settings import TIME_MIN_DELTA
 from lib.request.connect import Connect as Request
+from lib.request.templates import getPageTemplate
 from plugins.dbms.firebird.syntax import Syntax as Firebird
 from plugins.dbms.postgresql.syntax import Syntax as PostgreSQL
 from plugins.dbms.mssqlserver.syntax import Syntax as MSSQLServer
@@ -258,20 +259,22 @@ def checkSqlInjection(place, parameter, value):
 
             # For each test's <where>
             for where in test.where:
+                templatePayload = None
+
                 # Threat the parameter original value according to the
                 # test's <where> tag
                 if where == 1:
                     origValue = value
-                    kb.pageTemplate = kb.originalPage
                 elif where == 2:
                     origValue = "-%s" % randomInt()
                     # Use different page template than the original one
                     # as we are changing parameters value, which will result
                     # most definitely with a different content
-                    kb.pageTemplate, _ = Request.queryPage(agent.payload(place, parameter, value, origValue), place, content=True)
+                    templatePayload = agent.payload(place, parameter, value, origValue)
                 elif where == 3:
                     origValue = ""
-                    kb.pageTemplate = kb.originalPage
+
+                kb.pageTemplate = getPageTemplate(templatePayload, place)
 
                 # Forge request payload by prepending with boundary's
                 # prefix and appending the boundary's suffix to the
@@ -396,6 +399,7 @@ def checkSqlInjection(place, parameter, value):
                     injection.data[stype].where = where
                     injection.data[stype].vector = vector
                     injection.data[stype].comment = comment
+                    injection.data[stype].templatePayload = templatePayload
 
                     if hasattr(test, "details"):
                         for detailKey, detailValue in test.details.items():
diff --git a/lib/request/inject.py b/lib/request/inject.py
index 93f80850d..cd29890c9 100644
--- a/lib/request/inject.py
+++ b/lib/request/inject.py
@@ -101,7 +101,7 @@ def __goInferenceProxy(expression, fromUser=False, expected=None, batch=False, r
 
     if kb.injection.data[1].vector is not None:
         vector = agent.cleanupPayload(kb.injection.data[1].vector)
-        kb.pageTemplate = getPageTemplate(kb.injection.data[1].payload, kb.injection.place)
+        kb.pageTemplate = getPageTemplate(kb.injection.data[1].templatePayload, kb.injection.place)
     else:
         vector = queries[kb.misc.testedDbms].inference.query
         kb.pageTemplate = kb.originalPage
diff --git a/lib/request/templates.py b/lib/request/templates.py
index 382de18ea..0049173bf 100644
--- a/lib/request/templates.py
+++ b/lib/request/templates.py
@@ -11,8 +11,9 @@ from lib.core.data import kb
 from lib.request.connect import Connect as Request
 
 def getPageTemplate(payload, place):
-    retVal = None
-    if (payload, place) not in kb.pageTemplates:
-        kb.pageTemplates[(payload, place)], _ = Request.queryPage(payload, place, content=True)
-    retVal = kb.pageTemplates[(payload, place)]
+    retVal = kb.originalPage
+    if payload and place:
+        if (payload, place) not in kb.pageTemplates:
+            kb.pageTemplates[(payload, place)], _ = Request.queryPage(payload, place, content=True)
+        retVal = kb.pageTemplates[(payload, place)]
     return retVal