diff --git a/v2/internal/progress/progress.go b/v2/internal/progress/progress.go index dc4dce02..2ae847d6 100644 --- a/v2/internal/progress/progress.go +++ b/v2/internal/progress/progress.go @@ -20,29 +20,34 @@ type Progress struct { stdCaptureMutex *sync.Mutex stdout *strings.Builder stderr *strings.Builder + + colorizer aurora.Aurora } // Creates and returns a new progress tracking object. -func NewProgress(group *sync.WaitGroup) *Progress { +func NewProgress(noColor bool) *Progress { p := &Progress{ progress: mpb.New( - mpb.WithWaitGroup(group), mpb.WithOutput(os.Stderr), mpb.PopCompletedMode(), ), stdCaptureMutex: &sync.Mutex{}, stdout: &strings.Builder{}, stderr: &strings.Builder{}, + colorizer: aurora.NewAurora(!noColor), } return p } // Creates and returns a progress bar that tracks request progress for a specific template. func (p *Progress) SetupTemplateProgressbar(templateIndex int, templateCount int, name string, requestCount int64) { - barName := "[" + aurora.Green(name).String() + "]" + color := p.colorizer + barName := "[" + color.Green(name).String() + "]" if templateIndex > -1 && templateCount > -1 { - barName = aurora.Sprintf("[%d/%d] ", aurora.Bold(aurora.Cyan(templateIndex)), aurora.Cyan(templateCount)) + barName + barName = color.Sprintf("[%d/%d] ", + color.Bold(color.Cyan(templateIndex)), + color.Cyan(templateCount)) + barName } bar := p.setupProgressbar(barName, requestCount) @@ -62,15 +67,17 @@ func (p *Progress) SetupTemplateProgressbar(templateIndex int, templateCount int // Creates and returns a progress bar that tracks all the requests progress. // This is only useful when multiple templates are processed within the same run. func (p *Progress) SetupGlobalProgressbar(hostCount int64, templateCount int, requestCount int64) { + color := p.colorizer + hostPlural := "host" if hostCount > 1 { hostPlural = "hosts" } - barName := "[" + aurora.Sprintf( - aurora.Cyan("%d templates, %d %s"), - aurora.Bold(aurora.Cyan(templateCount)), - aurora.Bold(aurora.Cyan(hostCount)), + barName := "[" + color.Sprintf( + color.Cyan("%d templates, %d %s"), + color.Bold(color.Cyan(templateCount)), + color.Bold(color.Cyan(hostCount)), hostPlural) + "]" bar := p.setupProgressbar(barName, requestCount) @@ -118,17 +125,19 @@ func (p *Progress) Wait() { // Creates and returns a progress bar. func (p *Progress) setupProgressbar(name string, total int64) *mpb.Bar { + color := p.colorizer return p.progress.AddBar( total, mpb.BarNoPop(), mpb.BarRemoveOnComplete(), mpb.PrependDecorators( decor.Name(name, decor.WCSyncSpaceR), - decor.CountersNoUnit(aurora.Blue(" %d/%d").String(), decor.WCSyncSpace), - decor.NewPercentage(aurora.Bold("%d").String(), decor.WCSyncSpace), + decor.CountersNoUnit(color.Blue(" %d/%d").String(), decor.WCSyncSpace), + decor.NewPercentage(color.Bold("%d").String(), decor.WCSyncSpace), ), mpb.AppendDecorators( - decor.AverageSpeed(0, aurora.Yellow("%.2f r/s ").String(), decor.WCSyncSpace), + decor.AverageSpeed(0, color.Yellow("%.2f r/s ").String(), decor.WCSyncSpace), + decor.Elapsed(decor.ET_STYLE_GO, decor.WCSyncSpace), decor.AverageETA(decor.ET_STYLE_GO, decor.WCSyncSpace), ), ) diff --git a/v2/internal/runner/runner.go b/v2/internal/runner/runner.go index 7859f74f..1f293192 100644 --- a/v2/internal/runner/runner.go +++ b/v2/internal/runner/runner.go @@ -33,6 +33,9 @@ type Runner struct { templatesConfig *nucleiConfig // options contains configuration options for runner options *Options + + // progress tracking + progress *progress.Progress } // New creates a new client for running enumeration process. @@ -99,6 +102,10 @@ func New(options *Options) (*Runner, error) { } runner.output = output } + + // Creates the progress tracking object + runner.progress = progress.NewProgress(runner.options.NoColor) + return runner, nil } @@ -125,18 +132,18 @@ func (r *Runner) RunEnumeration() { r.options.Templates = newPath } + p := r.progress + // Single yaml provided if strings.HasSuffix(r.options.Templates, ".yaml") { t, err := r.parse(r.options.Templates) - // track progress - p := progress.NewProgress(nil) - switch t.(type) { case *templates.Template: var results bool template := t.(*templates.Template) + // track single template progress p.SetupTemplateProgressbar(-1, -1, template.ID, r.inputCount * template.GetHTTPRequestsCount()) // process http requests @@ -211,8 +218,7 @@ func (r *Runner) RunEnumeration() { } } - // track progress - p := progress.NewProgress(nil) + // track global progress p.SetupGlobalProgressbar(r.inputCount, len(matches), r.inputCount * totalRequests) var results bool @@ -222,6 +228,7 @@ func (r *Runner) RunEnumeration() { case *templates.Template: template := t.(*templates.Template) + // track template progress p.SetupTemplateProgressbar(i, totalTemplates, template.ID, r.inputCount * template.GetHTTPRequestsCount()) for _, request := range template.RequestsDNS { diff --git a/v2/pkg/workflows/var.go b/v2/pkg/workflows/var.go index 80791cca..ad61f525 100644 --- a/v2/pkg/workflows/var.go +++ b/v2/pkg/workflows/var.go @@ -35,7 +35,7 @@ func (n *NucleiVar) Call(args ...tengo.Object) (ret tengo.Object, err error) { var gotResult bool // track progress - p := progress.NewProgress(nil) + p := progress.NewProgress(false) for _, template := range n.Templates { if template.HTTPOptions != nil {