From ff17d12cedf02ace02ed19e24c51876c05f55917 Mon Sep 17 00:00:00 2001
From: Tarun Koyalwar <45962551+tarunKoyalwar@users.noreply.github.com>
Date: Wed, 28 Dec 2022 02:18:00 +0530
Subject: [PATCH] adds scan strategy (#3075)

---
 v2/cmd/nuclei/main.go        |  5 +++
 v2/internal/runner/runner.go |  2 +-
 v2/pkg/core/execute.go       | 63 +++++++++++++++++++++++++++++-------
 v2/pkg/types/types.go        |  2 ++
 4 files changed, 60 insertions(+), 12 deletions(-)

diff --git a/v2/cmd/nuclei/main.go b/v2/cmd/nuclei/main.go
index 98f4e9d1c..fb3fe6bd7 100644
--- a/v2/cmd/nuclei/main.go
+++ b/v2/cmd/nuclei/main.go
@@ -246,6 +246,11 @@ on extensive configurability, massive extensibility and ease of use.`)
 		flagSet.StringVar(&options.ProjectPath, "project-path", os.TempDir(), "set a specific project path"),
 		flagSet.BoolVarP(&options.StopAtFirstMatch, "stop-at-first-match", "spm", false, "stop processing HTTP requests after the first match (may break template/workflow logic)"),
 		flagSet.BoolVar(&options.Stream, "stream", false, "stream mode - start elaborating without sorting the input"),
+		flagSet.EnumVarP(&options.ScanStrategy, "scan-strategy", "ss", goflags.EnumVariable(0), "strategy to use while scanning(auto/host-spray/template-spray)", goflags.AllowdTypes{
+			"auto":           goflags.EnumVariable(0),
+			"host-spray":     goflags.EnumVariable(1),
+			"template-spray": goflags.EnumVariable(2),
+		}),
 		flagSet.DurationVarP(&options.InputReadTimeout, "input-read-timeout", "irt", time.Duration(3*time.Minute), "timeout on input read"),
 		flagSet.BoolVarP(&options.DisableHTTPProbe, "no-httpx", "nh", false, "disable httpx probing for non-url input"),
 		flagSet.BoolVar(&options.DisableStdin, "no-stdin", false, "disable stdin processing"),
diff --git a/v2/internal/runner/runner.go b/v2/internal/runner/runner.go
index 8ce6e2c40..56c316f63 100644
--- a/v2/internal/runner/runner.go
+++ b/v2/internal/runner/runner.go
@@ -627,7 +627,7 @@ func (r *Runner) executeTemplatesInput(store *loader.Store, engine *core.Engine)
 	// tracks global progress and captures stdout/stderr until p.Wait finishes
 	r.progress.Init(r.hmapInputProvider.Count(), templateCount, totalRequests)
 
-	results := engine.ExecuteWithOpts(finalTemplates, r.hmapInputProvider, true)
+	results := engine.ExecuteScanWithOpts(finalTemplates, r.hmapInputProvider, true)
 	return results, nil
 }
 
diff --git a/v2/pkg/core/execute.go b/v2/pkg/core/execute.go
index 073fd37a1..f3b003f3f 100644
--- a/v2/pkg/core/execute.go
+++ b/v2/pkg/core/execute.go
@@ -5,6 +5,7 @@ import (
 	"go.uber.org/atomic"
 
 	"github.com/projectdiscovery/gologger"
+	"github.com/projectdiscovery/nuclei/v2/pkg/core/inputs"
 	"github.com/projectdiscovery/nuclei/v2/pkg/output"
 	"github.com/projectdiscovery/nuclei/v2/pkg/protocols/common/contextargs"
 	"github.com/projectdiscovery/nuclei/v2/pkg/templates"
@@ -18,20 +19,13 @@ import (
 // All the execution logic for the templates/workflows happens in this part
 // of the engine.
 func (e *Engine) Execute(templates []*templates.Template, target InputProvider) *atomic.Bool {
-	return e.ExecuteWithOpts(templates, target, false)
+	return e.ExecuteScanWithOpts(templates, target, false)
 }
 
-// ExecuteWithOpts executes with the full options
-func (e *Engine) ExecuteWithOpts(templatesList []*templates.Template, target InputProvider, noCluster bool) *atomic.Bool {
-	var finalTemplates []*templates.Template
-	if !noCluster {
-		finalTemplates, _ = templates.ClusterTemplates(templatesList, e.executerOpts)
-	} else {
-		finalTemplates = templatesList
-	}
-
+// executeTemplateSpray executes scan using template spray strategy where targets are iterated over each template
+func (e *Engine) executeTemplateSpray(templatesList []*templates.Template, target InputProvider) *atomic.Bool {
 	results := &atomic.Bool{}
-	for _, template := range finalTemplates {
+	for _, template := range templatesList {
 		templateType := template.Type()
 
 		var wg *sizedwaitgroup.SizedWaitGroup
@@ -59,6 +53,53 @@ func (e *Engine) ExecuteWithOpts(templatesList []*templates.Template, target Inp
 	return results
 }
 
+// executeHostSpray executes scan using host spray strategy where templates are iterated over each target
+func (e *Engine) executeHostSpray(templatesList []*templates.Template, target InputProvider) *atomic.Bool {
+	results := &atomic.Bool{}
+	hostwg := sizedwaitgroup.New(e.options.BulkSize)
+	target.Scan(func(value *contextargs.MetaInput) bool {
+		host := inputs.SimpleInputProvider{
+			Inputs: []*contextargs.MetaInput{
+				value,
+			},
+		}
+		hostwg.Add()
+		go func(result *atomic.Bool) {
+			defer hostwg.Done()
+			status := e.executeTemplateSpray(templatesList, &host)
+			results.CompareAndSwap(false, status.Load())
+		}(results)
+		return true
+	})
+	hostwg.Wait()
+	return results
+}
+
+// ExecuteScanWithOpts executes scan with given scanStatergy
+func (e *Engine) ExecuteScanWithOpts(templatesList []*templates.Template, target InputProvider, noCluster bool) *atomic.Bool {
+	var results *atomic.Bool
+
+	var finalTemplates []*templates.Template
+	if !noCluster {
+		finalTemplates, _ = templates.ClusterTemplates(templatesList, e.executerOpts)
+	} else {
+		finalTemplates = templatesList
+	}
+
+	if e.options.ScanStrategy == "auto" {
+		// TODO: this is only a placeholder, auto scan strategy should choose scan strategy
+		// based on no of hosts , templates , stream and other optimization parameters
+		e.options.ScanStrategy = "template-spray"
+	}
+	switch e.options.ScanStrategy {
+	case "template-spray":
+		results = e.executeTemplateSpray(finalTemplates, target)
+	case "host-spray":
+		results = e.executeHostSpray(finalTemplates, target)
+	}
+	return results
+}
+
 // processSelfContainedTemplates execute a self-contained template.
 func (e *Engine) executeSelfContainedTemplateWithInput(template *templates.Template, results *atomic.Bool) {
 	match, err := template.Executer.Execute(contextargs.New())
diff --git a/v2/pkg/types/types.go b/v2/pkg/types/types.go
index 6fc743ed1..ff05773a3 100644
--- a/v2/pkg/types/types.go
+++ b/v2/pkg/types/types.go
@@ -336,6 +336,8 @@ type Options struct {
 	AwsBucketName string
 	// AWS Region name where aws s3 bucket is located
 	AwsRegion string
+	// Scan Strategy (auto,hosts-spray,templates-spray)
+	ScanStrategy string
 }
 
 func (options *Options) AddVarPayload(key string, value interface{}) {