Remove dependency on config in output.go (#250)

Refactor the output handlers to take the necessary writers as arguments.

https://github.com/zmap/zgrab2/pull/250
This commit is contained in:
David Adrian 2020-02-18 15:46:36 -05:00 committed by GitHub
parent ef33737c05
commit b2bf9cbc41
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 17 additions and 8 deletions

@ -89,7 +89,8 @@ func validateFrameworkConfiguration() {
log.Fatal(err) log.Fatal(err)
} }
} }
SetOutputFunc(OutputResultsFile) outputFunc := OutputResultsWriterFunc(config.outputFile)
SetOutputFunc(outputFunc)
if config.MetaFileName == "-" { if config.MetaFileName == "-" {
config.metaFile = os.Stderr config.metaFile = os.Stderr

@ -3,6 +3,7 @@ package zgrab2
import ( import (
"bufio" "bufio"
"fmt" "fmt"
"io"
) )
// FlagMap is a function that maps a single-bit bitmask (i.e. a number of the // FlagMap is a function that maps a single-bit bitmask (i.e. a number of the
@ -133,16 +134,23 @@ func WidenMapKeys(input map[int]string) map[uint64]string {
// results or error. // results or error.
type OutputResultsFunc func(results <-chan []byte) error type OutputResultsFunc func(results <-chan []byte) error
// OutputResultsFile is an OutputResultsFunc that write results to // OutputResultsWriterFunc returns an OutputResultsFunc that wraps an io.Writer
// a filename provided on the command line. // in a buffered writer, and uses OutputResults.
func OutputResultsFile(results <-chan []byte) error { func OutputResultsWriterFunc(w io.Writer) OutputResultsFunc {
out := bufio.NewWriter(config.outputFile) buf := bufio.NewWriter(w)
defer out.Flush() return func(result <-chan []byte) error {
defer buf.Flush()
return OutputResults(buf, result)
}
}
// OutputResults writes results to a buffered Writer from a channel.
func OutputResults(w *bufio.Writer, results <-chan []byte) error {
for result := range results { for result := range results {
if _, err := out.Write(result); err != nil { if _, err := w.Write(result); err != nil {
return err return err
} }
if err := out.WriteByte('\n'); err != nil { if err := w.WriteByte('\n'); err != nil {
return err return err
} }
} }