mirror of
https://github.com/projectdiscovery/nuclei.git
synced 2025-12-17 19:35:27 +00:00
311 lines
9.7 KiB
Go
311 lines
9.7 KiB
Go
package core
|
|
|
|
import (
|
|
"github.com/remeh/sizedwaitgroup"
|
|
"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"
|
|
"github.com/projectdiscovery/nuclei/v2/pkg/templates/types"
|
|
generalTypes "github.com/projectdiscovery/nuclei/v2/pkg/types"
|
|
)
|
|
|
|
// 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)
|
|
}
|
|
|
|
// 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 templatesList {
|
|
templateType := template.Type()
|
|
|
|
var wg *sizedwaitgroup.SizedWaitGroup
|
|
if templateType == types.HeadlessProtocol {
|
|
wg = e.workPool.Headless
|
|
} else {
|
|
wg = e.workPool.Default
|
|
}
|
|
|
|
wg.Add()
|
|
go func(tpl *templates.Template) {
|
|
defer wg.Done()
|
|
|
|
switch {
|
|
case tpl.SelfContained:
|
|
// Self Contained requests are executed here separately
|
|
e.executeSelfContainedTemplateWithInput(tpl, results)
|
|
default:
|
|
// All other request types are executed here
|
|
e.executeModelWithInput(templateType, tpl, target, results)
|
|
}
|
|
}(template)
|
|
}
|
|
e.workPool.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{}
|
|
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())
|
|
if err != nil {
|
|
gologger.Warning().Msgf("[%s] Could not execute step: %s\n", e.executerOpts.Colorizer.BrightBlue(template.ID), err)
|
|
}
|
|
results.CompareAndSwap(false, match)
|
|
}
|
|
|
|
// executeModelWithInput executes a type of template with input
|
|
func (e *Engine) executeModelWithInput(templateType types.ProtocolType, template *templates.Template, target InputProvider, results *atomic.Bool) {
|
|
wg := e.workPool.InputPool(templateType)
|
|
|
|
var (
|
|
index uint32
|
|
)
|
|
|
|
e.executerOpts.ResumeCfg.Lock()
|
|
currentInfo, ok := e.executerOpts.ResumeCfg.Current[template.ID]
|
|
if !ok {
|
|
currentInfo = &generalTypes.ResumeInfo{}
|
|
e.executerOpts.ResumeCfg.Current[template.ID] = currentInfo
|
|
}
|
|
if currentInfo.InFlight == nil {
|
|
currentInfo.InFlight = make(map[uint32]struct{})
|
|
}
|
|
resumeFromInfo, ok := e.executerOpts.ResumeCfg.ResumeFrom[template.ID]
|
|
if !ok {
|
|
resumeFromInfo = &generalTypes.ResumeInfo{}
|
|
e.executerOpts.ResumeCfg.ResumeFrom[template.ID] = resumeFromInfo
|
|
}
|
|
e.executerOpts.ResumeCfg.Unlock()
|
|
|
|
// track progression
|
|
cleanupInFlight := func(index uint32) {
|
|
currentInfo.Lock()
|
|
delete(currentInfo.InFlight, index)
|
|
currentInfo.Unlock()
|
|
}
|
|
|
|
target.Scan(func(scannedValue *contextargs.MetaInput) bool {
|
|
// Best effort to track the host progression
|
|
// skips indexes lower than the minimum in-flight at interruption time
|
|
var skip bool
|
|
if resumeFromInfo.Completed { // the template was completed
|
|
gologger.Debug().Msgf("[%s] Skipping \"%s\": Resume - Template already completed\n", template.ID, scannedValue)
|
|
skip = true
|
|
} else if index < resumeFromInfo.SkipUnder { // index lower than the sliding window (bulk-size)
|
|
gologger.Debug().Msgf("[%s] Skipping \"%s\": Resume - Target already processed\n", template.ID, scannedValue)
|
|
skip = true
|
|
} else if _, isInFlight := resumeFromInfo.InFlight[index]; isInFlight { // the target wasn't completed successfully
|
|
gologger.Debug().Msgf("[%s] Repeating \"%s\": Resume - Target wasn't completed\n", template.ID, scannedValue)
|
|
// skip is already false, but leaving it here for clarity
|
|
skip = false
|
|
} else if index > resumeFromInfo.DoAbove { // index above the sliding window (bulk-size)
|
|
// skip is already false - but leaving it here for clarity
|
|
skip = false
|
|
}
|
|
|
|
currentInfo.Lock()
|
|
currentInfo.InFlight[index] = struct{}{}
|
|
currentInfo.Unlock()
|
|
|
|
// Skip if the host has had errors
|
|
if e.executerOpts.HostErrorsCache != nil && e.executerOpts.HostErrorsCache.Check(scannedValue.ID()) {
|
|
return true
|
|
}
|
|
|
|
wg.WaitGroup.Add()
|
|
go func(index uint32, skip bool, value *contextargs.MetaInput) {
|
|
defer wg.WaitGroup.Done()
|
|
defer cleanupInFlight(index)
|
|
if skip {
|
|
return
|
|
}
|
|
|
|
var match bool
|
|
var err error
|
|
switch templateType {
|
|
case types.WorkflowProtocol:
|
|
match = e.executeWorkflow(value, template.CompiledWorkflow)
|
|
default:
|
|
ctxArgs := contextargs.New()
|
|
ctxArgs.MetaInput = value
|
|
match, err = template.Executer.Execute(ctxArgs)
|
|
}
|
|
if err != nil {
|
|
gologger.Warning().Msgf("[%s] Could not execute step: %s\n", e.executerOpts.Colorizer.BrightBlue(template.ID), err)
|
|
}
|
|
results.CompareAndSwap(false, match)
|
|
}(index, skip, scannedValue)
|
|
|
|
index++
|
|
return true
|
|
})
|
|
wg.WaitGroup.Wait()
|
|
|
|
// on completion marks the template as completed
|
|
currentInfo.Lock()
|
|
currentInfo.Completed = true
|
|
currentInfo.Unlock()
|
|
}
|
|
|
|
// ExecuteWithResults a list of templates with results
|
|
func (e *Engine) ExecuteWithResults(templatesList []*templates.Template, target InputProvider, callback func(*output.ResultEvent)) *atomic.Bool {
|
|
results := &atomic.Bool{}
|
|
for _, template := range templatesList {
|
|
templateType := template.Type()
|
|
|
|
var wg *sizedwaitgroup.SizedWaitGroup
|
|
if templateType == types.HeadlessProtocol {
|
|
wg = e.workPool.Headless
|
|
} else {
|
|
wg = e.workPool.Default
|
|
}
|
|
|
|
wg.Add()
|
|
go func(tpl *templates.Template) {
|
|
e.executeModelWithInputAndResult(templateType, tpl, target, results, callback)
|
|
wg.Done()
|
|
}(template)
|
|
}
|
|
e.workPool.Wait()
|
|
return results
|
|
}
|
|
|
|
// executeModelWithInputAndResult executes a type of template with input and result
|
|
func (e *Engine) executeModelWithInputAndResult(templateType types.ProtocolType, template *templates.Template, target InputProvider, results *atomic.Bool, callback func(*output.ResultEvent)) {
|
|
wg := e.workPool.InputPool(templateType)
|
|
|
|
target.Scan(func(scannedValue *contextargs.MetaInput) bool {
|
|
// Skip if the host has had errors
|
|
if e.executerOpts.HostErrorsCache != nil && e.executerOpts.HostErrorsCache.Check(scannedValue.ID()) {
|
|
return true
|
|
}
|
|
|
|
wg.WaitGroup.Add()
|
|
go func(value *contextargs.MetaInput) {
|
|
defer wg.WaitGroup.Done()
|
|
|
|
var match bool
|
|
var err error
|
|
switch templateType {
|
|
case types.WorkflowProtocol:
|
|
match = e.executeWorkflow(value, template.CompiledWorkflow)
|
|
default:
|
|
ctxArgs := contextargs.New()
|
|
ctxArgs.MetaInput = value
|
|
err = template.Executer.ExecuteWithResults(ctxArgs, func(event *output.InternalWrappedEvent) {
|
|
for _, result := range event.Results {
|
|
callback(result)
|
|
}
|
|
})
|
|
}
|
|
if err != nil {
|
|
gologger.Warning().Msgf("[%s] Could not execute step: %s\n", e.executerOpts.Colorizer.BrightBlue(template.ID), err)
|
|
}
|
|
results.CompareAndSwap(false, match)
|
|
}(scannedValue)
|
|
return true
|
|
})
|
|
wg.WaitGroup.Wait()
|
|
}
|
|
|
|
type ChildExecuter struct {
|
|
e *Engine
|
|
|
|
results *atomic.Bool
|
|
}
|
|
|
|
// Close closes the executer returning bool results
|
|
func (e *ChildExecuter) Close() *atomic.Bool {
|
|
e.e.workPool.Wait()
|
|
return e.results
|
|
}
|
|
|
|
// Execute executes a template and URLs
|
|
func (e *ChildExecuter) Execute(template *templates.Template, value *contextargs.MetaInput) {
|
|
templateType := template.Type()
|
|
|
|
var wg *sizedwaitgroup.SizedWaitGroup
|
|
if templateType == types.HeadlessProtocol {
|
|
wg = e.e.workPool.Headless
|
|
} else {
|
|
wg = e.e.workPool.Default
|
|
}
|
|
|
|
wg.Add()
|
|
go func(tpl *templates.Template) {
|
|
defer wg.Done()
|
|
|
|
ctxArgs := contextargs.New()
|
|
ctxArgs.MetaInput = value
|
|
match, err := template.Executer.Execute(ctxArgs)
|
|
if err != nil {
|
|
gologger.Warning().Msgf("[%s] Could not execute step: %s\n", e.e.executerOpts.Colorizer.BrightBlue(template.ID), err)
|
|
}
|
|
e.results.CompareAndSwap(false, match)
|
|
}(template)
|
|
}
|
|
|
|
// ExecuteWithOpts executes with the full options
|
|
func (e *Engine) ChildExecuter() *ChildExecuter {
|
|
return &ChildExecuter{
|
|
e: e,
|
|
results: &atomic.Bool{},
|
|
}
|
|
}
|