Added cloud scan progress tracking using stats (#3180)

* Added cloud scan progress tracking using stats

* Changed log messsage

* Fixed linting error

* Fixed bug in progress calculation logic

* Changed requests to input with cloud flag

* Changed progress name + removed redundant fields

v2/go.mod

@@ -77,7 +77,7 @@ require (
 	github.com/projectdiscovery/ratelimit v0.0.4
 	github.com/projectdiscovery/rdap v0.9.1-0.20221108103045-9865884d1917
 	github.com/projectdiscovery/sarif v0.0.1
-	github.com/projectdiscovery/tlsx v1.0.2
+	github.com/projectdiscovery/tlsx v1.0.1
 	github.com/projectdiscovery/uncover v1.0.2
 	github.com/projectdiscovery/utils v0.0.4-0.20230104145529-50cace956b0a
 	github.com/projectdiscovery/wappalyzergo v0.0.77
@@ -111,6 +111,8 @@ require (
 	github.com/prometheus/common v0.39.0 // indirect
 	github.com/prometheus/procfs v0.9.0 // indirect
 	github.com/skeema/knownhosts v1.1.0 // indirect
+	github.com/spacemonkeygo/openssl v0.0.0-20181017203307-c2dcc5cca94a // indirect
+	github.com/spacemonkeygo/spacelog v0.0.0-20180420211403-2296661a0572 // indirect
 	github.com/tidwall/btree v1.4.3 // indirect
 	github.com/tidwall/buntdb v1.2.10 // indirect
 	github.com/tidwall/gjson v1.14.3 // indirect
@@ -189,8 +191,6 @@ require (
 	github.com/rivo/uniseg v0.2.0 // indirect
 	github.com/rogpeppe/go-internal v1.9.0 // indirect
 	github.com/saintfish/chardet v0.0.0-20120816061221-3af4cd4741ca // indirect
-	github.com/spacemonkeygo/openssl v0.0.0-20181017203307-c2dcc5cca94a // indirect
-	github.com/spacemonkeygo/spacelog v0.0.0-20180420211403-2296661a0572 // indirect
 	github.com/tklauser/go-sysconf v0.3.11 // indirect
 	github.com/tklauser/numcpus v0.6.0 // indirect
 	github.com/trivago/tgo v1.0.7 // indirect

v2/go.sum

@@ -588,8 +588,8 @@ github.com/projectdiscovery/sliceutil v0.0.1 h1:YoCqCMcdwz+gqNfW5hFY8UvNHoA6SfyB
 github.com/projectdiscovery/sliceutil v0.0.1/go.mod h1:0wBmhU5uTDwMfrEZfvwH9qa5k60Q4shPVOC9E6LGsDI=
 github.com/projectdiscovery/stringsutil v0.0.2 h1:uzmw3IVLJSMW1kEg8eCStG/cGbYYZAja8BH3LqqJXMA=
 github.com/projectdiscovery/stringsutil v0.0.2/go.mod h1:EJ3w6bC5fBYjVou6ryzodQq37D5c6qbAYQpGmAy+DC0=
-github.com/projectdiscovery/tlsx v1.0.2 h1:2bbfPQLuMIhs6FPmGsIcAo3uJaB2E+9ssJtZ8ZPb970=
-github.com/projectdiscovery/tlsx v1.0.2/go.mod h1:WW+PdBImrqnMl18v4Brp3OsbnO4A1tqYPUcfiVtjNLM=
+github.com/projectdiscovery/tlsx v1.0.1 h1:rpWDka7yGYV8LFLMCdPMkXmvqN8RJMqbE6pJArBCbJA=
+github.com/projectdiscovery/tlsx v1.0.1/go.mod h1:WW+PdBImrqnMl18v4Brp3OsbnO4A1tqYPUcfiVtjNLM=
 github.com/projectdiscovery/uncover v1.0.2 h1:mRFzflYyvwKkHd3XKufMlDRrb6p1mjFZTSHoNAUpFwo=
 github.com/projectdiscovery/uncover v1.0.2/go.mod h1:lz4QYfArSA6jJkXyB71kN2/Pc7IW7nJB8c95n7xtwqY=
 github.com/projectdiscovery/utils v0.0.4-0.20230104145529-50cace956b0a h1:fHztw99lR4QO931no6Zsj8/RYGA4otFQH5BF8OqfTss=

v2/internal/runner/enumerate.go

@@ -2,6 +2,7 @@ package runner
 
 import (
 	"bytes"
+	"context"
 	"crypto/sha1"
 	"encoding/base64"
 	"encoding/hex"
@@ -13,6 +14,7 @@ import (
 	"time"
 
 	"github.com/klauspost/compress/zlib"
+	"github.com/pkg/errors"
 	"github.com/projectdiscovery/gologger"
 	"github.com/projectdiscovery/nuclei/v2/internal/runner/nucleicloud"
 	"github.com/projectdiscovery/nuclei/v2/pkg/catalog/loader"
@@ -84,7 +86,34 @@ func (r *Runner) runCloudEnumeration(store *loader.Store, cloudTemplates, cloudT
 	}
 	time.Sleep(3 * time.Second)
 
+	scanResponse, err := r.cloudClient.GetScan(taskID)
+	if err != nil {
+		return results, errors.Wrap(err, "could not get scan status")
+	}
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	// Start progress logging for the created scan
+	if r.progress != nil {
+		ticker := time.NewTicker(time.Duration(r.options.StatsInterval) * time.Second)
+		r.progress.Init(r.hmapInputProvider.Count(), int(scanResponse.Templates), int64(scanResponse.Total))
+		go func() {
+			defer ticker.Stop()
+			for {
+				select {
+				case <-ctx.Done():
+					return
+				case <-ticker.C:
+					if scanResponse, err := r.cloudClient.GetScan(taskID); err == nil {
+						r.progress.SetRequests(uint64(scanResponse.Current))
+					}
+				}
+			}
+		}()
+	}
+
 	err = r.cloudClient.GetResults(taskID, true, limit, func(re *output.ResultEvent) {
+		r.progress.IncrementMatched()
 		results.CompareAndSwap(false, true)
 		_ = count.Inc()
 

v2/internal/runner/nucleicloud/cloud.go

@@ -118,7 +118,7 @@ func (c *Client) GetResults(ID int64, checkProgress bool, limit int, callback fu
 			callback(&result)
 		}
 
-		//This is checked during scan is added else if no item found break out of loop.
+		// This is checked during scan is added else if no item found break out of loop.
 		if checkProgress {
 			if items.Finished && len(items.Items) == 0 {
 				break
@@ -151,6 +151,25 @@ func (c *Client) GetScans(limit int, from string) ([]GetScanRequest, error) {
 	return items, nil
 }
 
+func (c *Client) GetScan(id int64) (GetScanRequest, error) {
+	var items GetScanRequest
+	httpReq, err := retryablehttp.NewRequest(http.MethodGet, fmt.Sprintf("%s/scan/%d", c.baseURL, id), nil)
+	if err != nil {
+		return items, errors.Wrap(err, "could not make request")
+	}
+
+	resp, err := c.sendRequest(httpReq)
+	if err != nil {
+		return items, errors.Wrap(err, "could not do request")
+	}
+	defer resp.Body.Close()
+
+	if err := jsoniter.NewDecoder(resp.Body).Decode(&items); err != nil {
+		return items, errors.Wrap(err, "could not decode results")
+	}
+	return items, nil
+}
+
 // Delete a scan and it's issues by the scan id.
 func (c *Client) DeleteScan(id int64) (DeleteScanResults, error) {
 	deletescan := DeleteScanResults{}

v2/internal/runner/runner.go

@@ -223,7 +223,12 @@ func New(options *types.Options) (*Runner, error) {
 	}
 	// Creates the progress tracking object
 	var progressErr error
-	runner.progress, progressErr = progress.NewStatsTicker(options.StatsInterval, options.EnableProgressBar, options.StatsJSON, options.Metrics, options.MetricsPort)
+	statsInterval := options.StatsInterval
+	if options.Cloud && !options.EnableProgressBar {
+		statsInterval = -1
+		options.EnableProgressBar = true
+	}
+	runner.progress, progressErr = progress.NewStatsTicker(statsInterval, options.EnableProgressBar, options.StatsJSON, options.Metrics, options.Cloud, options.MetricsPort)
 	if progressErr != nil {
 		return nil, progressErr
 	}

v2/pkg/core/workflow_execute_test.go

@@ -16,7 +16,7 @@ import (
 )
 
 func TestWorkflowsSimple(t *testing.T) {
-	progressBar, _ := progress.NewStatsTicker(0, false, false, false, 0)
+	progressBar, _ := progress.NewStatsTicker(0, false, false, false, false, 0)
 
 	workflow := &workflows.Workflow{Options: &protocols.ExecuterOptions{Options: &types.Options{TemplateThreads: 10}}, Workflows: []*workflows.WorkflowTemplate{
 		{Executers: []*workflows.ProtocolExecuterPair{{
@@ -30,7 +30,7 @@ func TestWorkflowsSimple(t *testing.T) {
 }
 
 func TestWorkflowsSimpleMultiple(t *testing.T) {
-	progressBar, _ := progress.NewStatsTicker(0, false, false, false, 0)
+	progressBar, _ := progress.NewStatsTicker(0, false, false, false, false, 0)
 
 	var firstInput, secondInput string
 	workflow := &workflows.Workflow{Options: &protocols.ExecuterOptions{Options: &types.Options{TemplateThreads: 10}}, Workflows: []*workflows.WorkflowTemplate{
@@ -55,7 +55,7 @@ func TestWorkflowsSimpleMultiple(t *testing.T) {
 }
 
 func TestWorkflowsSubtemplates(t *testing.T) {
-	progressBar, _ := progress.NewStatsTicker(0, false, false, false, 0)
+	progressBar, _ := progress.NewStatsTicker(0, false, false, false, false, 0)
 
 	var firstInput, secondInput string
 	workflow := &workflows.Workflow{Options: &protocols.ExecuterOptions{Options: &types.Options{TemplateThreads: 10}}, Workflows: []*workflows.WorkflowTemplate{
@@ -81,7 +81,7 @@ func TestWorkflowsSubtemplates(t *testing.T) {
 }
 
 func TestWorkflowsSubtemplatesNoMatch(t *testing.T) {
-	progressBar, _ := progress.NewStatsTicker(0, false, false, false, 0)
+	progressBar, _ := progress.NewStatsTicker(0, false, false, false, false, 0)
 
 	var firstInput, secondInput string
 	workflow := &workflows.Workflow{Options: &protocols.ExecuterOptions{Options: &types.Options{TemplateThreads: 10}}, Workflows: []*workflows.WorkflowTemplate{
@@ -105,7 +105,7 @@ func TestWorkflowsSubtemplatesNoMatch(t *testing.T) {
 }
 
 func TestWorkflowsSubtemplatesWithMatcher(t *testing.T) {
-	progressBar, _ := progress.NewStatsTicker(0, false, false, false, 0)
+	progressBar, _ := progress.NewStatsTicker(0, false, false, false, false, 0)
 
 	var firstInput, secondInput string
 	workflow := &workflows.Workflow{Options: &protocols.ExecuterOptions{Options: &types.Options{TemplateThreads: 10}}, Workflows: []*workflows.WorkflowTemplate{
@@ -134,7 +134,7 @@ func TestWorkflowsSubtemplatesWithMatcher(t *testing.T) {
 }
 
 func TestWorkflowsSubtemplatesWithMatcherNoMatch(t *testing.T) {
-	progressBar, _ := progress.NewStatsTicker(0, false, false, false, 0)
+	progressBar, _ := progress.NewStatsTicker(0, false, false, false, false, 0)
 
 	var firstInput, secondInput string
 	workflow := &workflows.Workflow{Options: &protocols.ExecuterOptions{Options: &types.Options{TemplateThreads: 10}}, Workflows: []*workflows.WorkflowTemplate{

v2/pkg/progress/progress.go

@@ -26,6 +26,8 @@ type Progress interface {
 	AddToTotal(delta int64)
 	// IncrementRequests increments the requests counter by 1.
 	IncrementRequests()
+	// SetRequests sets the counter by incrementing it with a delta
+	SetRequests(count uint64)
 	// IncrementMatched increments the matched counter by 1.
 	IncrementMatched()
 	// IncrementErrorsBy increments the error counter by count.
@@ -39,6 +41,7 @@ var _ Progress = &StatsTicker{}
 
 // StatsTicker is a progress instance for showing program stats
 type StatsTicker struct {
+	cloud        bool
 	active       bool
 	outputJSON   bool
 	server       *http.Server
@@ -47,9 +50,9 @@ type StatsTicker struct {
 }
 
 // NewStatsTicker creates and returns a new progress tracking object.
-func NewStatsTicker(duration int, active, outputJSON, metrics bool, port int) (Progress, error) {
+func NewStatsTicker(duration int, active, outputJSON, metrics, cloud bool, port int) (Progress, error) {
 	var tickDuration time.Duration
-	if active {
+	if active && duration != -1 {
 		tickDuration = time.Duration(duration) * time.Second
 	} else {
 		tickDuration = -1
@@ -61,6 +64,7 @@ func NewStatsTicker(duration int, active, outputJSON, metrics bool, port int) (P
 	if err != nil {
 		return nil, err
 	}
+	progress.cloud = cloud
 	progress.active = active
 	progress.stats = stats
 	progress.tickDuration = tickDuration
@@ -99,7 +103,7 @@ func (p *StatsTicker) Init(hostCount int64, rulesCount int, requestCount int64)
 		if p.outputJSON {
 			printCallbackFunc = printCallbackJSON
 		} else {
-			printCallbackFunc = printCallback
+			printCallbackFunc = p.makePrintCallback()
 		}
 		if err := p.stats.Start(printCallbackFunc, p.tickDuration); err != nil {
 			gologger.Warning().Msgf("Couldn't start statistics: %s", err)
@@ -117,6 +121,13 @@ func (p *StatsTicker) IncrementRequests() {
 	p.stats.IncrementCounter("requests", 1)
 }
 
+// SetRequests sets the counter by incrementing it with a delta
+func (p *StatsTicker) SetRequests(count uint64) {
+	value, _ := p.stats.GetCounter("requests")
+	delta := count - value
+	p.stats.IncrementCounter("requests", int(delta))
+}
+
 // IncrementMatched increments the matched counter by 1.
 func (p *StatsTicker) IncrementMatched() {
 	p.stats.IncrementCounter("matched", 1)
@@ -134,60 +145,66 @@ func (p *StatsTicker) IncrementFailedRequestsBy(count int64) {
 	p.stats.IncrementCounter("errors", int(count))
 }
 
-func printCallback(stats clistats.StatisticsClient) {
-	builder := &strings.Builder{}
+func (p *StatsTicker) makePrintCallback() func(stats clistats.StatisticsClient) {
+	return func(stats clistats.StatisticsClient) {
+		builder := &strings.Builder{}
 
-	var duration time.Duration
-	if startedAt, ok := stats.GetStatic("startedAt"); ok {
-		if startedAtTime, ok := startedAt.(time.Time); ok {
-			duration = time.Since(startedAtTime)
-			builder.WriteString(fmt.Sprintf("[%s]", fmtDuration(duration)))
+		var duration time.Duration
+		if startedAt, ok := stats.GetStatic("startedAt"); ok {
+			if startedAtTime, ok := startedAt.(time.Time); ok {
+				duration = time.Since(startedAtTime)
+				builder.WriteString(fmt.Sprintf("[%s]", fmtDuration(duration)))
+			}
 		}
+
+		if templates, ok := stats.GetStatic("templates"); ok {
+			builder.WriteString(" | Templates: ")
+			builder.WriteString(clistats.String(templates))
+		}
+
+		if hosts, ok := stats.GetStatic("hosts"); ok {
+			builder.WriteString(" | Hosts: ")
+			builder.WriteString(clistats.String(hosts))
+		}
+
+		requests, okRequests := stats.GetCounter("requests")
+		total, okTotal := stats.GetCounter("total")
+
+		if okRequests && okTotal && duration > 0 && !p.cloud {
+			builder.WriteString(" | RPS: ")
+			builder.WriteString(clistats.String(uint64(float64(requests) / duration.Seconds())))
+		}
+
+		if matched, ok := stats.GetCounter("matched"); ok {
+			builder.WriteString(" | Matched: ")
+			builder.WriteString(clistats.String(matched))
+		}
+
+		if errors, ok := stats.GetCounter("errors"); ok && !p.cloud {
+			builder.WriteString(" | Errors: ")
+			builder.WriteString(clistats.String(errors))
+		}
+
+		if okRequests && okTotal {
+			if p.cloud {
+				builder.WriteString(" | Task: ")
+			} else {
+				builder.WriteString(" | Requests: ")
+			}
+			builder.WriteString(clistats.String(requests))
+			builder.WriteRune('/')
+			builder.WriteString(clistats.String(total))
+			builder.WriteRune(' ')
+			builder.WriteRune('(')
+			//nolint:gomnd // this is not a magic number
+			builder.WriteString(clistats.String(uint64(float64(requests) / float64(total) * 100.0)))
+			builder.WriteRune('%')
+			builder.WriteRune(')')
+			builder.WriteRune('\n')
+		}
+
+		fmt.Fprintf(os.Stderr, "%s", builder.String())
 	}
-
-	if templates, ok := stats.GetStatic("templates"); ok {
-		builder.WriteString(" | Templates: ")
-		builder.WriteString(clistats.String(templates))
-	}
-
-	if hosts, ok := stats.GetStatic("hosts"); ok {
-		builder.WriteString(" | Hosts: ")
-		builder.WriteString(clistats.String(hosts))
-	}
-
-	requests, okRequests := stats.GetCounter("requests")
-	total, okTotal := stats.GetCounter("total")
-
-	if okRequests && okTotal && duration > 0 {
-		builder.WriteString(" | RPS: ")
-		builder.WriteString(clistats.String(uint64(float64(requests) / duration.Seconds())))
-	}
-
-	if matched, ok := stats.GetCounter("matched"); ok {
-		builder.WriteString(" | Matched: ")
-		builder.WriteString(clistats.String(matched))
-	}
-
-	if errors, ok := stats.GetCounter("errors"); ok {
-		builder.WriteString(" | Errors: ")
-		builder.WriteString(clistats.String(errors))
-	}
-
-	if okRequests && okTotal {
-		builder.WriteString(" | Requests: ")
-		builder.WriteString(clistats.String(requests))
-		builder.WriteRune('/')
-		builder.WriteString(clistats.String(total))
-		builder.WriteRune(' ')
-		builder.WriteRune('(')
-		//nolint:gomnd // this is not a magic number
-		builder.WriteString(clistats.String(uint64(float64(requests) / float64(total) * 100.0)))
-		builder.WriteRune('%')
-		builder.WriteRune(')')
-		builder.WriteRune('\n')
-	}
-
-	fmt.Fprintf(os.Stderr, "%s", builder.String())
 }
 
 func printCallbackJSON(stats clistats.StatisticsClient) {
@@ -256,7 +273,7 @@ func (p *StatsTicker) Stop() {
 		if p.outputJSON {
 			printCallbackJSON(p.stats)
 		} else {
-			printCallback(p.stats)
+			p.makePrintCallback()(p.stats)
 		}
 		if err := p.stats.Stop(); err != nil {
 			gologger.Warning().Msgf("Couldn't stop statistics: %s", err)

v2/pkg/templates/compile_test.go

@@ -32,7 +32,7 @@ var executerOpts protocols.ExecuterOptions
 func setup() {
 	options := testutils.DefaultOptions
 	testutils.Init(options)
-	progressImpl, _ := progress.NewStatsTicker(0, false, false, false, 0)
+	progressImpl, _ := progress.NewStatsTicker(0, false, false, false, false, 0)
 
 	executerOpts = protocols.ExecuterOptions{
 		Output:       testutils.NewMockOutputWriter(),

v2/pkg/testutils/testutils.go

@@ -79,7 +79,7 @@ type TemplateInfo struct {
 
 // NewMockExecuterOptions creates a new mock executeroptions struct
 func NewMockExecuterOptions(options *types.Options, info *TemplateInfo) *protocols.ExecuterOptions {
-	progressImpl, _ := progress.NewStatsTicker(0, false, false, false, 0)
+	progressImpl, _ := progress.NewStatsTicker(0, false, false, false, false, 0)
 	executerOpts := &protocols.ExecuterOptions{
 		TemplateID:   info.ID,
 		TemplateInfo: info.Info,
@@ -159,6 +159,9 @@ func (m *MockProgressClient) AddToTotal(delta int64) {}
 // IncrementRequests increments the requests counter by 1.
 func (m *MockProgressClient) IncrementRequests() {}
 
+// SetRequests sets the counter by incrementing it with a delta
+func (m *MockProgressClient) SetRequests(count uint64) {}
+
 // IncrementMatched increments the matched counter by 1.
 func (m *MockProgressClient) IncrementMatched() {}