From bb277c70b0f0904c48525c0ab02bfa2abf2d4ef7 Mon Sep 17 00:00:00 2001
From: Mzack9999 <mzack9999@protonmail.com>
Date: Mon, 23 Nov 2020 21:37:52 +0100
Subject: [PATCH] adding metadata for exact collab match

---
 v2/go.mod                           |  2 +-
 v2/go.sum                           |  2 ++
 v2/pkg/collaborator/collaborator.go | 24 +++++++++++++++---------
 v2/pkg/collaborator/util.go         |  9 +++++++++
 v2/pkg/executer/executer_http.go    |  5 +++++
 5 files changed, 32 insertions(+), 10 deletions(-)
 create mode 100644 v2/pkg/collaborator/util.go

diff --git a/v2/go.mod b/v2/go.mod
index dfcfed3b..0884574b 100644
--- a/v2/go.mod
+++ b/v2/go.mod
@@ -14,7 +14,7 @@ require (
 	github.com/miekg/dns v1.1.35
 	github.com/pkg/errors v0.9.1
 	github.com/projectdiscovery/clistats v0.0.5
-	github.com/projectdiscovery/collaborator v0.0.1
+	github.com/projectdiscovery/collaborator v0.0.2-0.20201122173402-7afbb52febed
 	github.com/projectdiscovery/fastdialer v0.0.1
 	github.com/projectdiscovery/gologger v1.0.1
 	github.com/projectdiscovery/hmap v0.0.1
diff --git a/v2/go.sum b/v2/go.sum
index 57159baf..db688e69 100644
--- a/v2/go.sum
+++ b/v2/go.sum
@@ -57,6 +57,8 @@ github.com/projectdiscovery/clistats v0.0.5 h1:vcvOR9PrFRawO/7FWD6pER9nYVSoSTD2F
 github.com/projectdiscovery/clistats v0.0.5/go.mod h1:lV6jUHAv2bYWqrQstqW8iVIydKJhWlVaLl3Xo9ioVGg=
 github.com/projectdiscovery/collaborator v0.0.1 h1:dbQ5BCL/a3c+BB9cGtrGgiLs23+EfSzoaTzX/pxqiTI=
 github.com/projectdiscovery/collaborator v0.0.1/go.mod h1:J1z0fC7Svutz3LJqoRyTHA3F0Suh4livmkYv8MnKw20=
+github.com/projectdiscovery/collaborator v0.0.2-0.20201122173402-7afbb52febed h1:0B4ccP68IL6MNHz99dAFeh0E3AY7fU5wVYwZJVX8ZG4=
+github.com/projectdiscovery/collaborator v0.0.2-0.20201122173402-7afbb52febed/go.mod h1:J1z0fC7Svutz3LJqoRyTHA3F0Suh4livmkYv8MnKw20=
 github.com/projectdiscovery/fastdialer v0.0.1 h1:MgBkJ/zkciFu/PcbAz0DYGiZn2aqv6b39NvfXxfN8qg=
 github.com/projectdiscovery/fastdialer v0.0.1/go.mod h1:d24GUzSb93wOY7lu4gJmXAzfomqAGEcRrInEVrM6zbc=
 github.com/projectdiscovery/gologger v1.0.1 h1:FzoYQZnxz9DCvSi/eg5A6+ET4CQ0CDUs27l6Exr8zMQ=
diff --git a/v2/pkg/collaborator/collaborator.go b/v2/pkg/collaborator/collaborator.go
index f4d34c9e..cfa43440 100644
--- a/v2/pkg/collaborator/collaborator.go
+++ b/v2/pkg/collaborator/collaborator.go
@@ -2,6 +2,7 @@ package collaborator
 
 import (
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/projectdiscovery/collaborator"
@@ -17,6 +18,7 @@ var DefaultPollInterval time.Duration = time.Second * time.Duration(PollSeconds)
 var DefaultCollaborator BurpCollaborator = BurpCollaborator{Collab: collaborator.NewBurpCollaborator()}
 
 type BurpCollaborator struct {
+	sync.RWMutex
 	options *Options // unused
 	Collab  *collaborator.BurpCollaborator
 }
@@ -41,19 +43,23 @@ func (b *BurpCollaborator) Poll() {
 	}
 }
 
-func (b *BurpCollaborator) Has(s string) bool {
+func (b *BurpCollaborator) Has(s string) (found bool) {
+	foundAt := 0
 	for _, r := range b.Collab.RespBuffer {
 		for i := 0; i < len(r.Responses); i++ {
-			// search in dns
-			if strings.Contains(r.Responses[i].Data.RawRequestDecoded, s) {
-				return true
-			}
-			// search in http
-			if strings.Contains(r.Responses[i].Data.RequestDecoded, s) {
-				return true
+			// search in dns - http - smtp
+			b.RLock()
+			found = strings.Contains(r.Responses[i].Data.RawRequestDecoded, s) || strings.Contains(r.Responses[i].Data.RequestDecoded, s) || strings.Contains(r.Responses[i].Data.MessageDecoded, s)
+			b.RUnlock()
+			if found {
+				b.Lock()
+				r.Responses = removeMatch(r.Responses, foundAt)
+				b.Unlock()
+				break
 			}
+
 		}
 	}
 
-	return false
+	return
 }
diff --git a/v2/pkg/collaborator/util.go b/v2/pkg/collaborator/util.go
new file mode 100644
index 00000000..a6e35675
--- /dev/null
+++ b/v2/pkg/collaborator/util.go
@@ -0,0 +1,9 @@
+package collaborator
+
+import (
+	"github.com/projectdiscovery/collaborator"
+)
+
+func removeMatch(responses []collaborator.BurpResponse, index int) []collaborator.BurpResponse {
+	return append(responses[:index], responses[index+1:]...)
+}
diff --git a/v2/pkg/executer/executer_http.go b/v2/pkg/executer/executer_http.go
index 7bfbccfd..d2ed5d14 100644
--- a/v2/pkg/executer/executer_http.go
+++ b/v2/pkg/executer/executer_http.go
@@ -530,6 +530,11 @@ func (e *HTTPExecuter) handleHTTP(reqURL string, request *requests.HTTPRequest,
 	// hardcode stopping storing data after defaultMaxHistorydata items
 	if len(result.historyData) < defaultMaxHistorydata {
 		result.Lock()
+		// update history data with current reqURL and hostname
+		result.historyData["reqURL"] = reqURL
+		if parsed, err := url.Parse(reqURL); err == nil {
+			result.historyData["Hostname"] = parsed.Host
+		}
 		result.historyData = generators.MergeMaps(result.historyData, matchers.HTTPToMap(resp, body, headers, duration, format))
 		result.Unlock()
 	}