nuclei/v2/internal/progress/progress.go

package progress

import (
	"fmt"
	"github.com/logrusorgru/aurora"
	"github.com/vbauerster/mpb/v5"
	"github.com/vbauerster/mpb/v5/cwriter"
	"github.com/vbauerster/mpb/v5/decor"
	"io"
	"os"
	"regexp"
	"strconv"
	"strings"
	"sync"
)

// Encapsulates progress tracking.
type Progress struct {
	progress        *mpb.Progress
	gbar            *mpb.Bar
	total           int64
	initialTotal    int64
	totalMutex      *sync.Mutex
	captureData     *captureData
	stdCaptureMutex *sync.Mutex
	stdout          *strings.Builder
	stderr          *strings.Builder
	colorizer       aurora.Aurora
	termWidth int
}

// Creates and returns a new progress tracking object.
func NewProgress(noColor bool) *Progress {
	w := cwriter.New(os.Stdout)
	tw, err := w.GetWidth()
	if err != nil {
		panic("Couldn't determine available terminal width.")
	}

	p := &Progress{
		progress: mpb.New(
			mpb.WithOutput(os.Stderr),
			mpb.PopCompletedMode(),
		),
		totalMutex:      &sync.Mutex{},
		stdCaptureMutex: &sync.Mutex{},
		stdout:          &strings.Builder{},
		stderr:          &strings.Builder{},
		colorizer:       aurora.NewAurora(!noColor),
		termWidth: tw,
	}
	return p
}

// 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) InitProgressbar(hostCount int64, templateCount int, requestCount int64) {
	if p.gbar != nil {
		panic("A global progressbar is already present.")
	}

	color := p.colorizer

	barName := color.Sprintf(
		color.Cyan("%d %s, %d %s"),
		color.Bold(color.Cyan(templateCount)),
		pluralize(int64(templateCount), "template", "templates"),
		color.Bold(color.Cyan(hostCount)),
		pluralize(hostCount, "host", "hosts"))

	p.gbar = p.setupProgressbar("["+barName+"]", requestCount, 0)
}

func pluralize(count int64, singular, plural string) string {
	if count > 1 {
		return plural
	}
	return singular
}

// Update total progress request count
func (p *Progress) AddToTotal(delta int64) {
	p.totalMutex.Lock()
	p.total += delta
	p.gbar.SetTotal(p.total, false)
	p.totalMutex.Unlock()
}

// Update progress tracking information and increments the request counter by one unit.
func (p *Progress) Update() {
	p.gbar.Increment()
}

// 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) {
	// mimic dropping by incrementing the completed requests
	p.gbar.IncrInt64(count)

}

// Ensures that a progress bar's total count is up-to-date if during an enumeration there were uncompleted requests and
// wait for all the progress bars to finish.
func (p *Progress) Wait() {
	p.totalMutex.Lock()
	if p.total == 0 {
		p.gbar.Abort(true)
	} else if p.initialTotal != p.total {
		p.gbar.SetTotal(p.total, true)
	}
	p.totalMutex.Unlock()
	p.progress.Wait()
}

// Creates and returns a progress bar.
func (p *Progress) setupProgressbar(name string, total int64, priority int) *mpb.Bar {
	color := p.colorizer

	p.total = total
	p.initialTotal = total

	return p.progress.AddBar(
		total,
		mpb.BarPriority(priority),
		mpb.BarNoPop(),
		mpb.BarRemoveOnComplete(),
		mpb.PrependDecorators(
			decor.Name(name, decor.WCSyncSpaceR),
			decor.CountersNoUnit(color.BrightBlue(" %d/%d").String(), decor.WCSyncSpace),
			decor.NewPercentage(color.Bold("%d").String(), decor.WCSyncSpace),
		),
		mpb.AppendDecorators(
			decor.AverageSpeed(0, color.BrightYellow("%.2f").Bold().String()+color.BrightYellow("r/s").String(), decor.WCSyncSpace),
			decor.Elapsed(decor.ET_STYLE_GO, decor.WCSyncSpace),
			decor.AverageETA(decor.ET_STYLE_GO, decor.WCSyncSpace),
		),
	)
}

// Starts capturing stdout and stderr instead of producing visual output that may interfere with the progress bars.
func (p *Progress) StartStdCapture() {
	p.stdCaptureMutex.Lock()
	p.captureData = startStdCapture()
}

// Stops capturing stdout and stderr and store both output to be shown later.
func (p *Progress) StopStdCapture() {
	stopStdCapture(p.captureData)
	p.stdout.Write(p.captureData.DataStdOut.Bytes())
	p.stderr.Write(p.captureData.DataStdErr.Bytes())

	//
	var r = regexp.MustCompile("(.{" + strconv.Itoa(p.termWidth) + "})")
	multiline := r.ReplaceAllString(p.stdout.String(), "$1\n")
	arr := strings.Split(multiline, "\n")
	for _, msg := range arr {
		if len(msg) > 0 {
			p.progress.Add(0, makeLogBar(msg)).SetTotal(0, true)
		}
	}
	p.stdout.Reset()
	//

	p.stdCaptureMutex.Unlock()
}

// Writes the captured stdout data to stdout, if any.
func (p *Progress) ShowStdOut() {
	//if p.stdout.Len() > 0 {
	//	fmt.Fprint(os.Stdout, p.stdout.String())
	//}
}

// Writes the captured stderr data to stderr, if any.
func (p *Progress) ShowStdErr() {
	//if p.stderr.Len() > 0 {
	//	fmt.Fprint(os.Stderr, p.stderr.String())
	//}
}

func makeLogBar(msg string) mpb.BarFiller {
	return mpb.BarFillerFunc(func(w io.Writer, _ int, st decor.Statistics) {
		fmt.Fprintf(w, msg)
	})
}