nuclei/v2/pkg/core/execute_options.go

125 lines
4.1 KiB
Go

package core
import (
"sync"
"sync/atomic"
"github.com/remeh/sizedwaitgroup"
"github.com/projectdiscovery/nuclei/v2/pkg/output"
"github.com/projectdiscovery/nuclei/v2/pkg/protocols/common/contextargs"
"github.com/projectdiscovery/nuclei/v2/pkg/templates"
"github.com/projectdiscovery/nuclei/v2/pkg/templates/types"
"github.com/projectdiscovery/nuclei/v2/pkg/types/scanstrategy"
stringsutil "github.com/projectdiscovery/utils/strings"
)
// Execute takes a list of templates/workflows that have been compiled
// and executes them based on provided concurrency options.
//
// 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.ExecuteScanWithOpts(templates, target, false)
}
// ExecuteWithResults a list of templates with results
func (e *Engine) ExecuteWithResults(templatesList []*templates.Template, target InputProvider, callback func(*output.ResultEvent)) *atomic.Bool {
e.Callback = callback
return e.ExecuteScanWithOpts(templatesList, target, false)
}
// ExecuteScanWithOpts executes scan with given scanStatergy
func (e *Engine) ExecuteScanWithOpts(templatesList []*templates.Template, target InputProvider, noCluster bool) *atomic.Bool {
results := &atomic.Bool{}
selfcontainedWg := &sync.WaitGroup{}
var finalTemplates []*templates.Template
if !noCluster {
finalTemplates, _ = templates.ClusterTemplates(templatesList, e.executerOpts)
} else {
finalTemplates = templatesList
}
if stringsutil.EqualFoldAny(e.options.ScanStrategy, scanstrategy.Auto.String(), "") {
// 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 = scanstrategy.TemplateSpray.String()
}
filtered := []*templates.Template{}
selfContained := []*templates.Template{}
// Filter Self Contained templates since they are not bound to target
for _, v := range finalTemplates {
if v.SelfContained {
selfContained = append(selfContained, v)
} else {
filtered = append(filtered, v)
}
}
// Execute All SelfContained in parallel
e.executeAllSelfContained(selfContained, results, selfcontainedWg)
strategyResult := &atomic.Bool{}
switch e.options.ScanStrategy {
case scanstrategy.TemplateSpray.String():
strategyResult = e.executeTemplateSpray(filtered, target)
case scanstrategy.HostSpray.String():
strategyResult = e.executeHostSpray(filtered, target)
}
results.CompareAndSwap(false, strategyResult.Load())
selfcontainedWg.Wait()
return results
}
// 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{}
// wp is workpool that contains different waitgroups for
// headless and non-headless templates
wp := e.GetWorkPool()
for _, template := range templatesList {
templateType := template.Type()
var wg *sizedwaitgroup.SizedWaitGroup
if templateType == types.HeadlessProtocol {
wg = wp.Headless
} else {
wg = wp.Default
}
wg.Add()
go func(tpl *templates.Template) {
defer wg.Done()
// All other request types are executed here
// Note: executeTemplateWithTargets creates goroutines and blocks
// given template is executed on all targets
e.executeTemplateWithTargets(tpl, target, results)
}(template)
}
wp.Wait()
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{}
wp := sizedwaitgroup.New(e.options.BulkSize + e.options.HeadlessBulkSize)
target.Scan(func(value *contextargs.MetaInput) bool {
wp.Add()
go func(targetval *contextargs.MetaInput) {
defer wp.Done()
e.executeTemplatesOnTarget(templatesList, targetval, results)
}(value)
return true
})
wp.Wait()
return results
}