From 613d26266a5021239cb38bfb0b45642b62d3e445 Mon Sep 17 00:00:00 2001
From: David Adrian <davidcadrian@gmail.com>
Date: Wed, 19 Feb 2020 14:18:06 -0500
Subject: [PATCH] Extract the creation and Marshal of a Grab object

This allows these steps to be reusable outside of the context of the
grabTarget function.
---
 processing.go | 63 ++++++++++++++++++++++++++++++---------------------
 1 file changed, 37 insertions(+), 26 deletions(-)

diff --git a/processing.go b/processing.go
index a61c2356..7ac14234 100644
--- a/processing.go
+++ b/processing.go
@@ -113,6 +113,41 @@ func (target *ScanTarget) OpenUDP(flags *BaseFlags, udp *UDPFlags) (net.Conn, er
 	return NewTimeoutConnection(nil, conn, flags.Timeout, 0, 0, flags.BytesReadLimit), nil
 }
 
+// BuildGrabFromInputResponse constructs a Grab object for a target, given the
+// scan responses.
+func BuildGrabFromInputResponse(t *ScanTarget, responses map[string]ScanResponse) *Grab {
+	var ipstr string
+
+	if t.IP != nil {
+		ipstr = t.IP.String()
+	}
+	return &Grab{
+		IP:     ipstr,
+		Domain: t.Domain,
+		Data:   responses,
+	}
+}
+
+// EncodeGrab serializes a Grab to JSON, handling the debug fields if necessary.
+func EncodeGrab(raw *Grab, includeDebug bool) ([]byte, error) {
+	var outputData interface{}
+	if includeDebug {
+		outputData = raw
+	} else {
+		// If the caller doesn't explicitly request debug data, strip it out.
+		// TODO: Migrate this to the ZMap fork of sheriff, once it's more
+		// stable.
+		processor := output.Processor{Verbose: false}
+		stripped, err := processor.Process(raw)
+		if err != nil {
+			log.Debugf("Error processing results: %v", err)
+			stripped = raw
+		}
+		outputData = stripped
+	}
+	return json.Marshal(outputData)
+}
+
 // grabTarget calls handler for each action
 func grabTarget(input ScanTarget, m *Monitor) []byte {
 	moduleResult := make(map[string]ScanResponse)
@@ -140,32 +175,8 @@ func grabTarget(input ScanTarget, m *Monitor) []byte {
 		}
 	}
 
-	var ipstr string
-	if input.IP == nil {
-		ipstr = ""
-	} else {
-		s := input.IP.String()
-		ipstr = s
-	}
-
-	raw := Grab{IP: ipstr, Domain: input.Domain, Data: moduleResult}
-
-	var outputData interface{} = raw
-
-	if !includeDebugOutput() {
-		// If the caller doesn't explicitly request debug data, strip it out.
-		// Take advantage of the fact that we can skip the (expensive) call to
-		// process if debug output is included (TODO: until Process does anything else)
-		processor := output.Processor{Verbose: false}
-		stripped, err := processor.Process(raw)
-		if err != nil {
-			log.Debugf("Error processing results: %v", err)
-			stripped = raw
-		}
-		outputData = stripped
-	}
-
-	result, err := json.Marshal(outputData)
+	raw := BuildGrabFromInputResponse(&input, moduleResult)
+	result, err := EncodeGrab(raw, includeDebugOutput())
 	if err != nil {
 		log.Fatalf("unable to marshal data: %s", err)
 	}