Skip to content

Commit

Permalink
Add support for importing the ZGrab2 main (#224)
Browse files Browse the repository at this point in the history
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
  • Loading branch information
dadrian authored Sep 25, 2019
1 parent dda796c commit 8427a23
Show file tree
Hide file tree
Showing 3 changed files with 160 additions and 142 deletions.
154 changes: 154 additions & 0 deletions bin/bin.go
Original file line number Diff line number Diff line change
@@ -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) {

This comment has been minimized.

Copy link
@Alpha35

Alpha35 Apr 15, 2020

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())
}
}
3 changes: 2 additions & 1 deletion cmd/zgrab2/summary.go → bin/summary.go
Original file line number Diff line number Diff line change
@@ -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"`
Expand Down
145 changes: 4 additions & 141 deletions cmd/zgrab2/main.go
Original file line number Diff line number Diff line change
@@ -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()
}

0 comments on commit 8427a23

Please sign in to comment.