diff --git a/v2/internal/runner/runner.go b/v2/internal/runner/runner.go index 100e364f..84f2837a 100644 --- a/v2/internal/runner/runner.go +++ b/v2/internal/runner/runner.go @@ -9,8 +9,6 @@ import ( "github.com/logrusorgru/aurora" "github.com/pkg/errors" - "github.com/remeh/sizedwaitgroup" - "go.uber.org/atomic" "go.uber.org/ratelimit" "gopkg.in/yaml.v2" @@ -20,7 +18,7 @@ import ( "github.com/projectdiscovery/nuclei/v2/pkg/catalog/config" "github.com/projectdiscovery/nuclei/v2/pkg/catalog/loader" "github.com/projectdiscovery/nuclei/v2/pkg/core" - "github.com/projectdiscovery/nuclei/v2/pkg/engine/inputs/hmap" + "github.com/projectdiscovery/nuclei/v2/pkg/engine/inputs/hybrid" "github.com/projectdiscovery/nuclei/v2/pkg/model/types/severity" "github.com/projectdiscovery/nuclei/v2/pkg/output" "github.com/projectdiscovery/nuclei/v2/pkg/parsers" @@ -34,7 +32,6 @@ import ( "github.com/projectdiscovery/nuclei/v2/pkg/reporting" "github.com/projectdiscovery/nuclei/v2/pkg/reporting/exporters/markdown" "github.com/projectdiscovery/nuclei/v2/pkg/reporting/exporters/sarif" - "github.com/projectdiscovery/nuclei/v2/pkg/templates" "github.com/projectdiscovery/nuclei/v2/pkg/types" "github.com/projectdiscovery/nuclei/v2/pkg/utils" "github.com/projectdiscovery/nuclei/v2/pkg/utils/stats" @@ -52,7 +49,7 @@ type Runner struct { colorizer aurora.Aurora issuesClient *reporting.Client addColor func(severity.Severity) string - hmapInputProvider *hmap.Input + hmapInputProvider *hybrid.Input browser *engine.Browser ratelimiter ratelimit.Limiter hostErrors *hosterrorscache.Cache @@ -112,7 +109,7 @@ func New(options *types.Options) (*Runner, error) { } // Initialize the input source - hmapInput, err := hmap.New(options) + hmapInput, err := hybrid.New(options) if err != nil { return nil, errors.Wrap(err, "could not create input provider") } @@ -259,6 +256,7 @@ func (r *Runner) RunEnumeration() error { ProjectFile: r.projectFile, Browser: r.browser, HostErrorsCache: cache, + Colorizer: r.colorizer, } engine := core.New(r.options) engine.SetExecuterOptions(executerOpts) @@ -351,9 +349,6 @@ func (r *Runner) RunEnumeration() error { gologger.Info().Msgf("Workflows loaded for scan: %d", len(store.Workflows())) } - // pre-parse all the templates, apply filters - finalTemplates := []*templates.Template{} - var unclusteredRequests int64 for _, template := range store.Templates() { // workflows will dynamically adjust the totals while running, as @@ -373,6 +368,12 @@ func (r *Runner) RunEnumeration() error { } } + // Cluster the templates first because we want info on how many + // templates did we cluster for showing to user in CLI + originalTemplatesCount := len(store.Templates()) + finalTemplates, clusterCount := engine.ClusterTemplates(store.Templates()) + finalTemplates = append(finalTemplates, store.Workflows()...) + var totalRequests int64 for _, t := range finalTemplates { if len(t.Workflows) > 0 { @@ -391,32 +392,10 @@ func (r *Runner) RunEnumeration() error { return errors.New("no valid templates were found") } - /* - TODO does it make sense to run the logic below if there are no targets specified? - Can we safely assume the user is just experimenting with the template/workflow filters before running them? - */ - - results := &atomic.Bool{} - wgtemplates := sizedwaitgroup.New(r.options.TemplateThreads) - // tracks global progress and captures stdout/stderr until p.Wait finishes r.progress.Init(r.hmapInputProvider.Count(), templateCount, totalRequests) - for _, t := range finalTemplates { - wgtemplates.Add() - go func(template *templates.Template) { - defer wgtemplates.Done() - - if template.SelfContained { - results.CAS(false, r.processSelfContainedTemplates(template)) - } else if len(template.Workflows) > 0 { - results.CAS(false, r.processWorkflowWithList(template)) - } else { - results.CAS(false, r.processTemplateWithList(template)) - } - }(t) - } - wgtemplates.Wait() + results := engine.ExecuteWithOpts(finalTemplates, r.hmapInputProvider, true) if r.interactsh != nil { matched := r.interactsh.Close() diff --git a/v2/pkg/core/engine.go b/v2/pkg/core/engine.go index d7fdc6bb..45554624 100644 --- a/v2/pkg/core/engine.go +++ b/v2/pkg/core/engine.go @@ -24,6 +24,10 @@ type Engine struct { // // An example InputProvider is provided in form of hmap input provider. type InputProvider interface { + // Count returns the number of items for input provider + Count() int64 + // Scan calls a callback function till the input provider is exhausted + Scan(callback func(value string)) } // New returns a new Engine instance diff --git a/v2/pkg/core/execute.go b/v2/pkg/core/execute.go index 112be4e2..3593123a 100644 --- a/v2/pkg/core/execute.go +++ b/v2/pkg/core/execute.go @@ -3,9 +3,12 @@ package core import ( "fmt" + "github.com/projectdiscovery/gologger" "github.com/projectdiscovery/nuclei/v2/pkg/protocols/common/clusterer" "github.com/projectdiscovery/nuclei/v2/pkg/templates" + "github.com/remeh/sizedwaitgroup" "github.com/rs/xid" + "go.uber.org/atomic" ) // Execute takes a list of templates/workflows that have been compiled @@ -13,27 +16,86 @@ import ( // // All the execution logic for the templates/workflows happens in this part // of the engine. -func (e *Engine) Execute(templates []*templates.Template) { - finalTemplates, clusterCount := e.clusterTemplates(templates) +func (e *Engine) Execute(templates []*templates.Template, input InputProvider) *atomic.Bool { + return e.ExecuteWithOpts(templates, input, false) +} +// ExecuteWithOpts is execute with the full options +func (e *Engine) ExecuteWithOpts(templatesList []*templates.Template, input InputProvider, noCluster bool) *atomic.Bool { + var finalTemplates []*templates.Template + if !noCluster { + finalTemplates, _ = e.ClusterTemplates(templatesList) + } else { + finalTemplates = templatesList + } + + results := &atomic.Bool{} for _, template := range finalTemplates { templateType := template.Type() + var wg *sizedwaitgroup.SizedWaitGroup + if templateType == "headless" { + wg = e.workPool.Headless + } else { + wg = e.workPool.Default + } + + wg.Add() switch { case template.SelfContained: - // Self Contained requests are executed here - - case templateType == "workflow": - // Workflows requests are executed here - + // Self Contained requests are executed here separately + e.executeSelfContainedTemplateWithInput(template, results) default: // All other request types are executed here + e.executeModelWithInput(templateType, template, input, results) } } + e.workPool.Wait() + return results } -// clusterTemplates performs identical http requests clustering for a list of templates -func (e *Engine) clusterTemplates(templatesList []*templates.Template) ([]*templates.Template, int) { +// processSelfContainedTemplates execute a self-contained template. +func (e *Engine) executeSelfContainedTemplateWithInput(template *templates.Template, results *atomic.Bool) { + match, err := template.Executer.Execute("") + if err != nil { + gologger.Warning().Msgf("[%s] Could not execute step: %s\n", e.executerOpts.Colorizer.BrightBlue(template.ID), err) + } + results.CAS(false, match) +} + +// executeModelWithInput executes a type of template with input +func (e *Engine) executeModelWithInput(templateType string, template *templates.Template, input InputProvider, results *atomic.Bool) { + wg := e.workPool.InputPool(templateType) + + input.Scan(func(scannedValue string) { + // Skip if the host has had errors + if e.executerOpts.HostErrorsCache != nil && e.executerOpts.HostErrorsCache.Check(scannedValue) { + return + } + + wg.Waitgroup.Add() + go func(value string) { + defer wg.Waitgroup.Done() + + var match bool + var err error + switch templateType { + case "workflow": + match = e.executeWorkflow(value, template.CompiledWorkflow) + default: + match, err = template.Executer.Execute(value) + } + if err != nil { + gologger.Warning().Msgf("[%s] Could not execute step: %s\n", e.executerOpts.Colorizer.BrightBlue(template.ID), err) + } + results.CAS(false, match) + }(scannedValue) + }) + wg.Waitgroup.Wait() +} + +// ClusterTemplates performs identical http requests clustering for a list of templates +func (e *Engine) ClusterTemplates(templatesList []*templates.Template) ([]*templates.Template, int) { if e.options.OfflineHTTP { return templatesList, 0 } @@ -65,85 +127,3 @@ func (e *Engine) clusterTemplates(templatesList []*templates.Template) ([]*templ } return finalTemplatesList, clusterCount } - -/* -import ( - "github.com/projectdiscovery/gologger" - "github.com/projectdiscovery/nuclei/v2/pkg/templates" - "github.com/remeh/sizedwaitgroup" - "go.uber.org/atomic" -) - -// processSelfContainedTemplates execute a self-contained template. -func (r *Runner) processSelfContainedTemplates(template *templates.Template) bool { - match, err := template.Executer.Execute("") - if err != nil { - gologger.Warning().Msgf("[%s] Could not execute step: %s\n", r.colorizer.BrightBlue(template.ID), err) - } - return match -} - -// processTemplateWithList execute a template against the list of user provided targets -func (r *Runner) processTemplateWithList(template *templates.Template) bool { - results := &atomic.Bool{} - wg := sizedwaitgroup.New(r.options.BulkSize) - processItem := func(k, _ []byte) error { - URL := string(k) - - // Skip if the host has had errors - if r.hostErrors != nil && r.hostErrors.Check(URL) { - return nil - } - wg.Add() - go func(URL string) { - defer wg.Done() - - match, err := template.Executer.Execute(URL) - if err != nil { - gologger.Warning().Msgf("[%s] Could not execute step: %s\n", r.colorizer.BrightBlue(template.ID), err) - } - results.CAS(false, match) - }(URL) - return nil - } - if r.options.Stream { - _ = r.hostMapStream.Scan(processItem) - } else { - r.hostMap.Scan(processItem) - } - - wg.Wait() - return results.Load() -} - -// processTemplateWithList process a template on the URL list -func (r *Runner) processWorkflowWithList(template *templates.Template) bool { - results := &atomic.Bool{} - wg := sizedwaitgroup.New(r.options.BulkSize) - - processItem := func(k, _ []byte) error { - URL := string(k) - - // Skip if the host has had errors - if r.hostErrors != nil && r.hostErrors.Check(URL) { - return nil - } - wg.Add() - go func(URL string) { - defer wg.Done() - match := template.CompiledWorkflow.RunWorkflow(URL) - results.CAS(false, match) - }(URL) - return nil - } - - if r.options.Stream { - _ = r.hostMapStream.Scan(processItem) - } else { - r.hostMap.Scan(processItem) - } - - wg.Wait() - return results.Load() -} -*/ diff --git a/v2/pkg/core/inputs/hmap/hmap.go b/v2/pkg/core/inputs/hybrid/hmap.go similarity index 88% rename from v2/pkg/core/inputs/hmap/hmap.go rename to v2/pkg/core/inputs/hybrid/hmap.go index 2910d045..1eb4f5e4 100644 --- a/v2/pkg/core/inputs/hmap/hmap.go +++ b/v2/pkg/core/inputs/hybrid/hmap.go @@ -1,6 +1,6 @@ -// Package hmap implements a hybrid hmap/filekv backed input provider +// Package hybrid implements a hybrid hmap/filekv backed input provider // for nuclei that can either stream or store results using different kv stores. -package hmap +package hybrid import ( "bufio" @@ -118,3 +118,16 @@ func (i *Input) normalizeStoreInputValue(value string) { func (i *Input) Count() int64 { return i.inputCount } + +// Scan calls an input provider till the callback is exhausted +func (i *Input) Scan(callback func(value string)) { + callbackFunc := func(k, _ []byte) error { + callback(string(k)) + return nil + } + if i.hostMapStream != nil { + _ = i.hostMapStream.Scan(callbackFunc) + } else { + i.hostMap.Scan(callbackFunc) + } +} diff --git a/v2/pkg/core/workflow_execute.go b/v2/pkg/core/workflow_execute.go index 59a5ade5..8c255d7a 100644 --- a/v2/pkg/core/workflow_execute.go +++ b/v2/pkg/core/workflow_execute.go @@ -1,15 +1,22 @@ package core -/* -// RunWorkflow runs a workflow on an input and returns true or false -func (w *Workflow) RunWorkflow(input string) bool { +import ( + "github.com/projectdiscovery/gologger" + "github.com/projectdiscovery/nuclei/v2/pkg/output" + "github.com/projectdiscovery/nuclei/v2/pkg/workflows" + "github.com/remeh/sizedwaitgroup" + "go.uber.org/atomic" +) + +// executeWorkflow runs a workflow on an input and returns true or false +func (e *Engine) executeWorkflow(input string, w *workflows.Workflow) bool { results := &atomic.Bool{} swg := sizedwaitgroup.New(w.Options.Options.TemplateThreads) for _, template := range w.Workflows { swg.Add() - func(template *WorkflowTemplate) { - if err := w.runWorkflowStep(template, input, results, &swg); err != nil { + func(template *workflows.WorkflowTemplate) { + if err := e.runWorkflowStep(template, input, results, &swg, w); err != nil { gologger.Warning().Msgf("[%s] Could not execute workflow step: %s\n", template.Template, err) } swg.Done() @@ -21,7 +28,7 @@ func (w *Workflow) RunWorkflow(input string) bool { // runWorkflowStep runs a workflow step for the workflow. It executes the workflow // in a recursive manner running all subtemplates and matchers. -func (w *Workflow) runWorkflowStep(template *WorkflowTemplate, input string, results *atomic.Bool, swg *sizedwaitgroup.SizedWaitGroup) error { +func (e *Engine) runWorkflowStep(template *workflows.WorkflowTemplate, input string, results *atomic.Bool, swg *sizedwaitgroup.SizedWaitGroup, w *workflows.Workflow) error { var firstMatched bool var err error var mainErr error @@ -84,8 +91,8 @@ func (w *Workflow) runWorkflowStep(template *WorkflowTemplate, input string, res for _, subtemplate := range matcher.Subtemplates { swg.Add() - go func(subtemplate *WorkflowTemplate) { - if err := w.runWorkflowStep(subtemplate, input, results, swg); err != nil { + go func(subtemplate *workflows.WorkflowTemplate) { + if err := e.runWorkflowStep(subtemplate, input, results, swg, w); err != nil { gologger.Warning().Msgf("[%s] Could not execute workflow step: %s\n", subtemplate.Template, err) } swg.Done() @@ -108,8 +115,8 @@ func (w *Workflow) runWorkflowStep(template *WorkflowTemplate, input string, res for _, subtemplate := range template.Subtemplates { swg.Add() - go func(template *WorkflowTemplate) { - if err := w.runWorkflowStep(template, input, results, swg); err != nil { + go func(template *workflows.WorkflowTemplate) { + if err := e.runWorkflowStep(template, input, results, swg, w); err != nil { gologger.Warning().Msgf("[%s] Could not execute workflow step: %s\n", template.Template, err) } swg.Done() @@ -117,4 +124,4 @@ func (w *Workflow) runWorkflowStep(template *WorkflowTemplate, input string, res } } return mainErr -}*/ +} diff --git a/v2/pkg/core/workflow_execute_test.go b/v2/pkg/core/workflow_execute_test.go index 62cd4084..a00ce604 100644 --- a/v2/pkg/core/workflow_execute_test.go +++ b/v2/pkg/core/workflow_execute_test.go @@ -1,6 +1,5 @@ package core -/* import ( "testing" @@ -17,13 +16,14 @@ import ( func TestWorkflowsSimple(t *testing.T) { progressBar, _ := progress.NewStatsTicker(0, false, false, false, 0) - workflow := &workflows.Workflow{Options: &protocols.ExecuterOptions{Options: &types.Options{TemplateThreads: 10}}, Workflows: []*WorkflowTemplate{ + workflow := &workflows.Workflow{Options: &protocols.ExecuterOptions{Options: &types.Options{TemplateThreads: 10}}, Workflows: []*workflows.WorkflowTemplate{ {Executers: []*workflows.ProtocolExecuterPair{{ Executer: &mockExecuter{result: true}, Options: &protocols.ExecuterOptions{Progress: progressBar}}, }}, }} - matched := workflow.RunWorkflow("https://test.com") + engine := &Engine{} + matched := engine.executeWorkflow("https://test.com", workflow) require.True(t, matched, "could not get correct match value") } @@ -31,7 +31,7 @@ func TestWorkflowsSimpleMultiple(t *testing.T) { progressBar, _ := progress.NewStatsTicker(0, false, false, false, 0) var firstInput, secondInput string - workflow := &workflows.Workflow{Options: &protocols.ExecuterOptions{Options: &types.Options{TemplateThreads: 10}}, Workflows: []*WorkflowTemplate{ + workflow := &workflows.Workflow{Options: &protocols.ExecuterOptions{Options: &types.Options{TemplateThreads: 10}}, Workflows: []*workflows.WorkflowTemplate{ {Executers: []*workflows.ProtocolExecuterPair{{ Executer: &mockExecuter{result: true, executeHook: func(input string) { firstInput = input @@ -44,7 +44,8 @@ func TestWorkflowsSimpleMultiple(t *testing.T) { }}, }} - matched := workflow.RunWorkflow("https://test.com") + engine := &Engine{} + matched := engine.executeWorkflow("https://test.com", workflow) require.True(t, matched, "could not get correct match value") require.Equal(t, "https://test.com", firstInput, "could not get correct first input") @@ -55,7 +56,7 @@ func TestWorkflowsSubtemplates(t *testing.T) { progressBar, _ := progress.NewStatsTicker(0, false, false, false, 0) var firstInput, secondInput string - workflow := &workflows.Workflow{Options: &protocols.ExecuterOptions{Options: &types.Options{TemplateThreads: 10}}, Workflows: []*WorkflowTemplate{ + workflow := &workflows.Workflow{Options: &protocols.ExecuterOptions{Options: &types.Options{TemplateThreads: 10}}, Workflows: []*workflows.WorkflowTemplate{ {Executers: []*workflows.ProtocolExecuterPair{{ Executer: &mockExecuter{result: true, executeHook: func(input string) { firstInput = input @@ -69,7 +70,8 @@ func TestWorkflowsSubtemplates(t *testing.T) { }}}}, }} - matched := workflow.RunWorkflow("https://test.com") + engine := &Engine{} + matched := engine.executeWorkflow("https://test.com", workflow) require.True(t, matched, "could not get correct match value") require.Equal(t, "https://test.com", firstInput, "could not get correct first input") @@ -80,7 +82,7 @@ func TestWorkflowsSubtemplatesNoMatch(t *testing.T) { progressBar, _ := progress.NewStatsTicker(0, false, false, false, 0) var firstInput, secondInput string - workflow := &workflows.Workflow{Options: &protocols.ExecuterOptions{Options: &types.Options{TemplateThreads: 10}}, Workflows: []*WorkflowTemplate{ + workflow := &workflows.Workflow{Options: &protocols.ExecuterOptions{Options: &types.Options{TemplateThreads: 10}}, Workflows: []*workflows.WorkflowTemplate{ {Executers: []*workflows.ProtocolExecuterPair{{ Executer: &mockExecuter{result: false, executeHook: func(input string) { firstInput = input @@ -92,7 +94,8 @@ func TestWorkflowsSubtemplatesNoMatch(t *testing.T) { }}}}, }} - matched := workflow.RunWorkflow("https://test.com") + engine := &Engine{} + matched := engine.executeWorkflow("https://test.com", workflow) require.False(t, matched, "could not get correct match value") require.Equal(t, "https://test.com", firstInput, "could not get correct first input") @@ -103,7 +106,7 @@ func TestWorkflowsSubtemplatesWithMatcher(t *testing.T) { progressBar, _ := progress.NewStatsTicker(0, false, false, false, 0) var firstInput, secondInput string - workflow := &workflows.Workflow{Options: &protocols.ExecuterOptions{Options: &types.Options{TemplateThreads: 10}}, Workflows: []*WorkflowTemplate{ + workflow := &workflows.Workflow{Options: &protocols.ExecuterOptions{Options: &types.Options{TemplateThreads: 10}}, Workflows: []*workflows.WorkflowTemplate{ {Executers: []*workflows.ProtocolExecuterPair{{ Executer: &mockExecuter{result: true, executeHook: func(input string) { firstInput = input @@ -120,7 +123,8 @@ func TestWorkflowsSubtemplatesWithMatcher(t *testing.T) { }}}}}}, }} - matched := workflow.RunWorkflow("https://test.com") + engine := &Engine{} + matched := engine.executeWorkflow("https://test.com", workflow) require.True(t, matched, "could not get correct match value") require.Equal(t, "https://test.com", firstInput, "could not get correct first input") @@ -148,7 +152,8 @@ func TestWorkflowsSubtemplatesWithMatcherNoMatch(t *testing.T) { }}}}}}, }} - matched := workflow.RunWorkflow("https://test.com") + engine := &Engine{} + matched := engine.executeWorkflow("https://test.com", workflow) require.False(t, matched, "could not get correct match value") require.Equal(t, "https://test.com", firstInput, "could not get correct first input") @@ -189,4 +194,3 @@ func (m *mockExecuter) ExecuteWithResults(input string, callback protocols.Outpu } return nil } -*/ diff --git a/v2/pkg/core/workpool.go b/v2/pkg/core/workpool.go index ef50f70f..23b50d68 100644 --- a/v2/pkg/core/workpool.go +++ b/v2/pkg/core/workpool.go @@ -1,7 +1,7 @@ package core import ( - "github.com/projectdiscovery/nuclei/v2/pkg/templates" + "github.com/remeh/sizedwaitgroup" ) // WorkPool implements an execution pool for executing different @@ -10,7 +10,9 @@ import ( // It also allows Configuration of such requirements. This is used // for per-module like separate headless concurrency etc. type WorkPool struct { - config WorkPoolConfig + Headless *sizedwaitgroup.SizedWaitGroup + Default *sizedwaitgroup.SizedWaitGroup + config WorkPoolConfig } // WorkPoolConfig is the configuration for workpool @@ -27,9 +29,35 @@ type WorkPoolConfig struct { // NewWorkPool returns a new WorkPool instance func NewWorkPool(config WorkPoolConfig) *WorkPool { - return &WorkPool{config: config} + headlessWg := sizedwaitgroup.New(config.HeadlessTypeConcurrency) + defaultWg := sizedwaitgroup.New(config.TypeConcurrency) + + return &WorkPool{ + config: config, + Headless: &headlessWg, + Default: &defaultWg, + } } -func (w *WorkPool) Execute(templates []*templates.Template) { - +// Wait waits for all the workpool waitgroups to finish +func (w *WorkPool) Wait() { + w.Default.Wait() + w.Headless.Wait() +} + +// InputWorkPool is a workpool per-input +type InputWorkPool struct { + Waitgroup *sizedwaitgroup.SizedWaitGroup +} + +// InputPool returns a workpool for an input type +func (w *WorkPool) InputPool(templateType string) *InputWorkPool { + var count int + if templateType == "headless" { + count = w.config.HeadlessInputConcurrency + } else { + count = w.config.InputConcurrency + } + swg := sizedwaitgroup.New(count) + return &InputWorkPool{Waitgroup: &swg} } diff --git a/v2/pkg/protocols/protocols.go b/v2/pkg/protocols/protocols.go index 83a76f7c..d877328f 100644 --- a/v2/pkg/protocols/protocols.go +++ b/v2/pkg/protocols/protocols.go @@ -3,6 +3,7 @@ package protocols import ( "go.uber.org/ratelimit" + "github.com/logrusorgru/aurora" "github.com/projectdiscovery/nuclei/v2/pkg/catalog" "github.com/projectdiscovery/nuclei/v2/pkg/model" "github.com/projectdiscovery/nuclei/v2/pkg/operators" @@ -61,6 +62,7 @@ type ExecuterOptions struct { Operators []*operators.Operators // only used by offlinehttp module + Colorizer aurora.Aurora WorkflowLoader model.WorkflowLoader }