From e38267a61e631429725f79aa6dfa79c131899a1e Mon Sep 17 00:00:00 2001
From: Louis-Philippe Huberdeau <lp@delvelabs.ca>
Date: Tue, 18 Jul 2017 15:46:52 -0400
Subject: [PATCH] Include tracking properties in the HAR to identify which test
 the requests were associated to

---
 lib/controller/checks.py |  7 +++++++
 lib/core/threads.py      |  2 --
 lib/utils/har.py         | 22 +++++++++++++++++-----
 3 files changed, 24 insertions(+), 7 deletions(-)

diff --git a/lib/controller/checks.py b/lib/controller/checks.py
index 3e4698c5c..14a17f96b 100644
--- a/lib/controller/checks.py
+++ b/lib/controller/checks.py
@@ -162,6 +162,13 @@ def checkSqlInjection(place, parameter, value):
             unionExtended = False
             trueCode, falseCode = None, None
 
+            if conf.httpCollector is not None:
+                conf.httpCollector.setExtendedArguments({
+                    "_title": title,
+                    "_place": place,
+                    "_parameter": parameter,
+                })
+
             if stype == PAYLOAD.TECHNIQUE.UNION:
                 configUnion(test.request.char)
 
diff --git a/lib/core/threads.py b/lib/core/threads.py
index b3566b955..8f89fb1b8 100644
--- a/lib/core/threads.py
+++ b/lib/core/threads.py
@@ -38,8 +38,6 @@ class _ThreadData(threading.local):
         Resets thread data model
         """
 
-        self.requestCollector = None
-
         self.disableStdOut = False
         self.hashDBCursor = None
         self.inTransaction = False
diff --git a/lib/utils/har.py b/lib/utils/har.py
index 3d3579053..2efb33c80 100644
--- a/lib/utils/har.py
+++ b/lib/utils/har.py
@@ -29,9 +29,15 @@ class HTTPCollectorFactory:
 class HTTPCollector:
     def __init__(self):
         self.messages = BigArray()
+        self.extendedArguments = {}
+
+    def setExtendedArguments(self, arguments):
+        self.extendedArguments = arguments
 
     def collectRequest(self, requestMessage, responseMessage, startTime=None, endTime=None):
-        self.messages.append(RawPair(requestMessage, responseMessage, startTime, endTime))
+        self.messages.append(RawPair(requestMessage, responseMessage,
+                                     startTime=startTime, endTime=endTime,
+                                     extendedArguments=self.extendedArguments))
 
     def obtain(self):
         return {"log": {
@@ -41,24 +47,28 @@ class HTTPCollector:
         }}
 
 class RawPair:
-    def __init__(self, request, response, startTime=None, endTime=None):
+    def __init__(self, request, response, startTime=None, endTime=None, extendedArguments=None):
         self.request = request
         self.response = response
         self.startTime = startTime
         self.endTime = endTime
+        self.extendedArguments = extendedArguments or {}
 
     def toEntry(self):
-        return Entry(request=Request.parse(self.request), response=Response.parse(self.response), startTime=self.startTime, endTime=self.endTime)
+        return Entry(request=Request.parse(self.request), response=Response.parse(self.response),
+                     startTime=self.startTime, endTime=self.endTime,
+                     extendedArguments=self.extendedArguments)
 
 class Entry:
-    def __init__(self, request, response, startTime, endTime):
+    def __init__(self, request, response, startTime, endTime, extendedArguments):
         self.request = request
         self.response = response
         self.startTime = startTime or 0
         self.endTime = endTime or 0
+        self.extendedArguments = extendedArguments
 
     def toDict(self):
-        return {
+        out = {
             "request": self.request.toDict(),
             "response": self.response.toDict(),
             "cache": {},
@@ -66,6 +76,8 @@ class Entry:
             "time": int(1000 * (self.endTime - self.startTime)),
             "startedDateTime": "%s%s" % (datetime.datetime.fromtimestamp(self.startTime).isoformat(), time.strftime("%z")) if self.startTime else None
         }
+        out.update(self.extendedArguments)
+        return out
 
 class Request:
     def __init__(self, method, path, httpVersion, headers, postBody=None, raw=None, comment=None):