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

11
main.go
View File

@ -12,7 +12,6 @@ import (
var ( var (
outputDirectory string outputDirectory string
printStats bool
) )
func main() { func main() {
@ -22,15 +21,10 @@ func main() {
net.RegisterListenFlags() net.RegisterListenFlags()
net.RegisterCryptFlags() net.RegisterCryptFlags()
logistic.RegisterPackerFlags() logistic.RegisterPackerFlags()
stats.RegisterStatsFlags()
RegisterGlobalFlags() RegisterGlobalFlags()
flag.Parse() 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 // Read peers file
net.ReadPeers() net.ReadPeers()
@ -45,9 +39,7 @@ func main() {
go watchdog.WatchFuzzers(outputDirectory) go watchdog.WatchFuzzers(outputDirectory)
// Start stat printer // Start stat printer
if printStats {
go stats.PrintStats() go stats.PrintStats()
}
// Listen for incoming connections // Listen for incoming connections
listenErr := net.Listen(outputDirectory) listenErr := net.Listen(outputDirectory)
@ -59,5 +51,4 @@ func main() {
// Registers flags which are required by multiple modules and need to be handled here // Registers flags which are required by multiple modules and need to be handled here
func RegisterGlobalFlags() { func RegisterGlobalFlags() {
flag.StringVar(&outputDirectory, "fuzzer-directory", "", "The output directory of the fuzzer(s)") 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 package stats
import ( import (
"flag"
"fmt" "fmt"
"github.com/dustin/go-humanize" "github.com/dustin/go-humanize"
"time" "time"
@ -15,6 +16,14 @@ type Stat struct {
// statPipe is a channel used to // statPipe is a channel used to
var stats Stat 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 // PushStat pushes the given stat
func PushStat(s Stat) { func PushStat(s Stat) {
stats.SentBytes += s.SentBytes stats.SentBytes += s.SentBytes
@ -23,6 +32,11 @@ func PushStat(s Stat) {
// PrintStats periodically prints the collected statistics // PrintStats periodically prints the collected statistics
func PrintStats() { func PrintStats() {
// Check if we should print stats
if !printStats {
return
}
t := time.NewTicker(2 * time.Second) t := time.NewTicker(2 * time.Second)
for { for {