adds scan strategy (#3075)

This commit is contained in:
Tarun Koyalwar 2022-12-28 02:18:00 +05:30 committed by GitHub
parent af2aaacbc9
commit ff17d12ced
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 60 additions and 12 deletions

View File

@ -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"),

View File

@ -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
}

View File

@ -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())

View File

@ -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{}) {