Experimental single progress bar

v2/internal/progress/progress.go

@@ -16,6 +16,7 @@ import (
 
 type Progress struct {
 	progress *mpb.Progress
+	Bar *mpb.Bar
 	captureData *captureData
 	termWidth int
 }
@@ -34,16 +35,20 @@ func NewProgress(group *sync.WaitGroup) *Progress {
 			mpb.PopCompletedMode(),
 		),
 		termWidth: tw,
+		Bar: nil,
 	}
 	return p
 }
 
-func (p *Progress) NewBar(name string, total int64, URL string) *mpb.Bar {
+func (p *Progress) NewBar(name string, total int64) *mpb.Bar {
+	barname := "[" + aurora.Green(name).String() + "]"
+
 	return p.progress.AddBar(
 		total,
 		mpb.BarNoPop(),
+		//mpb.BarQueueAfter(p.Bar),
 		mpb.PrependDecorators(
-			decor.Name("[" + aurora.Green(URL).String() + " / " + aurora.Magenta(name).String() + "]"),
+			decor.Name(barname),
 			decor.CountersNoUnit(aurora.Blue(" %d/%d").String()),
 			decor.NewPercentage(aurora.Bold("%d").String(), decor.WCSyncSpace),
 		),

v2/internal/runner/runner.go

@@ -133,6 +133,9 @@ func (r *Runner) RunEnumeration() {
 		r.options.Templates = newPath
 	}
 
+	// track progress
+	p := progress.NewProgress(nil)
+
 	// Single yaml provided
 	if strings.HasSuffix(r.options.Templates, ".yaml") {
 		t, err := r.parse(r.options.Templates)
@@ -140,13 +143,17 @@ func (r *Runner) RunEnumeration() {
 		case *templates.Template:
 			var results bool
 			template := t.(*templates.Template)
+
+			p.Bar = p.NewBar(template.ID, r.inputCount * r.getHTTPRequestsCount(template))
+
 			// process http requests
 			for _, request := range template.RequestsHTTP {
-				results = r.processTemplateRequest(template, request)
+				results = r.processTemplateWithList(p, template, request)
 			}
+
 			// process dns requests
 			for _, request := range template.RequestsDNS {
-				dnsResults := r.processTemplateRequest(template, request)
+				dnsResults := r.processTemplateWithList(p, template, request)
 				if !results {
 					results = dnsResults
 				}
@@ -200,22 +207,28 @@ func (r *Runner) RunEnumeration() {
 		case *templates.Template:
 			template := t.(*templates.Template)
 			for _, request := range template.RequestsDNS {
-				dnsResults := r.processTemplateRequest(template, request)
+				dnsResults := r.processTemplateWithList(p, template, request)
 				if dnsResults {
 					results = dnsResults
 				}
 			}
+
+			p.Bar = p.NewBar(template.ID, r.inputCount * r.getHTTPRequestsCount(template))
+
 			for _, request := range template.RequestsHTTP {
-				httpResults := r.processTemplateRequest(template, request)
+				httpResults := r.processTemplateWithList(p, template, request)
 				if httpResults {
 					results = httpResults
 				}
 			}
+
 		case *workflows.Workflow:
 			workflow := t.(*workflows.Workflow)
 			r.ProcessWorkflowWithList(workflow)
 		default:
+			p.StartStdCapture()
 			gologger.Errorf("Could not parse file '%s': %s\n", r.options.Templates, err)
+			p.StopStdCaptureAndShow()
 		}
 	}
 	if !results {
@@ -224,9 +237,13 @@ func (r *Runner) RunEnumeration() {
 			r.output.Close()
 			os.Remove(outputFile)
 		}
+		p.StartStdCapture()
 		gologger.Infof("No results found for the template. Happy hacking!")
+		p.StopStdCaptureAndShow()
 	}
 
+	p.Wait()
+	return
 }
 
 // processTemplateWithList processes a template and runs the enumeration on all the targets
@@ -236,7 +253,9 @@ func (r *Runner) processTemplateWithList(p *progress.Progress, template *templat
 	if template.Info.Severity != "" {
 		message += " [" + template.Info.Severity + "]"
 	}
+	p.StartStdCapture()
 	gologger.Infof("%s\n", message)
+	p.StopStdCaptureAndShow()
 
 	var writer *bufio.Writer
 	if r.output != nil {
@@ -273,7 +292,9 @@ func (r *Runner) processTemplateWithList(p *progress.Progress, template *templat
 		})
 	}
 	if err != nil {
+		p.StartStdCapture()
 		gologger.Warningf("Could not create http client: %s\n", err)
+		p.StopStdCaptureAndShow()
 		return false
 	}
 
@@ -310,8 +331,7 @@ func (r *Runner) processTemplateWithList(p *progress.Progress, template *templat
 	}
 	close(limiter)
 
-	// Wait for both the WaitGroup and all the bars to complete
-	p.Wait()
+	wg.Wait()
 
 	// See if we got any results from the executors
 	var results bool

v2/pkg/executor/executer_http.go

@@ -101,16 +101,13 @@ func (e *HTTPExecutor) ExecuteHTTP(p *progress.Progress, URL string) error {
 		return errors.Wrap(err, "could not make http request")
 	}
 
-	// track progress
-	bar := p.NewBar(e.template.ID, e.GetRequestCount(), URL)
-
 	// Send the request to the target servers
 mainLoop:
 	for compiledRequest := range compiledRequest {
 		start := time.Now()
 
 		if compiledRequest.Error != nil {
-			bar.Abort(true)
+			p.Bar.Abort(true)
 			return errors.Wrap(err, "could not make http request")
 		}
 		e.setCustomHeaders(compiledRequest)
@@ -123,7 +120,7 @@ mainLoop:
 
 			dumpedRequest, err := httputil.DumpRequest(req.Request, true)
 			if err != nil {
-				bar.Abort(true)
+				p.Bar.Abort(true)
 				return errors.Wrap(err, "could not dump http request")
 			}
 			p.StartStdCapture()
@@ -136,7 +133,7 @@ mainLoop:
 			if resp != nil {
 				resp.Body.Close()
 			}
-			bar.Abort(true)
+			p.Bar.Abort(true)
 			return errors.Wrap(err, "could not make http request")
 		}
 
@@ -147,7 +144,7 @@ mainLoop:
 
 			dumpedResponse, err := httputil.DumpResponse(resp, true)
 			if err != nil {
-				bar.Abort(true)
+				p.Bar.Abort(true)
 				return errors.Wrap(err, "could not dump http response")
 			}
 			p.StartStdCapture()
@@ -159,7 +156,7 @@ mainLoop:
 		if err != nil {
 			io.Copy(ioutil.Discard, resp.Body)
 			resp.Body.Close()
-			bar.Abort(true)
+			p.Bar.Abort(true)
 			return errors.Wrap(err, "could not read http body")
 		}
 		resp.Body.Close()
@@ -168,7 +165,7 @@ mainLoop:
 		// so in case we have to manually do it
 		data, err = requests.HandleDecompression(compiledRequest.Request, data)
 		if err != nil {
-			bar.Abort(true)
+			p.Bar.Abort(true)
 			return errors.Wrap(err, "could not decompress http body")
 		}
 
@@ -188,8 +185,8 @@ mainLoop:
 			if !matcher.Match(resp, body, headers) {
 				// If the condition is AND we haven't matched, try next request.
 				if matcherCondition == matchers.ANDCondition {
-					bar.IncrBy(1)
-					bar.DecoratorEwmaUpdate(time.Since(start))
+					p.Bar.IncrBy(1)
+					p.Bar.DecoratorEwmaUpdate(time.Since(start))
 					continue mainLoop
 				}
 			} else {
@@ -230,8 +227,8 @@ mainLoop:
 			atomic.CompareAndSwapUint32(&e.results, 0, 1)
 		}
 
-		bar.Increment()
-		bar.DecoratorEwmaUpdate(time.Since(start))
+		p.Bar.Increment()
+		p.Bar.DecoratorEwmaUpdate(time.Since(start))
 	}
 
 	p.StartStdCapture()