Add support for importing the ZGrab2 main (#224)

Move ZGrab2's main function to a library, and call it in cmd/zgrab2 after importing all of our modules. Consumes of ZGrab2 as a library can use the same approach to provide custom sets of modules, without having to hack the build system or reimplement main. #224
zmap · Sep 25, 2019 · 8427a23 · Alpha35 · Apr 15, 2020 · 8427a23
1 parent dda796c
commit 8427a23
Show file tree

Hide file tree

Showing 3 changed files with 160 additions and 142 deletions.
diff --git a/bin/bin.go b/bin/bin.go
@@ -0,0 +1,154 @@
+package bin
+
+import (
+	"encoding/json"
+	"os"
+	"runtime/pprof"
+	"time"
+
+	"fmt"
+	"runtime"
+	"strings"
+
+	log "github.com/sirupsen/logrus"
+	flags "github.com/zmap/zflags"
+	"github.com/zmap/zgrab2"
+)
+
+// Get the value of the ZGRAB2_MEMPROFILE variable (or the empty string).
+// This may include {TIMESTAMP} or {NANOS}, which should be replaced using
+// getFormattedFile().
+func getMemProfileFile() string {
+	return os.Getenv("ZGRAB2_MEMPROFILE")
+}
+
+// Get the value of the ZGRAB2_CPUPROFILE variable (or the empty string).
+// This may include {TIMESTAMP} or {NANOS}, which should be replaced using
+// getFormattedFile().
+func getCPUProfileFile() string {
+	return os.Getenv("ZGRAB2_CPUPROFILE")
+}
+
+// Replace instances in formatString of {TIMESTAMP} with when formatted as
+// YYYYMMDDhhmmss, and {NANOS} as the decimal nanosecond offset.
+func getFormattedFile(formatString string, when time.Time) string {
+	timestamp := when.Format("20060102150405")
+	nanos := fmt.Sprintf("%d", when.Nanosecond())
+	ret := strings.Replace(formatString, "{TIMESTAMP}", timestamp, -1)
+	ret = strings.Replace(ret, "{NANOS}", nanos, -1)
+	return ret
+}
+
+// If memory profiling is enabled (ZGRAB2_MEMPROFILE is not empty), perform a GC
+// then write the heap profile to the profile file.
+func dumpHeapProfile() {
+	if file := getMemProfileFile(); file != "" {
+		now := time.Now()
+		fullFile := getFormattedFile(file, now)
+		f, err := os.Create(fullFile)
+		if err != nil {
+			log.Fatal("could not create heap profile: ", err)
+		}
+		runtime.GC()
+		if err := pprof.WriteHeapProfile(f); err != nil {
+			log.Fatal("could not write heap profile: ", err)
+		}
+		f.Close()
+	}
+}
+
+// If CPU profiling is enabled (ZGRAB2_CPUPROFILE is not empty), start tracking
+// CPU profiling in the configured file. Caller is responsible for invoking
+// stopCPUProfile() when finished.
+func startCPUProfile() {
+	if file := getCPUProfileFile(); file != "" {
+		now := time.Now()
+		fullFile := getFormattedFile(file, now)
+		f, err := os.Create(fullFile)
+		if err != nil {
+			log.Fatal("could not create CPU profile: ", err)
+		}
+		if err := pprof.StartCPUProfile(f); err != nil {
+			log.Fatal("could not start CPU profile: ", err)
+		}
+	}
+}
+
+// If CPU profiling is enabled (ZGRAB2_CPUPROFILE is not empty), stop profiling
+// CPU usage.
+func stopCPUProfile() {
+	if getCPUProfileFile() != "" {
+		pprof.StopCPUProfile()
+	}
+}
+
+// ZGrab2Main should be called by func main() in a binary. The caller is
+// responsible for importing any modules in use. This allows clients to easily
+// include custom sets of scan modules by creating new main packages with custom
+// sets of ZGrab modules imported with side-effects.
+func ZGrab2Main() {
+	startCPUProfile()
+	defer stopCPUProfile()
+	defer dumpHeapProfile()
+	_, moduleType, flag, err := zgrab2.ParseCommandLine(os.Args[1:])
+
+	// Blanked arg is positional arguments
+	if err != nil {
+		// Outputting help is returned as an error. Exit successfuly on help output.
+		flagsErr, ok := err.(*flags.Error)
+		if ok && flagsErr.Type == flags.ErrHelp {
+			return
+		}
+
+		// Didn't output help. Unknown parsing error.
+		log.Fatalf("could not parse flags: %s", err)
+	}
+
+	if m, ok := flag.(*zgrab2.MultipleCommand); ok {
+		iniParser := zgrab2.NewIniParser()
+		var modTypes []string
+		var flagsReturned []interface{}
+		if m.ConfigFileName == "-" {
+			modTypes, flagsReturned, err = iniParser.Parse(os.Stdin)
+		} else {
+			modTypes, flagsReturned, err = iniParser.ParseFile(m.ConfigFileName)
+		}
+		if err != nil {
+			log.Fatalf("could not parse multiple: %s", err)
+		}
+		if len(modTypes) != len(flagsReturned) {
+			log.Fatalf("error parsing flags")
+		}
+		for i, fl := range flagsReturned {
+			f, _ := fl.(zgrab2.ScanFlags)
+			mod := zgrab2.GetModule(modTypes[i])
+			s := mod.NewScanner()
+			s.Init(f)
+			zgrab2.RegisterScan(s.GetName(), s)
+		}
+	} else {
+		mod := zgrab2.GetModule(moduleType)
+		s := mod.NewScanner()
+		s.Init(flag)
+		zgrab2.RegisterScan(moduleType, s)
+	}
+	monitor := zgrab2.MakeMonitor()
+	monitor.Callback = func(_ string) {
+		dumpHeapProfile()
+	}
+	start := time.Now()
+	log.Infof("started grab at %s", start.Format(time.RFC3339))
+	zgrab2.Process(monitor)
+	end := time.Now()
+	log.Infof("finished grab at %s", end.Format(time.RFC3339))
+	s := Summary{
+		StatusesPerModule: monitor.GetStatuses(),
+		StartTime:         start.Format(time.RFC3339),
+		EndTime:           end.Format(time.RFC3339),
+		Duration:          end.Sub(start).String(),
+	}
+	enc := json.NewEncoder(zgrab2.GetMetaFile())
+	if err := enc.Encode(&s); err != nil {
+		log.Fatalf("unable to write summary: %s", err.Error())
+	}
+}
diff --git a/cmd/zgrab2/summary.go → bin/summary.go b/cmd/zgrab2/summary.go → bin/summary.go
@@ -1,7 +1,8 @@
-package main
+package bin
 
 import "github.com/zmap/zgrab2"
 
+// Summary holds the results of a run of a ZGrab2 binary.
 type Summary struct {
 	StatusesPerModule map[string]*zgrab2.State `json:"statuses"`
 	StartTime         string                   `json:"start"`

diff --git a/cmd/zgrab2/main.go b/cmd/zgrab2/main.go
@@ -1,149 +1,12 @@
 package main
 
 import (
-	"encoding/json"
-	"os"
-	"time"
-	"runtime/pprof"
-
-	flags "github.com/zmap/zflags"
-	log "github.com/sirupsen/logrus"
-	"github.com/zmap/zgrab2"
+	"github.com/zmap/zgrab2/bin"
 	_ "github.com/zmap/zgrab2/modules"
-	"fmt"
-	"runtime"
-	"strings"
 )
 
-// Get the value of the ZGRAB2_MEMPROFILE variable (or the empty string).
-// This may include {TIMESTAMP} or {NANOS}, which should be replaced using
-// getFormattedFile().
-func getMemProfileFile() string {
-	return os.Getenv("ZGRAB2_MEMPROFILE")
-}
-
-// Get the value of the ZGRAB2_CPUPROFILE variable (or the empty string).
-// This may include {TIMESTAMP} or {NANOS}, which should be replaced using
-// getFormattedFile().
-func getCPUProfileFile() string {
-	return os.Getenv("ZGRAB2_CPUPROFILE")
-}
-
-// Replace instances in formatString of {TIMESTAMP} with when formatted as
-// YYYYMMDDhhmmss, and {NANOS} as the decimal nanosecond offset.
-func getFormattedFile(formatString string, when time.Time) string {
-	timestamp := when.Format("20060102150405")
-	nanos := fmt.Sprintf("%d", when.Nanosecond())
-	ret := strings.Replace(formatString, "{TIMESTAMP}", timestamp, -1)
-	ret = strings.Replace(ret, "{NANOS}", nanos, -1)
-	return ret
-}
-
-// If memory profiling is enabled (ZGRAB2_MEMPROFILE is not empty), perform a GC
-// then write the heap profile to the profile file.
-func dumpHeapProfile() {
-	if file := getMemProfileFile(); file != "" {
-		now := time.Now()
-		fullFile := getFormattedFile(file, now)
-		f, err := os.Create(fullFile)
-		if err != nil {
-			log.Fatal("could not create heap profile: ", err)
-		}
-		runtime.GC()
-		if err := pprof.WriteHeapProfile(f); err != nil {
-			log.Fatal("could not write heap profile: ", err)
-		}
-		f.Close()
-	}
-}
-
-// If CPU profiling is enabled (ZGRAB2_CPUPROFILE is not empty), start tracking
-// CPU profiling in the configured file. Caller is responsible for invoking
-// stopCPUProfile() when finished.
-func startCPUProfile() {
-	if file := getCPUProfileFile(); file != "" {
-		now := time.Now()
-		fullFile := getFormattedFile(file, now)
-		f, err := os.Create(fullFile)
-		if err != nil {
-			log.Fatal("could not create CPU profile: ", err)
-		}
-		if err := pprof.StartCPUProfile(f); err != nil {
-			log.Fatal("could not start CPU profile: ", err)
-		}
-	}
-}
-
-// If CPU profiling is enabled (ZGRAB2_CPUPROFILE is not empty), stop profiling
-// CPU usage.
-func stopCPUProfile() {
-	if getCPUProfileFile() != "" {
-		pprof.StopCPUProfile()
-	}
-}
+// main wraps the "true" main, bin.ZGrab2Main(), after importing all scan
+// modules in ZGrab2.
 func main() {
-	startCPUProfile()
-	defer stopCPUProfile()
-	defer dumpHeapProfile()
-	_, moduleType, flag, err := zgrab2.ParseCommandLine(os.Args[1:])
-
-	// Blanked arg is positional arguments
-	if err != nil {
-		// Outputting help is returned as an error. Exit successfuly on help output.
-		flagsErr, ok := err.(*flags.Error)
-		if ok && flagsErr.Type == flags.ErrHelp {
-			return
-		}
-
-		// Didn't output help. Unknown parsing error.
-		log.Fatalf("could not parse flags: %s", err)
-	}
-
-	if m, ok := flag.(*zgrab2.MultipleCommand); ok {
-		iniParser := zgrab2.NewIniParser()
-		var modTypes []string
-		var flagsReturned []interface{}
-		if m.ConfigFileName == "-" {
-			modTypes, flagsReturned, err = iniParser.Parse(os.Stdin)
-		} else {
-			modTypes, flagsReturned, err = iniParser.ParseFile(m.ConfigFileName)
-		}
-		if err != nil {
-			log.Fatalf("could not parse multiple: %s", err)
-		}
-		if len(modTypes) != len(flagsReturned) {
-			log.Fatalf("error parsing flags")
-		}
-		for i, fl := range flagsReturned {
-			f, _ := fl.(zgrab2.ScanFlags)
-			mod := zgrab2.GetModule(modTypes[i])
-			s := mod.NewScanner()
-			s.Init(f)
-			zgrab2.RegisterScan(s.GetName(), s)
-		}
-	} else {
-		mod := zgrab2.GetModule(moduleType)
-		s := mod.NewScanner()
-		s.Init(flag)
-		zgrab2.RegisterScan(moduleType, s)
-	}
-	monitor := zgrab2.MakeMonitor()
-	monitor.Callback = func(_ string) {
-		dumpHeapProfile()
-	}
-	start := time.Now()
-	log.Infof("started grab at %s", start.Format(time.RFC3339))
-	zgrab2.Process(monitor)
-	end := time.Now()
-	log.Infof("finished grab at %s", end.Format(time.RFC3339))
-	s := Summary{
-		StatusesPerModule: monitor.GetStatuses(),
-		StartTime:         start.Format(time.RFC3339),
-		EndTime:           end.Format(time.RFC3339),
-		Duration:          end.Sub(start).String(),
-	}
-	enc := json.NewEncoder(zgrab2.GetMetaFile())
-	if err := enc.Encode(&s); err != nil {
-		log.Fatalf("unable to write summary: %s", err.Error())
-	}
+	bin.ZGrab2Main()
 }