daffainfo
/
nuclei
mirror of https://github.com/daffainfo/nuclei.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Use manual progressbar rendering to allow syncing with other output 

Refactor into an interface to return a no-op impl, avoid the need to
conditionally invoke any of the methods.
dev
Manuel Bua 2020-07-31 23:07:33 +02:00
parent 058c4ffbbf
commit 02238c491b
6 changed files with 155 additions and 143 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

161
v2/internal/progress/progress.go
View File

@ -3,21 +3,31 @@ package progress
import (
"fmt"
"github.com/logrusorgru/aurora"
"github.com/projectdiscovery/gologger"
"github.com/vbauerster/mpb/v5"
"github.com/vbauerster/mpb/v5/cwriter"
"github.com/vbauerster/mpb/v5/decor"
"io"
"os"
"regexp"
"strconv"
"strings"
"sync"
"time"
)
// Encapsulates progress tracking.
type IProgress interface {
setupProgressbar(name string, total int64, priority int) *mpb.Bar
InitProgressbar(hostCount int64, templateCount int, requestCount int64)
AddToTotal(delta int64)
Update()
render()
Drop(count int64)
Wait()
StartStdCapture()
StopStdCapture()
}
type Progress struct {
progress *mpb.Progress
gbar *mpb.Bar
bar *mpb.Bar
total int64
initialTotal int64
totalMutex *sync.Mutex
@ -26,28 +36,32 @@ type Progress struct {
stdout *strings.Builder
stderr *strings.Builder
colorizer aurora.Aurora
termWidth int
renderChan chan time.Time
renderMutex *sync.Mutex
firstTimeOutput bool
}
// 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.")
func NewProgress(noColor bool, active bool) IProgress {
if !active {
return &NoOpProgress{}
}
renderChan := make(chan time.Time)
p := &Progress{
progress: mpb.New(
mpb.WithOutput(os.Stderr),
mpb.PopCompletedMode(),
mpb.WithManualRefresh(renderChan),
),
totalMutex: &sync.Mutex{},
stdCaptureMutex: &sync.Mutex{},
stdout: &strings.Builder{},
stderr: &strings.Builder{},
colorizer: aurora.NewAurora(!noColor),
termWidth: tw,
renderChan: renderChan,
renderMutex: &sync.Mutex{},
firstTimeOutput: true,
}
return p
}
@ -55,7 +69,7 @@ func NewProgress(noColor bool) *Progress {
// 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 {
if p.bar != nil {
panic("A global progressbar is already present.")
}
@ -68,35 +82,29 @@ func (p *Progress) InitProgressbar(hostCount int64, templateCount int, requestCo
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
p.bar = p.setupProgressbar("["+barName+"]", requestCount, 0)
}
// Update total progress request count
func (p *Progress) AddToTotal(delta int64) {
p.totalMutex.Lock()
p.total += delta
p.gbar.SetTotal(p.total, false)
p.bar.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()
p.bar.Increment()
p.render()
}
// 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)
p.bar.IncrInt64(count)
p.render()
}
// Ensures that a progress bar's total count is up-to-date if during an enumeration there were uncompleted requests and
@ -104,14 +112,63 @@ func (p *Progress) Drop(count int64) {
func (p *Progress) Wait() {
p.totalMutex.Lock()
if p.total == 0 {
p.gbar.Abort(true)
p.bar.Abort(true)
} else if p.initialTotal != p.total {
p.gbar.SetTotal(p.total, true)
p.bar.SetTotal(p.total, true)
}
p.totalMutex.Unlock()
p.progress.Wait()
}
// 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())
p.renderMutex.Lock()
{
hasStdout := p.stdout.Len() > 0
hasStderr := p.stderr.Len() > 0
hasOutput := hasStdout || hasStderr
if hasOutput {
if p.firstTimeOutput {
// trigger a render event
p.renderChan <- time.Time{}
gologger.Infof("Waiting for your terminal to settle..")
// no way to sync to it? :(
time.Sleep(time.Millisecond * 250)
p.firstTimeOutput = false
}
// go back one line and clean it all
fmt.Fprint(os.Stderr, "\u001b[1A\u001b[2K")
if hasStdout {
fmt.Fprint(os.Stdout, p.stdout.String())
}
if hasStderr {
fmt.Fprint(os.Stderr, p.stderr.String())
}
// make space for the progressbar to render itself
fmt.Fprintln(os.Stderr, "")
// always trigger a render event to try ensure it's visible even with fast output
p.renderChan <- time.Time{}
}
}
p.renderMutex.Unlock()
p.stdout.Reset()
p.stderr.Reset()
p.stdCaptureMutex.Unlock()
}
// Creates and returns a progress bar.
func (p *Progress) setupProgressbar(name string, total int64, priority int) *mpb.Bar {
color := p.colorizer
@ -137,49 +194,15 @@ func (p *Progress) setupProgressbar(name string, total int64, priority int) *mpb
)
}
// 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()
func (p *Progress) render() {
p.renderMutex.Lock()
p.renderChan <- time.Now()
p.renderMutex.Unlock()
}
// 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)
}
func pluralize(count int64, singular, plural string) string {
if count > 1 {
return plural
}
p.stdout.Reset()
//
p.stdCaptureMutex.Unlock()
return singular
}
// 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)
})
}

15
v2/internal/progress/progress_noop.go Normal file
View File

@ -0,0 +1,15 @@
package progress
import "github.com/vbauerster/mpb/v5"
type NoOpProgress struct{}
func (p *NoOpProgress) setupProgressbar(name string, total int64, priority int) *mpb.Bar { return nil }
func (p *NoOpProgress) InitProgressbar(hostCount int64, templateCount int, requestCount int64) {}
func (p *NoOpProgress) AddToTotal(delta int64) {}
func (p *NoOpProgress) Update() {}
func (p *NoOpProgress) render() {}
func (p *NoOpProgress) Drop(count int64) {}
func (p *NoOpProgress) Wait() {}
func (p *NoOpProgress) StartStdCapture() {}
func (p *NoOpProgress) StopStdCapture() {}

32
v2/internal/runner/runner.go
View File

@ -43,7 +43,7 @@ type Runner struct {
limiter chan struct{}
// progress tracking
progress *progress.Progress
progress progress.IProgress
// output coloring
colorizer aurora.Aurora
@ -148,10 +148,8 @@ func New(options *Options) (*Runner, error) {
runner.output = output
}
if !options.Silent && !options.DisableProgressBar {
// Creates the progress tracking object
runner.progress = progress.NewProgress(runner.options.NoColor)
}
// Creates the progress tracking object
runner.progress = progress.NewProgress(runner.options.NoColor, !options.Silent && !options.DisableProgressBar)
runner.limiter = make(chan struct{}, options.Threads)
@ -339,9 +337,7 @@ func (r *Runner) RunEnumeration() {
} else if totalRequests > 0 || hasWorkflows {
// track global progress
if p != nil {
p.InitProgressbar(r.inputCount, templateCount, totalRequests)
}
p.InitProgressbar(r.inputCount, templateCount, totalRequests)
for _, match := range allTemplates {
wgtemplates.Add(1)
@ -367,13 +363,7 @@ func (r *Runner) RunEnumeration() {
}
wgtemplates.Wait()
if p != nil {
p.Wait()
p.ShowStdErr()
p.ShowStdOut()
}
p.Wait()
}
if !results.Get() {
@ -388,7 +378,7 @@ func (r *Runner) RunEnumeration() {
}
// processTemplateWithList processes a template and runs the enumeration on all the targets
func (r *Runner) processTemplateWithList(p *progress.Progress, template *templates.Template, request interface{}) bool {
func (r *Runner) processTemplateWithList(p progress.IProgress, template *templates.Template, request interface{}) bool {
// Display the message for the template
message := fmt.Sprintf("[%s] Loaded template %s (@%s)", template.ID, template.Info.Name, template.Info.Author)
if template.Info.Severity != "" {
@ -441,9 +431,7 @@ func (r *Runner) processTemplateWithList(p *progress.Progress, template *templat
})
}
if err != nil {
if p != nil {
p.Drop(request.(*requests.BulkHTTPRequest).GetRequestCount())
}
p.Drop(request.(*requests.BulkHTTPRequest).GetRequestCount())
p.StartStdCapture()
gologger.Warningf("Could not create http client: %s\n", err)
p.StopStdCapture()
@ -489,7 +477,7 @@ func (r *Runner) processTemplateWithList(p *progress.Progress, template *templat
}
// ProcessWorkflowWithList coming from stdin or list of targets
func (r *Runner) ProcessWorkflowWithList(p *progress.Progress, workflow *workflows.Workflow) {
func (r *Runner) ProcessWorkflowWithList(p progress.IProgress, workflow *workflows.Workflow) {
var wg sync.WaitGroup
scanner := bufio.NewScanner(strings.NewReader(r.input))
for scanner.Scan() {
@ -513,7 +501,7 @@ func (r *Runner) ProcessWorkflowWithList(p *progress.Progress, workflow *workflo
}
// ProcessWorkflow towards an URL
func (r *Runner) ProcessWorkflow(p *progress.Progress, workflow *workflows.Workflow, URL string) error {
func (r *Runner) ProcessWorkflow(p progress.IProgress, workflow *workflows.Workflow, URL string) error {
script := tengo.NewScript([]byte(workflow.Logic))
script.SetImports(stdlib.GetModuleMap(stdlib.AllModuleNames()...))
var jar *cookiejar.Jar
@ -641,7 +629,9 @@ func (r *Runner) ProcessWorkflow(p *progress.Progress, workflow *workflows.Workf
_, err := script.RunContext(context.Background())
if err != nil {
p.StartStdCapture()
gologger.Errorf("Could not execute workflow '%s': %s\n", workflow.ID, err)
p.StopStdCapture()
return err
}
return nil

48
v2/pkg/executer/executer_dns.go
View File

@ -29,9 +29,9 @@ type DNSExecuter struct {
writer *bufio.Writer
outputMutex *sync.Mutex
coloredOutput bool
colorizer aurora.Aurora
decolorizer *regexp.Regexp
coloredOutput bool
colorizer aurora.Aurora
decolorizer *regexp.Regexp
}
// DefaultResolvers contains the list of resolvers known to be trusted.
@ -50,9 +50,9 @@ type DNSOptions struct {
DNSRequest *requests.DNSRequest
Writer *bufio.Writer
ColoredOutput bool
Colorizer aurora.Aurora
Decolorizer *regexp.Regexp
ColoredOutput bool
Colorizer aurora.Aurora
Decolorizer *regexp.Regexp
}
// NewDNSExecuter creates a new DNS executer from a template
@ -61,22 +61,22 @@ func NewDNSExecuter(options *DNSOptions) *DNSExecuter {
dnsClient := retryabledns.New(DefaultResolvers, options.DNSRequest.Retries)
executer := &DNSExecuter{
debug: options.Debug,
jsonOutput: options.JSON,
dnsClient: dnsClient,
template: options.Template,
dnsRequest: options.DNSRequest,
writer: options.Writer,
outputMutex: &sync.Mutex{},
coloredOutput: options.ColoredOutput,
colorizer: options.Colorizer,
decolorizer: options.Decolorizer,
debug: options.Debug,
jsonOutput: options.JSON,
dnsClient: dnsClient,
template: options.Template,
dnsRequest: options.DNSRequest,
writer: options.Writer,
outputMutex: &sync.Mutex{},
coloredOutput: options.ColoredOutput,
colorizer: options.Colorizer,
decolorizer: options.Decolorizer,
}
return executer
}
// ExecuteDNS executes the DNS request on a URL
func (e *DNSExecuter) ExecuteDNS(p *progress.Progress, URL string) (result Result) {
func (e *DNSExecuter) ExecuteDNS(p progress.IProgress, URL string) (result Result) {
// Parse the URL and return domain if URL.
var domain string
if isURL(URL) {
@ -89,9 +89,7 @@ func (e *DNSExecuter) ExecuteDNS(p *progress.Progress, URL string) (result Resul
compiledRequest, err := e.dnsRequest.MakeDNSRequest(domain)
if err != nil {
result.Error = errors.Wrap(err, "could not make dns request")
if p != nil {
p.Drop(1)
}
p.Drop(1)
return
}
@ -106,17 +104,15 @@ func (e *DNSExecuter) ExecuteDNS(p *progress.Progress, URL string) (result Resul
resp, err := e.dnsClient.Do(compiledRequest)
if err != nil {
result.Error = errors.Wrap(err, "could not send dns request")
if p != nil {
p.Drop(1)
}
p.Drop(1)
return
}
if p != nil {
p.Update()
}
p.Update()
p.StartStdCapture()
gologger.Verbosef("Sent DNS request to %s\n", "dns-request", URL)
p.StopStdCapture()
if e.debug {
p.StartStdCapture()

28
v2/pkg/executer/executer_http.go
View File

@ -42,9 +42,9 @@ type HTTPExecuter struct {
customHeaders requests.CustomHeaders
CookieJar *cookiejar.Jar
coloredOutput bool
colorizer aurora.Aurora
decolorizer *regexp.Regexp
coloredOutput bool
colorizer aurora.Aurora
decolorizer *regexp.Regexp
}
// HTTPOptions contains configuration options for the HTTP executer.
@ -62,9 +62,9 @@ type HTTPOptions struct {
CustomHeaders requests.CustomHeaders
CookieReuse bool
CookieJar *cookiejar.Jar
ColoredOutput bool
Colorizer aurora.Aurora
Decolorizer *regexp.Regexp
ColoredOutput bool
Colorizer aurora.Aurora
Decolorizer *regexp.Regexp
}
// NewHTTPExecuter creates a new HTTP executer from a template
@ -113,7 +113,7 @@ func NewHTTPExecuter(options *HTTPOptions) (*HTTPExecuter, error) {
}
// ExecuteHTTP executes the HTTP request on a URL
func (e *HTTPExecuter) ExecuteHTTP(p *progress.Progress, URL string) (result Result) {
func (e *HTTPExecuter) ExecuteHTTP(p progress.IProgress, URL string) (result Result) {
result.Matches = make(map[string]interface{})
result.Extractions = make(map[string]interface{})
dynamicvalues := make(map[string]interface{})
@ -130,25 +130,19 @@ func (e *HTTPExecuter) ExecuteHTTP(p *progress.Progress, URL string) (result Res
httpRequest, err := e.bulkHttpRequest.MakeHTTPRequest(URL, dynamicvalues, e.bulkHttpRequest.Current(URL))
if err != nil {
result.Error = errors.Wrap(err, "could not build http request")
if p != nil {
p.Drop(remaining)
}
p.Drop(remaining)
return
}
err = e.handleHTTP(p, URL, httpRequest, dynamicvalues, &result)
if err != nil {
result.Error = errors.Wrap(err, "could not handle http request")
if p != nil {
p.Drop(remaining)
}
p.Drop(remaining)
return
}
e.bulkHttpRequest.Increment(URL)
if p != nil {
p.Update()
}
p.Update()
remaining--
}
@ -159,7 +153,7 @@ func (e *HTTPExecuter) ExecuteHTTP(p *progress.Progress, URL string) (result Res
return
}
func (e *HTTPExecuter) handleHTTP(p *progress.Progress, URL string, request *requests.HttpRequest, dynamicvalues map[string]interface{}, result *Result) error {
func (e *HTTPExecuter) handleHTTP(p progress.IProgress, URL string, request *requests.HttpRequest, dynamicvalues map[string]interface{}, result *Result) error {
e.setCustomHeaders(request)
req := request.Request

14
v2/pkg/workflows/var.go
View File

@ -24,7 +24,7 @@ type NucleiVar struct {
type Template struct {
HTTPOptions *executer.HTTPOptions
DNSOptions *executer.DNSOptions
Progress *progress.Progress
Progress progress.IProgress
}
// TypeName of the variable
@ -57,9 +57,7 @@ func (n *NucleiVar) Call(args ...tengo.Object) (ret tengo.Object, err error) {
for _, template := range n.Templates {
p := template.Progress
if template.HTTPOptions != nil {
if p != nil {
p.AddToTotal(template.HTTPOptions.Template.GetHTTPRequestCount())
}
p.AddToTotal(template.HTTPOptions.Template.GetHTTPRequestCount())
for _, request := range template.HTTPOptions.Template.BulkRequestsHTTP {
// apply externally supplied payloads if any
request.Headers = generators.MergeMapsWithStrings(request.Headers, headers)
@ -68,9 +66,7 @@ func (n *NucleiVar) Call(args ...tengo.Object) (ret tengo.Object, err error) {
template.HTTPOptions.BulkHttpRequest = request
httpExecuter, err := executer.NewHTTPExecuter(template.HTTPOptions)
if err != nil {
if p != nil {
p.Drop(request.GetRequestCount())
}
p.Drop(request.GetRequestCount())
p.StartStdCapture()
gologger.Warningf("Could not compile request for template '%s': %s\n", template.HTTPOptions.Template.ID, err)
p.StopStdCapture()
@ -92,9 +88,7 @@ func (n *NucleiVar) Call(args ...tengo.Object) (ret tengo.Object, err error) {
}
if template.DNSOptions != nil {
if p != nil {
p.AddToTotal(template.DNSOptions.Template.GetDNSRequestCount())
}
p.AddToTotal(template.DNSOptions.Template.GetDNSRequestCount())
for _, request := range template.DNSOptions.Template.RequestsDNS {
template.DNSOptions.DNSRequest = request
dnsExecuter := executer.NewDNSExecuter(template.DNSOptions)