Skip to content

Commit

Permalink
Extract the creation and Marshal of a Grab object (#256)
Browse files Browse the repository at this point in the history
This allows these steps to be reusable outside of the context of the
grabTarget function.

#256
  • Loading branch information
dadrian authored Mar 17, 2020
1 parent 0bd36c5 commit 351f826
Showing 1 changed file with 37 additions and 26 deletions.
63 changes: 37 additions & 26 deletions processing.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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)
}
Expand Down

0 comments on commit 351f826

Please sign in to comment.