Provides a framework for modularized target input and result output functions.
This commit is contained in:
parent
4f046b3085
commit
1f5e4b0f22
14
config.go
14
config.go
|
@ -27,6 +27,18 @@ type Config struct {
|
|||
outputFile *os.File
|
||||
metaFile *os.File
|
||||
logFile *os.File
|
||||
inputTargets InputTargetsFunc
|
||||
outputResults OutputResultsFunc
|
||||
}
|
||||
|
||||
// SetInputFunc sets the target input function to the provided function.
|
||||
func SetInputFunc(f InputTargetsFunc) {
|
||||
config.inputTargets = f
|
||||
}
|
||||
|
||||
// SetOutputFunc sets the result output function to the provided function.
|
||||
func SetOutputFunc(f OutputResultsFunc) {
|
||||
config.outputResults = f
|
||||
}
|
||||
|
||||
func init() {
|
||||
|
@ -46,6 +58,7 @@ func validateFrameworkConfiguration() {
|
|||
}
|
||||
log.SetOutput(config.logFile)
|
||||
}
|
||||
SetInputFunc(InputTargetsCSV)
|
||||
|
||||
if config.InputFileName == "-" {
|
||||
config.inputFile = os.Stdin
|
||||
|
@ -64,6 +77,7 @@ func validateFrameworkConfiguration() {
|
|||
log.Fatal(err)
|
||||
}
|
||||
}
|
||||
SetOutputFunc(OutputResultsFile)
|
||||
|
||||
if config.MetaFileName == "-" {
|
||||
config.metaFile = os.Stderr
|
||||
|
|
12
input.go
12
input.go
|
@ -80,6 +80,12 @@ func duplicateIP(ip net.IP) net.IP {
|
|||
return dup
|
||||
}
|
||||
|
||||
// InputTargetsCSV is an InputTargetsFunc that calls GetTargetsCSV with
|
||||
// the CSV file provided on the command line.
|
||||
func InputTargetsCSV(ch chan<- ScanTarget) error {
|
||||
return GetTargetsCSV(config.inputFile, ch)
|
||||
}
|
||||
|
||||
// GetTargetsCSV reads targets from a CSV source, generates ScanTargets,
|
||||
// and delivers them to the provided channel.
|
||||
func GetTargetsCSV(source io.Reader, ch chan<- ScanTarget) error {
|
||||
|
@ -117,3 +123,9 @@ func GetTargetsCSV(source io.Reader, ch chan<- ScanTarget) error {
|
|||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// InputTargetsFunc is a function type for target input functions.
|
||||
//
|
||||
// A function of this type generates ScanTargets on the provided
|
||||
// channel. It returns nil if there are no further inputs or error.
|
||||
type InputTargetsFunc func(ch chan<- ScanTarget) error
|
||||
|
|
28
output.go
28
output.go
|
@ -1,6 +1,9 @@
|
|||
package zgrab2
|
||||
|
||||
import "fmt"
|
||||
import (
|
||||
"bufio"
|
||||
"fmt"
|
||||
)
|
||||
|
||||
// FlagMap is a function that maps a single-bit bitmask (i.e. a number of the
|
||||
// form (1 << x)) to a string representing that bit.
|
||||
|
@ -122,3 +125,26 @@ func WidenMapKeys(input map[int]string) map[uint64]string {
|
|||
}
|
||||
return ret
|
||||
}
|
||||
|
||||
// OutputResultsFunc is a function type for result output functions.
|
||||
//
|
||||
// A function of this type receives results on the provided channel
|
||||
// and outputs them somehow. It returns nil if there are no further
|
||||
// results or error.
|
||||
type OutputResultsFunc func(results <-chan []byte) error
|
||||
|
||||
// OutputResultsFile is an OutputResultsFunc that write results to
|
||||
// a filename provided on the command line.
|
||||
func OutputResultsFile(results <-chan []byte) error {
|
||||
out := bufio.NewWriter(config.outputFile)
|
||||
defer out.Flush()
|
||||
for result := range results {
|
||||
if _, err := out.Write(result); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := out.WriteByte('\n'); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
|
|
@ -1,7 +1,6 @@
|
|||
package zgrab2
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net"
|
||||
|
@ -146,7 +145,7 @@ func grabTarget(input ScanTarget, m *Monitor) []byte {
|
|||
return result
|
||||
}
|
||||
|
||||
// Process sets up an output encoder, input reader, and starts grab workers
|
||||
// Process sets up an output encoder, input reader, and starts grab workers.
|
||||
func Process(mon *Monitor) {
|
||||
workers := config.Senders
|
||||
processQueue := make(chan ScanTarget, workers*4)
|
||||
|
@ -160,16 +159,9 @@ func Process(mon *Monitor) {
|
|||
|
||||
// Start the output encoder
|
||||
go func() {
|
||||
out := bufio.NewWriter(config.outputFile)
|
||||
defer outputDone.Done()
|
||||
defer out.Flush()
|
||||
for result := range outputQueue {
|
||||
if _, err := out.Write(result); err != nil {
|
||||
log.Fatal(err)
|
||||
}
|
||||
if err := out.WriteByte('\n'); err != nil {
|
||||
log.Fatal(err)
|
||||
}
|
||||
if err := config.outputResults(outputQueue); err != nil {
|
||||
log.Fatal(err)
|
||||
}
|
||||
}()
|
||||
//Start all the workers
|
||||
|
@ -189,7 +181,7 @@ func Process(mon *Monitor) {
|
|||
}(i)
|
||||
}
|
||||
|
||||
if err := GetTargetsCSV(config.inputFile, processQueue); err != nil {
|
||||
if err := config.inputTargets(processQueue); err != nil {
|
||||
log.Fatal(err)
|
||||
}
|
||||
close(processQueue)
|
||||
|
|
|
@ -24,6 +24,12 @@ func NewIniParser() *flags.IniParser {
|
|||
return flags.NewIniParser(parser)
|
||||
}
|
||||
|
||||
// AddGroup exposes the parser's AddGroup function, allowing extension
|
||||
// of the global arguments.
|
||||
func AddGroup(shortDescription string, longDescription string, data interface{}) {
|
||||
parser.AddGroup(shortDescription, longDescription, data)
|
||||
}
|
||||
|
||||
// AddCommand adds a module to the parser and returns a pointer to
|
||||
// a flags.command object or an error
|
||||
func AddCommand(command string, shortDescription string, longDescription string, port int, m ScanModule) (*flags.Command, error) {
|
||||
|
|
Loading…
Reference in New Issue