Extract the creation and Marshal of a Grab object (#256)

This allows these steps to be reusable outside of the context of the
grabTarget function.

https://github.com/zmap/zgrab2/pull/256
This commit is contained in:
David Adrian 2020-03-17 11:50:54 -04:00 committed by GitHub
parent 0bd36c5267
commit 351f826d8a
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -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)
}