Added simple json based http metrics support

dev
Ice3man543 2020-12-17 20:33:42 +05:30
parent 019c61037c
commit 15708cb941
3 changed files with 82 additions and 30 deletions

View File

@ -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 {
progress := &Progress{}
stats, err := clistats.New()
if err != nil {
gologger.Warningf("Couldn't create progress engine: %s\n", err)
return nil, err
}
progress.active = active
progress.stats = stats
progress.tickDuration = tickDuration
}
return &progress
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) {
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 p.active {
if err := p.stats.Start(makePrintCallback(), p.tickDuration); err != nil {
gologger.Warningf("Couldn't start statistics: %s\n", err)
}
@ -57,26 +78,20 @@ 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))
}
}
// Update progress tracking information and increments the request counter by one unit.
func (p *Progress) Update() {
if p.active {
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))
}
}
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())
}

View File

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

View File

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