From 15708cb9411f3519b6b927d2243a56d077849cea Mon Sep 17 00:00:00 2001 From: Ice3man543 Date: Thu, 17 Dec 2020 20:33:42 +0530 Subject: [PATCH] Added simple json based http metrics support --- v2/internal/progress/progress.go | 102 ++++++++++++++++++++++--------- v2/internal/runner/options.go | 4 ++ v2/internal/runner/runner.go | 6 +- 3 files changed, 82 insertions(+), 30 deletions(-) diff --git a/v2/internal/progress/progress.go b/v2/internal/progress/progress.go index 6c12be52..112f5c4e 100644 --- a/v2/internal/progress/progress.go +++ b/v2/internal/progress/progress.go @@ -1,8 +1,13 @@ package progress import ( + "context" + "encoding/json" "fmt" + "net" + "net/http" "os" + "strconv" "strings" "time" @@ -13,12 +18,13 @@ import ( // Progress is a progress instance for showing program stats type Progress struct { active bool - stats clistats.StatisticsClient tickDuration time.Duration + stats clistats.StatisticsClient + server *http.Server } // NewProgress creates and returns a new progress tracking object. -func NewProgress(active bool) *Progress { +func NewProgress(active, metrics bool, port int) (*Progress, error) { var tickDuration time.Duration if active { tickDuration = 5 * time.Second @@ -26,29 +32,44 @@ func NewProgress(active bool) *Progress { tickDuration = -1 } - var progress Progress - if active { - stats, err := clistats.New() - if err != nil { - gologger.Warningf("Couldn't create progress engine: %s\n", err) - } - progress.active = active - progress.stats = stats - progress.tickDuration = tickDuration - } + progress := &Progress{} - return &progress + stats, err := clistats.New() + if err != nil { + return nil, err + } + progress.active = active + progress.stats = stats + progress.tickDuration = tickDuration + + if metrics { + http.HandleFunc("/metrics", func(w http.ResponseWriter, req *http.Request) { + metrics := progress.getMetrics() + _ = json.NewEncoder(w).Encode(metrics) + }) + progress.server = &http.Server{ + Addr: net.JoinHostPort("127.0.0.1", strconv.Itoa(port)), + Handler: http.DefaultServeMux, + } + go func() { + if err := progress.server.ListenAndServe(); err != nil { + gologger.Warningf("Could not serve metrics: %s\n", err) + } + }() + } + return progress, nil } // Init initializes the progress display mechanism by setting counters, etc. func (p *Progress) Init(hostCount int64, rulesCount int, requestCount int64) { + p.stats.AddStatic("templates", rulesCount) + p.stats.AddStatic("hosts", hostCount) + p.stats.AddStatic("startedAt", time.Now()) + p.stats.AddCounter("requests", uint64(0)) + p.stats.AddCounter("errors", uint64(0)) + p.stats.AddCounter("total", uint64(requestCount)) + if p.active { - p.stats.AddStatic("templates", rulesCount) - p.stats.AddStatic("hosts", hostCount) - p.stats.AddStatic("startedAt", time.Now()) - p.stats.AddCounter("requests", uint64(0)) - p.stats.AddCounter("errors", uint64(0)) - p.stats.AddCounter("total", uint64(requestCount)) if err := p.stats.Start(makePrintCallback(), p.tickDuration); err != nil { gologger.Warningf("Couldn't start statistics: %s\n", err) } @@ -57,25 +78,19 @@ func (p *Progress) Init(hostCount int64, rulesCount int, requestCount int64) { // AddToTotal adds a value to the total request count func (p *Progress) AddToTotal(delta int64) { - if p.active { - p.stats.IncrementCounter("total", int(delta)) - } + p.stats.IncrementCounter("total", int(delta)) } // Update progress tracking information and increments the request counter by one unit. func (p *Progress) Update() { - if p.active { - p.stats.IncrementCounter("requests", 1) - } + p.stats.IncrementCounter("requests", 1) } // Drop drops the specified number of requests from the progress bar total. // This may be the case when uncompleted requests are encountered and shouldn't be part of the total count. func (p *Progress) Drop(count int64) { - if p.active { - // mimic dropping by incrementing the completed requests - p.stats.IncrementCounter("errors", int(count)) - } + // mimic dropping by incrementing the completed requests + p.stats.IncrementCounter("errors", int(count)) } const bufferSize = 128 @@ -125,6 +140,34 @@ func makePrintCallback() func(stats clistats.StatisticsClient) { } } +// getMetrics returns a map of important metrics for client +func (p *Progress) getMetrics() map[string]interface{} { + results := make(map[string]interface{}) + + startedAt, _ := p.stats.GetStatic("startedAt") + duration := time.Since(startedAt.(time.Time)) + + results["startedAt"] = startedAt.(time.Time) + results["duration"] = fmtDuration(duration) + templates, _ := p.stats.GetStatic("templates") + results["templates"] = clistats.String(templates) + hosts, _ := p.stats.GetStatic("hosts") + results["hosts"] = clistats.String(hosts) + requests, _ := p.stats.GetCounter("requests") + results["requests"] = clistats.String(requests) + total, _ := p.stats.GetCounter("total") + results["total"] = clistats.String(total) + results["rps"] = clistats.String(uint64(float64(requests) / duration.Seconds())) + errors, _ := p.stats.GetCounter("errors") + results["errors"] = clistats.String(errors) + + //nolint:gomnd // this is not a magic number + percentData := (float64(requests) * float64(100)) / float64(total) + percent := clistats.String(uint64(percentData)) + results["percent"] = percent + return results +} + // fmtDuration formats the duration for the time elapsed func fmtDuration(d time.Duration) string { d = d.Round(time.Second) @@ -143,4 +186,5 @@ func (p *Progress) Stop() { gologger.Warningf("Couldn't stop statistics: %s\n", err) } } + p.server.Shutdown(context.Background()) } diff --git a/v2/internal/runner/options.go b/v2/internal/runner/options.go index 8dfa7f19..55ca127d 100644 --- a/v2/internal/runner/options.go +++ b/v2/internal/runner/options.go @@ -15,6 +15,8 @@ import ( // nolint // false positive, options are allocated once and are necessary as is type Options struct { MaxWorkflowDuration int // MaxWorkflowDuration is the maximum time a workflow can run for a URL + Metrics bool // Metrics enables display of metrics via an http endpoint + MetricsPort int // MetricsPort is the port to show metrics on Sandbox bool // Sandbox mode allows users to run isolated workflows with system commands disabled Debug bool // Debug mode allows debugging request/responses for the engine Silent bool // Silent suppresses any extra text and only writes found URLs on screen. @@ -69,6 +71,8 @@ func ParseOptions() *Options { options := &Options{} flag.BoolVar(&options.Sandbox, "sandbox", false, "Run workflows in isolated sandbox mode") + flag.BoolVar(&options.Metrics, "metrics", false, "Expose nuclei metrics on a port") + flag.IntVar(&options.MetricsPort, "metrics-port", 9092, "Port to expose nuclei metrics on") flag.IntVar(&options.MaxWorkflowDuration, "workflow-duration", 10, "Max time for workflow run on single URL in minutes") flag.StringVar(&options.Target, "target", "", "Target is a single target to scan using template") flag.Var(&options.Templates, "t", "Template input dir/file/files to run on host. Can be used multiple times. Supports globbing.") diff --git a/v2/internal/runner/runner.go b/v2/internal/runner/runner.go index 46c7d4e8..79ded4bf 100644 --- a/v2/internal/runner/runner.go +++ b/v2/internal/runner/runner.go @@ -173,7 +173,11 @@ func New(options *Options) (*Runner, error) { } // Creates the progress tracking object - runner.progress = progress.NewProgress(options.EnableProgressBar) + var progressErr error + runner.progress, progressErr = progress.NewProgress(options.EnableProgressBar, options.Metrics, options.MetricsPort) + if progressErr != nil { + return nil, progressErr + } // create project file if requested or load existing one if options.Project {