Merge pull request #196 from manuelbua/fix-186-experimental-progressb…

…ar-live-results

Fix 186 - Experimental progressbar live results

v2/internal/progress/progress.go

@@ -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,48 +88,99 @@ 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
 // wait for all the progress bars to finish.
 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()
-}
-
-// 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())
+func pluralize(count int64, singular, plural string) string {
+	if count > 1 {
+		return plural
 	}
+	return singular
 }
 
-// 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()
 	}
 }

v2/internal/progress/progress_noop.go

@@ -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()                                                                  {}

v2/internal/progress/stdcapture.go

@@ -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