diff --git a/processing.go b/processing.go index a61c235..7ac1423 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) }