mirror of
https://github.com/projectdiscovery/nuclei.git
synced 2026-11-08 10:10:18 -10:00
122 lines
4.0 KiB
Go
122 lines
4.0 KiB
Go
package core
|
|
|
|
import (
|
|
"sync"
|
|
|
|
"github.com/projectdiscovery/gologger"
|
|
"github.com/projectdiscovery/nuclei/v3/pkg/output"
|
|
"github.com/projectdiscovery/nuclei/v3/pkg/protocols"
|
|
"github.com/projectdiscovery/nuclei/v3/pkg/templates"
|
|
"github.com/projectdiscovery/nuclei/v3/pkg/types"
|
|
)
|
|
|
|
// TemplateExecutionState identifies a template execution lifecycle transition.
|
|
type TemplateExecutionState uint8
|
|
|
|
const (
|
|
// TemplateExecutionStarted is emitted immediately before a template starts on a target.
|
|
TemplateExecutionStarted TemplateExecutionState = iota + 1
|
|
// TemplateExecutionFinished is emitted after that template execution returns.
|
|
TemplateExecutionFinished
|
|
)
|
|
|
|
// TemplateExecutionEvent identifies one template execution on one target.
|
|
type TemplateExecutionEvent struct {
|
|
TemplateID string
|
|
TemplatePath string
|
|
Target string
|
|
State TemplateExecutionState
|
|
ContextErr error
|
|
}
|
|
|
|
// TemplateExecutionCallback observes template execution lifecycle transitions.
|
|
// Nuclei can invoke the callback concurrently; implementations must be concurrency-safe.
|
|
type TemplateExecutionCallback func(TemplateExecutionEvent)
|
|
|
|
// Engine is an executer 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 work pool, it is
|
|
// handled by the engine.
|
|
type Engine struct {
|
|
workPool *WorkPool
|
|
// originScoped tracks origin-scoped template attempts. A pair is finished
|
|
// only after a successful execution, so a failed first target does not
|
|
// suppress every later URL on that origin.
|
|
originScoped sync.Map
|
|
options *types.Options
|
|
executerOpts *protocols.ExecutorOptions
|
|
Callback func(*output.ResultEvent) // Executed on results
|
|
templateExecutionCallback TemplateExecutionCallback
|
|
Logger *gologger.Logger
|
|
}
|
|
|
|
// New returns a new Engine instance
|
|
func New(options *types.Options) *Engine {
|
|
engine := &Engine{
|
|
options: options,
|
|
Logger: options.Logger,
|
|
}
|
|
engine.workPool = engine.GetWorkPool()
|
|
return engine
|
|
}
|
|
|
|
func (e *Engine) GetWorkPoolConfig() WorkPoolConfig {
|
|
config := WorkPoolConfig{
|
|
InputConcurrency: e.options.BulkSize,
|
|
TypeConcurrency: e.options.CurrentTemplateThreads(),
|
|
HeadlessInputConcurrency: e.options.HeadlessBulkSize,
|
|
HeadlessTypeConcurrency: e.options.HeadlessTemplateThreads,
|
|
}
|
|
return config
|
|
}
|
|
|
|
// GetWorkPool returns a workpool from options
|
|
func (e *Engine) GetWorkPool() *WorkPool {
|
|
return NewWorkPool(e.GetWorkPoolConfig())
|
|
}
|
|
|
|
// 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.ExecutorOptions) {
|
|
e.executerOpts = options
|
|
}
|
|
|
|
// SetTemplateExecutionCallback registers a callback for template execution lifecycle events.
|
|
func (e *Engine) SetTemplateExecutionCallback(callback TemplateExecutionCallback) {
|
|
e.templateExecutionCallback = callback
|
|
}
|
|
|
|
func (e *Engine) templateExecutionStarted(template *templates.Template, target string) func(error) {
|
|
if e.templateExecutionCallback == nil {
|
|
return func(error) {}
|
|
}
|
|
event := TemplateExecutionEvent{
|
|
TemplateID: template.ID,
|
|
TemplatePath: template.Path,
|
|
Target: target,
|
|
State: TemplateExecutionStarted,
|
|
}
|
|
e.templateExecutionCallback(event)
|
|
return func(contextErr error) {
|
|
event.State = TemplateExecutionFinished
|
|
event.ContextErr = contextErr
|
|
e.templateExecutionCallback(event)
|
|
}
|
|
}
|
|
|
|
// ExecuterOptions returns protocols.ExecutorOptions for nuclei engine.
|
|
func (e *Engine) ExecuterOptions() *protocols.ExecutorOptions {
|
|
return e.executerOpts
|
|
}
|
|
|
|
// WorkPool returns the worker pool for the engine
|
|
func (e *Engine) WorkPool() *WorkPool {
|
|
// resize check point - nop if there are no changes
|
|
e.workPool.RefreshWithConfig(e.GetWorkPoolConfig())
|
|
return e.workPool
|
|
}
|