From f164ebcd88000e363a6acb0f972f8b867b08302a Mon Sep 17 00:00:00 2001
From: Kashif Khan <70996046+kashifkhan0771@users.noreply.github.com>
Date: Mon, 28 Oct 2024 23:09:14 +0500
Subject: [PATCH] Handle custom detector response and include in extra data
 (#3411)

* Handle custom detector response and include in extra data

* Added todo

* fixed panic

* simplicity is always good

* limit the response to 200 chars

* results should print now in output
---
 pkg/custom_detectors/custom_detectors.go | 34 ++++++++++++++++++------
 pkg/detectors/meraki/meraki.go           |  3 ++-
 2 files changed, 28 insertions(+), 9 deletions(-)

diff --git a/pkg/custom_detectors/custom_detectors.go b/pkg/custom_detectors/custom_detectors.go
index 602e7f89bd52..7579324dc583 100644
--- a/pkg/custom_detectors/custom_detectors.go
+++ b/pkg/custom_detectors/custom_detectors.go
@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"context"
 	"encoding/json"
+	"io"
 	"net/http"
 	"regexp"
 	"strings"
@@ -101,10 +102,6 @@ func (c *CustomRegexWebhook) FromData(ctx context.Context, verify bool, data []b
 	close(resultsCh)
 
 	for result := range resultsCh {
-		// NOTE: I don't believe this is being set anywhere else, hence the map assignment.
-		result.ExtraData = map[string]string{
-			"name": c.GetName(),
-		}
 		results = append(results, result)
 	}
 
@@ -129,6 +126,7 @@ func (c *CustomRegexWebhook) createResults(ctx context.Context, match map[string
 		DetectorType: detectorspb.DetectorType_CustomRegex,
 		DetectorName: c.GetName(),
 		Raw:          []byte(raw),
+		ExtraData:    map[string]string{},
 	}
 
 	if !verify {
@@ -166,14 +164,34 @@ func (c *CustomRegexWebhook) createResults(ctx context.Context, match map[string
 			}
 			req.Header.Add(key, strings.TrimLeft(value, "\t\n\v\f\r "))
 		}
-		res, err := httpClient.Do(req)
+		resp, err := httpClient.Do(req)
 		if err != nil {
 			continue
 		}
-		// TODO: Read response body.
-		res.Body.Close()
-		if res.StatusCode == http.StatusOK {
+		defer func() {
+			_, _ = io.Copy(io.Discard, resp.Body)
+			_ = resp.Body.Close()
+		}()
+
+		if resp.StatusCode == http.StatusOK {
+			// mark the result as verified
 			result.Verified = true
+
+			body, err := io.ReadAll(resp.Body)
+			if err != nil {
+				continue
+			}
+
+			// TODO: handle different content-type responses seperatly when implement custom detector configurations
+			responseStr := string(body)
+			// truncate to 200 characters if response length exceeds 200
+			if len(responseStr) > 200 {
+				responseStr = responseStr[:200]
+			}
+
+			// store the processed response in ExtraData
+			result.ExtraData["response"] = responseStr
+
 			break
 		}
 	}
diff --git a/pkg/detectors/meraki/meraki.go b/pkg/detectors/meraki/meraki.go
index 9aed559f4861..e326e84ced34 100644
--- a/pkg/detectors/meraki/meraki.go
+++ b/pkg/detectors/meraki/meraki.go
@@ -4,10 +4,11 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
-	regexp "github.com/wasilibs/go-re2"
 	"io"
 	"net/http"
 
+	regexp "github.com/wasilibs/go-re2"
+
 	"github.com/trufflesecurity/trufflehog/v3/pkg/common"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/detectors"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/pb/detectorspb"