Minor refactoring, track count of aborted requests for better progress

v2/internal/progress/progress.go

@@ -17,6 +17,8 @@ import (
 type Progress struct {
 	progress *mpb.Progress
 	bar *mpb.Bar
+	total int64
+	initialTotal int64
 	captureData *captureData
 	termWidth int
 }
@@ -59,6 +61,8 @@ func (p *Progress) SetupProgressBar(name string, total int64) *mpb.Bar {
 	)
 
 	p.bar = bar
+	p.total = total
+	p.initialTotal = total
 	return bar
 }
 
@@ -66,11 +70,16 @@ func (p *Progress) Update() {
 	p.bar.Increment()
 }
 
-func (p *Progress) Abort() {
-	p.bar.Abort(true)
+func (p *Progress) Abort(remaining int64) {
+	p.total -= remaining
+	p.bar.SetTotal(p.total, false)
 }
 
 func (p *Progress) Wait() {
+	if p.initialTotal != p.total {
+		p.bar.SetTotal(p.total, true)
+	}
+
 	p.progress.Wait()
 }
 

v2/internal/runner/runner.go

@@ -109,14 +109,6 @@ func (r *Runner) Close() {
 	os.Remove(r.tempFile)
 }
 
-func (r *Runner) getHTTPRequestsCount(t *templates.Template) int64 {
-	var count int64 = 0
-	for _, request := range t.RequestsHTTP {
-		count += request.GetRequestCount()
-	}
-	return count
-}
-
 // RunEnumeration sets up the input layer for giving input nuclei.
 // binary and runs the actual enumeration
 func (r *Runner) RunEnumeration() {
@@ -133,18 +125,19 @@ 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)
+
+		// track progress
+		p := progress.NewProgress(nil)
+
 		switch t.(type) {
 		case *templates.Template:
 			var results bool
 			template := t.(*templates.Template)
 
-			p.SetupProgressBar(template.ID, r.inputCount * r.getHTTPRequestsCount(template))
+			p.SetupProgressBar(template.ID, r.inputCount * template.GetHTTPRequestsCount())
 
 			// process http requests
 			for _, request := range template.RequestsHTTP {
@@ -159,6 +152,8 @@ func (r *Runner) RunEnumeration() {
 				}
 			}
 
+			p.Wait()
+
 			if !results {
 				if r.output != nil {
 					outputFile := r.output.Name()
@@ -200,6 +195,9 @@ func (r *Runner) RunEnumeration() {
 		gologger.Fatalf("Error, no templates found in directory: '%s'\n", r.options.Templates)
 	}
 
+	// track progress
+	p := progress.NewProgress(nil)
+
 	var results bool
 	for _, match := range matches {
 		t, err := r.parse(match)
@@ -213,7 +211,7 @@ func (r *Runner) RunEnumeration() {
 				}
 			}
 
-			p.SetupProgressBar(template.ID, r.inputCount * r.getHTTPRequestsCount(template))
+			p.SetupProgressBar(template.ID, r.inputCount * template.GetHTTPRequestsCount())
 
 			for _, request := range template.RequestsHTTP {
 				httpResults := r.processTemplateWithList(p, template, request)
@@ -231,18 +229,19 @@ func (r *Runner) RunEnumeration() {
 			p.StopStdCaptureAndShow()
 		}
 	}
+	p.Wait()
+
 	if !results {
 		if r.output != nil {
 			outputFile := r.output.Name()
 			r.output.Close()
 			os.Remove(outputFile)
 		}
-		p.StartStdCapture()
+		//p.StartStdCapture()
 		gologger.Infof("No results found for the template. Happy hacking!")
-		p.StopStdCaptureAndShow()
+		//p.StopStdCaptureAndShow()
 	}
 
-	p.Wait()
 	return
 }
 

v2/pkg/executor/executer_http.go

@@ -101,11 +101,13 @@ func (e *HTTPExecutor) ExecuteHTTP(p *progress.Progress, URL string) error {
 		return errors.Wrap(err, "could not compile http request")
 	}
 
+	remaining := e.template.GetHTTPRequestsCount()
+
 	// Send the request to the target servers
 mainLoop:
 	for compiledRequest := range compiledRequest {
 		if compiledRequest.Error != nil {
-			p.Abort()
+			p.Abort(remaining)
 			return errors.Wrap(err, "error in compiled http request")
 		}
 		e.setCustomHeaders(compiledRequest)
@@ -118,7 +120,7 @@ mainLoop:
 
 			dumpedRequest, err := httputil.DumpRequest(req.Request, true)
 			if err != nil {
-				p.Abort()
+				p.Abort(remaining)
 				return errors.Wrap(err, "could not dump http request")
 			}
 			p.StartStdCapture()
@@ -131,7 +133,7 @@ mainLoop:
 			if resp != nil {
 				resp.Body.Close()
 			}
-			p.Abort()
+			p.Abort(remaining)
 			return errors.Wrap(err, "could not issue http request")
 		}
 
@@ -142,7 +144,7 @@ mainLoop:
 
 			dumpedResponse, err := httputil.DumpResponse(resp, true)
 			if err != nil {
-				p.Abort()
+				p.Abort(remaining)
 				return errors.Wrap(err, "could not dump http response")
 			}
 			p.StartStdCapture()
@@ -154,7 +156,7 @@ mainLoop:
 		if err != nil {
 			io.Copy(ioutil.Discard, resp.Body)
 			resp.Body.Close()
-			p.Abort()
+			p.Abort(remaining)
 			return errors.Wrap(err, "could not read http body")
 		}
 		resp.Body.Close()
@@ -163,7 +165,7 @@ mainLoop:
 		// so in case we have to manually do it
 		data, err = requests.HandleDecompression(compiledRequest.Request, data)
 		if err != nil {
-			p.Abort()
+			p.Abort(remaining)
 			return errors.Wrap(err, "could not decompress http body")
 		}
 
@@ -184,6 +186,7 @@ mainLoop:
 				// If the condition is AND we haven't matched, try next request.
 				if matcherCondition == matchers.ANDCondition {
 					p.Update()
+					remaining--
 					continue mainLoop
 				}
 			} else {
@@ -225,6 +228,7 @@ mainLoop:
 		}
 
 		p.Update()
+		remaining--
 	}
 
 	p.StartStdCapture()

v2/pkg/templates/templates.go

@@ -25,3 +25,11 @@ type Info struct {
 	// Severity optionally describes the severity of the template
 	Severity string `yaml:"severity,omitempty"`
 }
+
+func (t* Template) GetHTTPRequestsCount() int64 {
+	var count int64 = 0
+	for _, request := range t.RequestsHTTP {
+		count += request.GetRequestCount()
+	}
+	return count
+}