From c16c93fe7ca698d3086b85fd86e0c1e8c99a222e Mon Sep 17 00:00:00 2001 From: Ice3man543 Date: Wed, 27 Oct 2021 16:50:36 +0530 Subject: [PATCH] refactor the modules to core --- v2/internal/runner/runner.go | 46 +++------------- v2/pkg/core/engine.go | 53 +++++++++++++++++++ v2/pkg/core/engine_test.go | 1 + .../processor.go => pkg/core/execute.go} | 46 +++++++++++++++- v2/pkg/{engine => core}/inputs/hmap/hmap.go | 0 .../execute.go => core/workflow_execute.go} | 12 ++--- .../workflow_execute_test.go} | 39 +++++++------- v2/pkg/{engine => core}/workpool.go | 6 +-- v2/pkg/engine/engine.go | 16 ------ v2/pkg/engine/engine_test.go | 1 - v2/pkg/protocols/protocols.go | 6 +++ v2/pkg/templates/workflows.go | 14 +---- 12 files changed, 139 insertions(+), 101 deletions(-) create mode 100644 v2/pkg/core/engine.go create mode 100644 v2/pkg/core/engine_test.go rename v2/{internal/runner/processor.go => pkg/core/execute.go} (61%) rename v2/pkg/{engine => core}/inputs/hmap/hmap.go (100%) rename v2/pkg/{workflows/execute.go => core/workflow_execute.go} (95%) rename v2/pkg/{workflows/execute_test.go => core/workflow_execute_test.go} (77%) rename v2/pkg/{engine => core}/workpool.go (90%) delete mode 100644 v2/pkg/engine/engine.go delete mode 100644 v2/pkg/engine/engine_test.go diff --git a/v2/internal/runner/runner.go b/v2/internal/runner/runner.go index d4d7f40c8..100e364f0 100644 --- a/v2/internal/runner/runner.go +++ b/v2/internal/runner/runner.go @@ -2,7 +2,6 @@ package runner import ( "bufio" - "fmt" "os" "path/filepath" "strings" @@ -11,7 +10,6 @@ import ( "github.com/logrusorgru/aurora" "github.com/pkg/errors" "github.com/remeh/sizedwaitgroup" - "github.com/rs/xid" "go.uber.org/atomic" "go.uber.org/ratelimit" "gopkg.in/yaml.v2" @@ -21,6 +19,7 @@ import ( "github.com/projectdiscovery/nuclei/v2/pkg/catalog" "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/model/types/severity" "github.com/projectdiscovery/nuclei/v2/pkg/output" @@ -28,7 +27,6 @@ import ( "github.com/projectdiscovery/nuclei/v2/pkg/progress" "github.com/projectdiscovery/nuclei/v2/pkg/projectfile" "github.com/projectdiscovery/nuclei/v2/pkg/protocols" - "github.com/projectdiscovery/nuclei/v2/pkg/protocols/common/clusterer" "github.com/projectdiscovery/nuclei/v2/pkg/protocols/common/hosterrorscache" "github.com/projectdiscovery/nuclei/v2/pkg/protocols/common/interactsh" "github.com/projectdiscovery/nuclei/v2/pkg/protocols/common/protocolinit" @@ -247,6 +245,9 @@ func (r *Runner) RunEnumeration() error { cache = hosterrorscache.New(r.options.MaxHostError, hosterrorscache.DefaultMaxHostsCount).SetVerbose(r.options.Verbose) } r.hostErrors = cache + + // Create the executer options which will be used throughout the execution + // stage by the nuclei engine modules. executerOpts := protocols.ExecuterOptions{ Output: r.output, Options: r.options, @@ -259,12 +260,13 @@ func (r *Runner) RunEnumeration() error { Browser: r.browser, HostErrorsCache: cache, } + engine := core.New(r.options) + engine.SetExecuterOptions(executerOpts) workflowLoader, err := parsers.NewLoader(&executerOpts) if err != nil { return errors.Wrap(err, "Could not create loader.") } - executerOpts.WorkflowLoader = workflowLoader loaderConfig := loader.Config{ @@ -370,42 +372,6 @@ func (r *Runner) RunEnumeration() error { r.logAvailableTemplate(template.Path) } } - templatesMap := make(map[string]*templates.Template) - for _, v := range store.Templates() { - templatesMap[v.Path] = v - } - originalTemplatesCount := len(store.Templates()) - clusterCount := 0 - clusters := clusterer.Cluster(templatesMap) - for _, cluster := range clusters { - if len(cluster) > 1 && !r.options.OfflineHTTP { - executerOpts := protocols.ExecuterOptions{ - Output: r.output, - Options: r.options, - Progress: r.progress, - Catalog: r.catalog, - RateLimiter: r.ratelimiter, - IssuesClient: r.issuesClient, - Browser: r.browser, - ProjectFile: r.projectFile, - Interactsh: r.interactsh, - HostErrorsCache: cache, - } - clusterID := fmt.Sprintf("cluster-%s", xid.New().String()) - - finalTemplates = append(finalTemplates, &templates.Template{ - ID: clusterID, - RequestsHTTP: cluster[0].RequestsHTTP, - Executer: clusterer.NewExecuter(cluster, &executerOpts), - TotalRequests: len(cluster[0].RequestsHTTP), - }) - clusterCount += len(cluster) - } else { - finalTemplates = append(finalTemplates, cluster...) - } - } - - finalTemplates = append(finalTemplates, store.Workflows()...) var totalRequests int64 for _, t := range finalTemplates { diff --git a/v2/pkg/core/engine.go b/v2/pkg/core/engine.go new file mode 100644 index 000000000..d7fdc6bbc --- /dev/null +++ b/v2/pkg/core/engine.go @@ -0,0 +1,53 @@ +package core + +import ( + "github.com/projectdiscovery/nuclei/v2/pkg/protocols" + "github.com/projectdiscovery/nuclei/v2/pkg/types" +) + +// Engine is an engine for running Nuclei Templates/Workflows. +// +// The engine contains multiple thread pools which allow using different +// concurrency values per protocol executed. +// +// The engine does most of the heavy lifting of execution, from clustering +// templates to leading to the final execution by the workpool, it is +// handled by the engine. +type Engine struct { + workPool *WorkPool + options *types.Options + executerOpts protocols.ExecuterOptions +} + +// InputProvider is an input provider interface for the nuclei execution +// engine. +// +// An example InputProvider is provided in form of hmap input provider. +type InputProvider interface { +} + +// New returns a new Engine instance +func New(options *types.Options) *Engine { + workPool := NewWorkPool(WorkPoolConfig{ + InputConcurrency: options.BulkSize, + TypeConcurrency: options.TemplateThreads, + HeadlessInputConcurrency: options.HeadlessBulkSize, + HeadlessTypeConcurrency: options.HeadlessTemplateThreads, + }) + engine := &Engine{ + options: options, + workPool: workPool, + } + return engine +} + +// SetExecuterOptions sets the executer options for the engine. This is required +// before using the engine to perform any execution. +func (e *Engine) SetExecuterOptions(options protocols.ExecuterOptions) { + e.executerOpts = options +} + +// ExecuterOptions returns protocols.ExecuterOptions for nuclei engine. +func (e *Engine) ExecuterOptions() protocols.ExecuterOptions { + return e.executerOpts +} diff --git a/v2/pkg/core/engine_test.go b/v2/pkg/core/engine_test.go new file mode 100644 index 000000000..9a8bc9592 --- /dev/null +++ b/v2/pkg/core/engine_test.go @@ -0,0 +1 @@ +package core diff --git a/v2/internal/runner/processor.go b/v2/pkg/core/execute.go similarity index 61% rename from v2/internal/runner/processor.go rename to v2/pkg/core/execute.go index 6b14b5f91..65c31558d 100644 --- a/v2/internal/runner/processor.go +++ b/v2/pkg/core/execute.go @@ -1,5 +1,48 @@ -package runner +package core +import ( + "fmt" + + "github.com/projectdiscovery/nuclei/v2/pkg/protocols/common/clusterer" + "github.com/projectdiscovery/nuclei/v2/pkg/templates" + "github.com/rs/xid" +) + +// 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 + } + + templatesMap := make(map[string]*templates.Template) + for _, v := range templatesList { + templatesMap[v.Path] = v + } + clusterCount := 0 + + finalTemplatesList := make([]*templates.Template, 0, len(templatesList)) + clusters := clusterer.Cluster(templatesMap) + for _, cluster := range clusters { + if len(cluster) > 1 { + executerOpts := e.ExecuterOptions() + + clusterID := fmt.Sprintf("cluster-%s", xid.New().String()) + + finalTemplatesList = append(finalTemplatesList, &templates.Template{ + ID: clusterID, + RequestsHTTP: cluster[0].RequestsHTTP, + Executer: clusterer.NewExecuter(cluster, &executerOpts), + TotalRequests: len(cluster[0].RequestsHTTP), + }) + clusterCount += len(cluster) + } else { + finalTemplatesList = append(finalTemplatesList, cluster...) + } + } + return finalTemplatesList, clusterCount +} + +/* import ( "github.com/projectdiscovery/gologger" "github.com/projectdiscovery/nuclei/v2/pkg/templates" @@ -79,3 +122,4 @@ func (r *Runner) processWorkflowWithList(template *templates.Template) bool { wg.Wait() return results.Load() } +*/ diff --git a/v2/pkg/engine/inputs/hmap/hmap.go b/v2/pkg/core/inputs/hmap/hmap.go similarity index 100% rename from v2/pkg/engine/inputs/hmap/hmap.go rename to v2/pkg/core/inputs/hmap/hmap.go diff --git a/v2/pkg/workflows/execute.go b/v2/pkg/core/workflow_execute.go similarity index 95% rename from v2/pkg/workflows/execute.go rename to v2/pkg/core/workflow_execute.go index c0710c6ad..59a5ade55 100644 --- a/v2/pkg/workflows/execute.go +++ b/v2/pkg/core/workflow_execute.go @@ -1,12 +1,6 @@ -package workflows - -import ( - "github.com/projectdiscovery/gologger" - "github.com/projectdiscovery/nuclei/v2/pkg/output" - "github.com/remeh/sizedwaitgroup" - "go.uber.org/atomic" -) +package core +/* // RunWorkflow runs a workflow on an input and returns true or false func (w *Workflow) RunWorkflow(input string) bool { results := &atomic.Bool{} @@ -123,4 +117,4 @@ func (w *Workflow) runWorkflowStep(template *WorkflowTemplate, input string, res } } return mainErr -} +}*/ diff --git a/v2/pkg/workflows/execute_test.go b/v2/pkg/core/workflow_execute_test.go similarity index 77% rename from v2/pkg/workflows/execute_test.go rename to v2/pkg/core/workflow_execute_test.go index 6d9ab6a09..62cd40840 100644 --- a/v2/pkg/workflows/execute_test.go +++ b/v2/pkg/core/workflow_execute_test.go @@ -1,5 +1,6 @@ -package workflows +package core +/* import ( "testing" @@ -10,13 +11,14 @@ import ( "github.com/projectdiscovery/nuclei/v2/pkg/progress" "github.com/projectdiscovery/nuclei/v2/pkg/protocols" "github.com/projectdiscovery/nuclei/v2/pkg/types" + "github.com/projectdiscovery/nuclei/v2/pkg/workflows" ) func TestWorkflowsSimple(t *testing.T) { progressBar, _ := progress.NewStatsTicker(0, false, false, false, 0) - workflow := &Workflow{Options: &protocols.ExecuterOptions{Options: &types.Options{TemplateThreads: 10}}, Workflows: []*WorkflowTemplate{ - {Executers: []*ProtocolExecuterPair{{ + workflow := &workflows.Workflow{Options: &protocols.ExecuterOptions{Options: &types.Options{TemplateThreads: 10}}, Workflows: []*WorkflowTemplate{ + {Executers: []*workflows.ProtocolExecuterPair{{ Executer: &mockExecuter{result: true}, Options: &protocols.ExecuterOptions{Progress: progressBar}}, }}, }} @@ -29,13 +31,13 @@ func TestWorkflowsSimpleMultiple(t *testing.T) { progressBar, _ := progress.NewStatsTicker(0, false, false, false, 0) var firstInput, secondInput string - workflow := &Workflow{Options: &protocols.ExecuterOptions{Options: &types.Options{TemplateThreads: 10}}, Workflows: []*WorkflowTemplate{ - {Executers: []*ProtocolExecuterPair{{ + workflow := &workflows.Workflow{Options: &protocols.ExecuterOptions{Options: &types.Options{TemplateThreads: 10}}, Workflows: []*WorkflowTemplate{ + {Executers: []*workflows.ProtocolExecuterPair{{ Executer: &mockExecuter{result: true, executeHook: func(input string) { firstInput = input }}, Options: &protocols.ExecuterOptions{Progress: progressBar}}, }}, - {Executers: []*ProtocolExecuterPair{{ + {Executers: []*workflows.ProtocolExecuterPair{{ Executer: &mockExecuter{result: true, executeHook: func(input string) { secondInput = input }}, Options: &protocols.ExecuterOptions{Progress: progressBar}}, @@ -53,14 +55,14 @@ func TestWorkflowsSubtemplates(t *testing.T) { progressBar, _ := progress.NewStatsTicker(0, false, false, false, 0) var firstInput, secondInput string - workflow := &Workflow{Options: &protocols.ExecuterOptions{Options: &types.Options{TemplateThreads: 10}}, Workflows: []*WorkflowTemplate{ - {Executers: []*ProtocolExecuterPair{{ + workflow := &workflows.Workflow{Options: &protocols.ExecuterOptions{Options: &types.Options{TemplateThreads: 10}}, Workflows: []*WorkflowTemplate{ + {Executers: []*workflows.ProtocolExecuterPair{{ Executer: &mockExecuter{result: true, executeHook: func(input string) { firstInput = input }, outputs: []*output.InternalWrappedEvent{ {OperatorsResult: &operators.Result{}, Results: []*output.ResultEvent{{}}}, }}, Options: &protocols.ExecuterOptions{Progress: progressBar}}, - }, Subtemplates: []*WorkflowTemplate{{Executers: []*ProtocolExecuterPair{{ + }, Subtemplates: []*workflows.WorkflowTemplate{{Executers: []*workflows.ProtocolExecuterPair{{ Executer: &mockExecuter{result: true, executeHook: func(input string) { secondInput = input }}, Options: &protocols.ExecuterOptions{Progress: progressBar}}, @@ -78,12 +80,12 @@ func TestWorkflowsSubtemplatesNoMatch(t *testing.T) { progressBar, _ := progress.NewStatsTicker(0, false, false, false, 0) var firstInput, secondInput string - workflow := &Workflow{Options: &protocols.ExecuterOptions{Options: &types.Options{TemplateThreads: 10}}, Workflows: []*WorkflowTemplate{ - {Executers: []*ProtocolExecuterPair{{ + workflow := &workflows.Workflow{Options: &protocols.ExecuterOptions{Options: &types.Options{TemplateThreads: 10}}, Workflows: []*WorkflowTemplate{ + {Executers: []*workflows.ProtocolExecuterPair{{ Executer: &mockExecuter{result: false, executeHook: func(input string) { firstInput = input }}, Options: &protocols.ExecuterOptions{Progress: progressBar}}, - }, Subtemplates: []*WorkflowTemplate{{Executers: []*ProtocolExecuterPair{{ + }, Subtemplates: []*workflows.WorkflowTemplate{{Executers: []*workflows.ProtocolExecuterPair{{ Executer: &mockExecuter{result: true, executeHook: func(input string) { secondInput = input }}, Options: &protocols.ExecuterOptions{Progress: progressBar}}, @@ -101,8 +103,8 @@ func TestWorkflowsSubtemplatesWithMatcher(t *testing.T) { progressBar, _ := progress.NewStatsTicker(0, false, false, false, 0) var firstInput, secondInput string - workflow := &Workflow{Options: &protocols.ExecuterOptions{Options: &types.Options{TemplateThreads: 10}}, Workflows: []*WorkflowTemplate{ - {Executers: []*ProtocolExecuterPair{{ + workflow := &workflows.Workflow{Options: &protocols.ExecuterOptions{Options: &types.Options{TemplateThreads: 10}}, Workflows: []*WorkflowTemplate{ + {Executers: []*workflows.ProtocolExecuterPair{{ Executer: &mockExecuter{result: true, executeHook: func(input string) { firstInput = input }, outputs: []*output.InternalWrappedEvent{ @@ -111,7 +113,7 @@ func TestWorkflowsSubtemplatesWithMatcher(t *testing.T) { Extracts: map[string][]string{}, }}, }}, Options: &protocols.ExecuterOptions{Progress: progressBar}}, - }, Matchers: []*Matcher{{Name: "tomcat", Subtemplates: []*WorkflowTemplate{{Executers: []*ProtocolExecuterPair{{ + }, Matchers: []*workflows.Matcher{{Name: "tomcat", Subtemplates: []*workflows.WorkflowTemplate{{Executers: []*workflows.ProtocolExecuterPair{{ Executer: &mockExecuter{result: true, executeHook: func(input string) { secondInput = input }}, Options: &protocols.ExecuterOptions{Progress: progressBar}}, @@ -129,8 +131,8 @@ func TestWorkflowsSubtemplatesWithMatcherNoMatch(t *testing.T) { progressBar, _ := progress.NewStatsTicker(0, false, false, false, 0) var firstInput, secondInput string - workflow := &Workflow{Options: &protocols.ExecuterOptions{Options: &types.Options{TemplateThreads: 10}}, Workflows: []*WorkflowTemplate{ - {Executers: []*ProtocolExecuterPair{{ + 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 }, outputs: []*output.InternalWrappedEvent{ @@ -139,7 +141,7 @@ func TestWorkflowsSubtemplatesWithMatcherNoMatch(t *testing.T) { Extracts: map[string][]string{}, }}, }}, Options: &protocols.ExecuterOptions{Progress: progressBar}}, - }, Matchers: []*Matcher{{Name: "apache", Subtemplates: []*WorkflowTemplate{{Executers: []*ProtocolExecuterPair{{ + }, Matchers: []*workflows.Matcher{{Name: "apache", Subtemplates: []*workflows.WorkflowTemplate{{Executers: []*workflows.ProtocolExecuterPair{{ Executer: &mockExecuter{result: true, executeHook: func(input string) { secondInput = input }}, Options: &protocols.ExecuterOptions{Progress: progressBar}}, @@ -187,3 +189,4 @@ func (m *mockExecuter) ExecuteWithResults(input string, callback protocols.Outpu } return nil } +*/ diff --git a/v2/pkg/engine/workpool.go b/v2/pkg/core/workpool.go similarity index 90% rename from v2/pkg/engine/workpool.go rename to v2/pkg/core/workpool.go index 04db44715..f307b205d 100644 --- a/v2/pkg/engine/workpool.go +++ b/v2/pkg/core/workpool.go @@ -1,4 +1,4 @@ -package engine +package core import ( "github.com/projectdiscovery/nuclei/v2/pkg/templates" @@ -25,8 +25,8 @@ type WorkPoolConfig struct { HeadlessTypeConcurrency int } -// New returns a new WorkPool instance -func New(config WorkPoolConfig) *WorkPool { +// NewWorkPool returns a new WorkPool instance +func NewWorkPool(config WorkPoolConfig) *WorkPool { return &WorkPool{config: config} } diff --git a/v2/pkg/engine/engine.go b/v2/pkg/engine/engine.go deleted file mode 100644 index 978e845ca..000000000 --- a/v2/pkg/engine/engine.go +++ /dev/null @@ -1,16 +0,0 @@ -package engine - -// Engine is an engine for running Nuclei Templates/Workflows. -// -// The engine contains multiple thread pools which allow using different -// concurrency values per protocol executed. This was something which was -// missing from the previous versions of nuclei. -type Engine struct { -} - -// InputProvider is an input provider interface for the nuclei execution -// engine. -// -// An example InputProvider is provided in form of hmap input provider. -type InputProvider interface { -} diff --git a/v2/pkg/engine/engine_test.go b/v2/pkg/engine/engine_test.go deleted file mode 100644 index 00a22ef62..000000000 --- a/v2/pkg/engine/engine_test.go +++ /dev/null @@ -1 +0,0 @@ -package engine diff --git a/v2/pkg/protocols/protocols.go b/v2/pkg/protocols/protocols.go index d77ba147e..83a76f7cf 100644 --- a/v2/pkg/protocols/protocols.go +++ b/v2/pkg/protocols/protocols.go @@ -64,6 +64,12 @@ type ExecuterOptions struct { WorkflowLoader model.WorkflowLoader } +// Copy returns a copy of the executeroptions structure +func (e ExecuterOptions) Copy() ExecuterOptions { + copy := e + return copy +} + // Request is an interface implemented any protocol based request generator. type Request interface { // Compile compiles the request generators preparing any requests possible. diff --git a/v2/pkg/templates/workflows.go b/v2/pkg/templates/workflows.go index 9d510f2e2..246c4bf8e 100644 --- a/v2/pkg/templates/workflows.go +++ b/v2/pkg/templates/workflows.go @@ -62,19 +62,7 @@ func parseWorkflowTemplate(workflow *workflows.WorkflowTemplate, preprocessor Pr return nil } for _, path := range paths { - opts := protocols.ExecuterOptions{ - Output: options.Output, - Options: options.Options, - Progress: options.Progress, - Catalog: options.Catalog, - Browser: options.Browser, - RateLimiter: options.RateLimiter, - IssuesClient: options.IssuesClient, - Interactsh: options.Interactsh, - ProjectFile: options.ProjectFile, - HostErrorsCache: options.HostErrorsCache, - } - template, err := Parse(path, preprocessor, opts) + template, err := Parse(path, preprocessor, options.Copy()) if err != nil { gologger.Warning().Msgf("Could not parse workflow template %s: %v\n", path, err) continue