mirror of https://github.com/daffainfo/nuclei.git
Respect color output user flag, remove unused Progress WaitGroup
parent
e877d613e4
commit
d0174c047c
|
@ -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),
|
||||
),
|
||||
)
|
||||
|
|
|
@ -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 {
|
||||
|
|
|
@ -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 {
|
||||
|
|
Loading…
Reference in New Issue