diff --git a/main.go b/main.go index 2d8ce4a..903e89d 100644 --- a/main.go +++ b/main.go @@ -12,7 +12,6 @@ import ( var ( outputDirectory string - printStats bool ) func main() { @@ -22,15 +21,10 @@ func main() { net.RegisterListenFlags() net.RegisterCryptFlags() logistic.RegisterPackerFlags() + stats.RegisterStatsFlags() RegisterGlobalFlags() flag.Parse() - // Check if we have the only required argument present - outputDirectory - if outputDirectory == "" { - fmt.Println("Please specify fuzzer-directory. See help (--help) for details.") - return - } - // Read peers file net.ReadPeers() @@ -45,9 +39,7 @@ func main() { go watchdog.WatchFuzzers(outputDirectory) // Start stat printer - if printStats { - go stats.PrintStats() - } + go stats.PrintStats() // Listen for incoming connections listenErr := net.Listen(outputDirectory) @@ -59,5 +51,4 @@ func main() { // Registers flags which are required by multiple modules and need to be handled here func RegisterGlobalFlags() { flag.StringVar(&outputDirectory, "fuzzer-directory", "", "The output directory of the fuzzer(s)") - flag.BoolVar(&printStats, "print-stats", true, "Print traffic statistics every few seconds") } diff --git a/stats/stats.go b/stats/stats.go index bc50114..f6b6605 100644 --- a/stats/stats.go +++ b/stats/stats.go @@ -1,6 +1,7 @@ package stats import ( + "flag" "fmt" "github.com/dustin/go-humanize" "time" @@ -15,6 +16,14 @@ type Stat struct { // statPipe is a channel used to var stats Stat +// printStats sets whether we should print stats or not +var printStats bool + +// RegisterStatsFlags registers all flags required by the stats module +func RegisterStatsFlags() { + flag.BoolVar(&printStats, "print-stats", true, "Print traffic statistics every few seconds") +} + // PushStat pushes the given stat func PushStat(s Stat) { stats.SentBytes += s.SentBytes @@ -23,6 +32,11 @@ func PushStat(s Stat) { // PrintStats periodically prints the collected statistics func PrintStats() { + // Check if we should print stats + if !printStats { + return + } + t := time.NewTicker(2 * time.Second) for {