From 67c094444ee0c63e89c263895bf2bf18d06916e6 Mon Sep 17 00:00:00 2001 From: Ice3man Date: Fri, 13 Jan 2023 13:41:05 +0530 Subject: [PATCH] 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 | 6 +- v2/go.sum | 4 +- v2/internal/runner/enumerate.go | 29 ++++++ v2/internal/runner/nucleicloud/cloud.go | 21 +++- v2/internal/runner/runner.go | 7 +- v2/pkg/core/workflow_execute_test.go | 12 +-- v2/pkg/progress/progress.go | 127 ++++++++++++++---------- v2/pkg/templates/compile_test.go | 2 +- v2/pkg/testutils/testutils.go | 5 +- 9 files changed, 143 insertions(+), 70 deletions(-) diff --git a/v2/go.mod b/v2/go.mod index eec02cf68..81628161b 100644 --- a/v2/go.mod +++ b/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 diff --git a/v2/go.sum b/v2/go.sum index ec511a528..4f4d2b8bc 100644 --- a/v2/go.sum +++ b/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= diff --git a/v2/internal/runner/enumerate.go b/v2/internal/runner/enumerate.go index d43885464..271141ca3 100644 --- a/v2/internal/runner/enumerate.go +++ b/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() diff --git a/v2/internal/runner/nucleicloud/cloud.go b/v2/internal/runner/nucleicloud/cloud.go index 266e31d70..60cc88c3c 100644 --- a/v2/internal/runner/nucleicloud/cloud.go +++ b/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{} diff --git a/v2/internal/runner/runner.go b/v2/internal/runner/runner.go index 9e09f1c01..b4c0b9c39 100644 --- a/v2/internal/runner/runner.go +++ b/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 } diff --git a/v2/pkg/core/workflow_execute_test.go b/v2/pkg/core/workflow_execute_test.go index b8858ff25..a64ec1f0b 100644 --- a/v2/pkg/core/workflow_execute_test.go +++ b/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{ diff --git a/v2/pkg/progress/progress.go b/v2/pkg/progress/progress.go index 5c21fc9d5..8f33ce6e0 100644 --- a/v2/pkg/progress/progress.go +++ b/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) diff --git a/v2/pkg/templates/compile_test.go b/v2/pkg/templates/compile_test.go index cf4cac708..7a060554f 100644 --- a/v2/pkg/templates/compile_test.go +++ b/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(), diff --git a/v2/pkg/testutils/testutils.go b/v2/pkg/testutils/testutils.go index 17596364a..23fbd0909 100644 --- a/v2/pkg/testutils/testutils.go +++ b/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() {}