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

Merge pull request #196 from manuelbua/fix-186-experimental-progressbar-live-results 

Fix 186 - Experimental progressbar live results
dev
bauthard 2020-08-02 13:28:03 +05:30 committed by GitHub
parent 2685fdd1b9 87dc8bc650
commit 858168b9a8
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
7 changed files with 218 additions and 178 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -3,47 +3,79 @@ package progress
import (
"fmt"
"github.com/logrusorgru/aurora"
"github.com/projectdiscovery/gologger"
"github.com/vbauerster/mpb/v5"
"github.com/vbauerster/mpb/v5/decor"
"io"
"os"
"strings"
"sync"
"time"
)
// global output refresh rate
const RefreshHz = 8
// Encapsulates progress tracking.
type IProgress interface {
InitProgressbar(hostCount int64, templateCount int, requestCount int64)
AddToTotal(delta int64)
Update()
Drop(count int64)
Wait()
}
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
progress *mpb.Progress
bar *mpb.Bar
total int64
initialTotal int64
totalMutex *sync.Mutex
colorizer aurora.Aurora
renderChan chan time.Time
captureData *captureData
stdCaptureMutex *sync.Mutex
stdOut *strings.Builder
stdErr *strings.Builder
stdStopRenderEvent chan bool
stdRenderEvent *time.Ticker
stdRenderWaitGroup *sync.WaitGroup
}
// Creates and returns a new progress tracking object.
func NewProgress(noColor bool) *Progress {
func NewProgress(noColor bool, active bool) IProgress {
if !active {
return &NoOpProgress{}
}
refreshMillis := int64(1. / float64(RefreshHz) * 1000.)
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),
totalMutex: &sync.Mutex{},
colorizer: aurora.NewAurora(!noColor),
renderChan: renderChan,
stdCaptureMutex: &sync.Mutex{},
stdOut: &strings.Builder{},
stdErr: &strings.Builder{},
stdStopRenderEvent: make(chan bool),
stdRenderEvent: time.NewTicker(time.Millisecond * time.Duration(refreshMillis)),
stdRenderWaitGroup: &sync.WaitGroup{},
}
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.
// Creates and returns a progress bar that tracks all the progress.
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.")
}
@ -56,35 +88,33 @@ 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)
}
p.bar = p.setupProgressbar("["+barName+"]", requestCount, 0)
func pluralize(count int64, singular, plural string) string {
if count > 1 {
return plural
}
return singular
// creates r/w pipes and divert stdout/stderr writers to them and start capturing their output
p.captureData = startCapture(p.stdCaptureMutex, p.stdOut, p.stdErr)
// starts rendering both the progressbar and the captured stdout/stderr data
p.renderStdData()
}
// 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()
}
// 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)
}
// Ensures that a progress bar's total count is up-to-date if during an enumeration there were uncompleted requests and
@ -92,12 +122,65 @@ 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()
// close the writers and wait for the EOF condition
stopCapture(p.captureData)
// stop the renderer and wait for it
p.stdStopRenderEvent <- true
p.stdRenderWaitGroup.Wait()
// drain any stdout/stderr data
p.drainStringBuilderTo(p.stdOut, os.Stdout)
p.drainStringBuilderTo(p.stdErr, os.Stderr)
}
func (p *Progress) renderStdData() {
// trigger a render event
p.renderChan <- time.Now()
gologger.Infof("Waiting for your terminal to settle..")
time.Sleep(time.Millisecond * 250)
p.stdRenderWaitGroup.Add(1)
go func(waitGroup *sync.WaitGroup) {
for {
select {
case <-p.stdStopRenderEvent:
waitGroup.Done()
return
case _ = <-p.stdRenderEvent.C:
p.stdCaptureMutex.Lock()
{
hasStdout := p.stdOut.Len() > 0
hasStderr := p.stdErr.Len() > 0
hasOutput := hasStdout || hasStderr
if hasOutput {
stdout := p.captureData.backupStdout
stderr := p.captureData.backupStderr
// go back one line and clean it all
fmt.Fprint(stderr, "\u001b[1A\u001b[2K")
p.drainStringBuilderTo(p.stdOut, stdout)
p.drainStringBuilderTo(p.stdErr, stderr)
// make space for the progressbar to render itself
fmt.Fprintln(stderr, "")
}
// always trigger a render event to try ensure it's visible even with fast output
p.renderChan <- time.Now()
}
p.stdCaptureMutex.Unlock()
}
}
}(p.stdRenderWaitGroup)
}
// Creates and returns a progress bar.
@ -125,30 +208,16 @@ 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 pluralize(count int64, singular, plural string) string {
if count > 1 {
return plural
}
return singular
}
// 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.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 (p *Progress) drainStringBuilderTo(builder *strings.Builder, writer io.Writer) {
if builder.Len() > 0 {
fmt.Fprint(writer, builder.String())
builder.Reset()
}
}

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

@ -0,0 +1,9 @@
package progress
type NoOpProgress struct{}
func (p *NoOpProgress) InitProgressbar(hostCount int64, templateCount int, requestCount int64) {}
func (p *NoOpProgress) AddToTotal(delta int64) {}
func (p *NoOpProgress) Update() {}
func (p *NoOpProgress) Drop(count int64) {}
func (p *NoOpProgress) Wait() {}

86
v2/internal/progress/stdcapture.go
View File

@ -2,28 +2,25 @@ package progress
/**
Inspired by the https://github.com/PumpkinSeed/cage module
*/
*/
import (
"bytes"
"bufio"
"github.com/projectdiscovery/gologger"
"io"
"os"
"strings"
"sync"
)
type captureData struct {
backupStdout *os.File
writerStdout *os.File
backupStderr *os.File
writerStderr *os.File
DataStdOut *bytes.Buffer
DataStdErr *bytes.Buffer
outStdout chan []byte
outStderr chan []byte
backupStdout *os.File
writerStdout *os.File
backupStderr *os.File
writerStderr *os.File
waitFinishRead *sync.WaitGroup
}
func startStdCapture() *captureData {
func startCapture(writeMutex *sync.Mutex, stdout *strings.Builder, stderr *strings.Builder) *captureData {
rStdout, wStdout, errStdout := os.Pipe()
if errStdout != nil {
panic(errStdout)
@ -41,54 +38,51 @@ func startStdCapture() *captureData {
backupStderr: os.Stderr,
writerStderr: wStderr,
outStdout: make(chan []byte),
outStderr: make(chan []byte),
DataStdOut: &bytes.Buffer{},
DataStdErr: &bytes.Buffer{},
waitFinishRead: &sync.WaitGroup{},
}
os.Stdout = c.writerStdout
os.Stderr = c.writerStderr
stdCopy := func(out chan<- []byte, reader *os.File) {
var buffer bytes.Buffer
_, _ = io.Copy(&buffer, reader)
if buffer.Len() > 0 {
out <- buffer.Bytes()
stdCopy := func(builder *strings.Builder, reader *os.File, waitGroup *sync.WaitGroup) {
r := bufio.NewReader(reader)
buf := make([]byte, 0, 4*1024)
for {
n, err := r.Read(buf[:cap(buf)])
buf = buf[:n]
if n == 0 {
if err == nil {
continue
}
if err == io.EOF {
waitGroup.Done()
break
}
waitGroup.Done()
gologger.Fatalf("stdcapture error: %s", err)
}
if err != nil && err != io.EOF {
waitGroup.Done()
gologger.Fatalf("stdcapture error: %s", err)
}
writeMutex.Lock()
builder.Write(buf)
writeMutex.Unlock()
}
close(out)
}
go stdCopy(c.outStdout, rStdout)
go stdCopy(c.outStderr, rStderr)
c.waitFinishRead.Add(2)
go stdCopy(stdout, rStdout, c.waitFinishRead)
go stdCopy(stderr, rStderr, c.waitFinishRead)
return c
}
func stopStdCapture(c *captureData) {
func stopCapture(c *captureData) {
_ = c.writerStdout.Close()
_ = c.writerStderr.Close()
var wg sync.WaitGroup
stdRead := func(in <-chan []byte, outData *bytes.Buffer) {
defer wg.Done()
for {
out, more := <-in
if more {
outData.Write(out)
} else {
return
}
}
}
wg.Add(2)
go stdRead(c.outStdout, c.DataStdOut)
go stdRead(c.outStderr, c.DataStdErr)
wg.Wait()
c.waitFinishRead.Wait()
os.Stdout = c.backupStdout
os.Stderr = c.backupStderr

34
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.EnableProgressBar {
// 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.EnableProgressBar)
runner.limiter = make(chan struct{}, options.Threads)
@ -338,11 +336,8 @@ func (r *Runner) RunEnumeration() {
gologger.Errorf("Could not find any valid input URLs.")
} else if totalRequests > 0 || hasWorkflows {
// track global progress
if p != nil {
p.InitProgressbar(r.inputCount, templateCount, totalRequests)
p.StartStdCapture()
}
// tracks global progress and captures stdout/stderr until p.Wait finishes
p.InitProgressbar(r.inputCount, templateCount, totalRequests)
for _, match := range allTemplates {
wgtemplates.Add(1)
@ -368,14 +363,7 @@ func (r *Runner) RunEnumeration() {
}
wgtemplates.Wait()
if p != nil {
p.Wait()
p.StopStdCapture()
p.ShowStdErr()
p.ShowStdOut()
}
p.Wait()
}
if !results.Get() {
@ -390,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 +429,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())
gologger.Warningf("Could not create http client: %s\n", err)
return false
}
@ -485,7 +471,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() {
@ -507,7 +493,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

46
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
}
@ -104,15 +102,11 @@ 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()
gologger.Verbosef("Sent DNS request to %s\n", "dns-request", URL)

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--
}
@ -157,7 +151,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())
gologger.Warningf("Could not compile request for template '%s': %s\n", template.HTTPOptions.Template.ID, err)
continue
}
@ -88,9 +84,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)