From 1f5e4b0f22e67498c7bef23eac8aa29d72eeff3b Mon Sep 17 00:00:00 2001
From: Alex Halderman <jhalderm@eecs.umich.edu>
Date: Thu, 5 Jul 2018 12:41:34 -0400
Subject: [PATCH] Provides a framework for modularized target input and result
 output functions.

---
 config.go     | 14 ++++++++++++++
 input.go      | 12 ++++++++++++
 output.go     | 28 +++++++++++++++++++++++++++-
 processing.go | 16 ++++------------
 utility.go    |  6 ++++++
 5 files changed, 63 insertions(+), 13 deletions(-)

diff --git a/config.go b/config.go
index 530b8dc..1903e89 100644
--- a/config.go
+++ b/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
diff --git a/input.go b/input.go
index 2066d4c..8830560 100644
--- a/input.go
+++ b/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
diff --git a/output.go b/output.go
index 85c11e8..5b6e0bc 100644
--- a/output.go
+++ b/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
+}
diff --git a/processing.go b/processing.go
index 0c0a3ec..7c7bd49 100644
--- a/processing.go
+++ b/processing.go
@@ -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)
diff --git a/utility.go b/utility.go
index dbf6cc1..8ac1e05 100644
--- a/utility.go
+++ b/utility.go
@@ -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) {