Move stats flag to stats module

This commit is contained in:
maride 2021-04-23 19:23:44 +02:00
parent 5cb6c1af9b
commit c7dff7c496
2 changed files with 16 additions and 11 deletions

13
main.go
View File

@ -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")
}

View File

@ -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 {