mirror of https://github.com/daffainfo/nuclei.git
Moved all important execution stuff to engine
parent
df78ea72c5
commit
d124dbacc7
|
@ -9,8 +9,6 @@ import (
|
||||||
|
|
||||||
"github.com/logrusorgru/aurora"
|
"github.com/logrusorgru/aurora"
|
||||||
"github.com/pkg/errors"
|
"github.com/pkg/errors"
|
||||||
"github.com/remeh/sizedwaitgroup"
|
|
||||||
"go.uber.org/atomic"
|
|
||||||
"go.uber.org/ratelimit"
|
"go.uber.org/ratelimit"
|
||||||
"gopkg.in/yaml.v2"
|
"gopkg.in/yaml.v2"
|
||||||
|
|
||||||
|
@ -20,7 +18,7 @@ import (
|
||||||
"github.com/projectdiscovery/nuclei/v2/pkg/catalog/config"
|
"github.com/projectdiscovery/nuclei/v2/pkg/catalog/config"
|
||||||
"github.com/projectdiscovery/nuclei/v2/pkg/catalog/loader"
|
"github.com/projectdiscovery/nuclei/v2/pkg/catalog/loader"
|
||||||
"github.com/projectdiscovery/nuclei/v2/pkg/core"
|
"github.com/projectdiscovery/nuclei/v2/pkg/core"
|
||||||
"github.com/projectdiscovery/nuclei/v2/pkg/engine/inputs/hmap"
|
"github.com/projectdiscovery/nuclei/v2/pkg/engine/inputs/hybrid"
|
||||||
"github.com/projectdiscovery/nuclei/v2/pkg/model/types/severity"
|
"github.com/projectdiscovery/nuclei/v2/pkg/model/types/severity"
|
||||||
"github.com/projectdiscovery/nuclei/v2/pkg/output"
|
"github.com/projectdiscovery/nuclei/v2/pkg/output"
|
||||||
"github.com/projectdiscovery/nuclei/v2/pkg/parsers"
|
"github.com/projectdiscovery/nuclei/v2/pkg/parsers"
|
||||||
|
@ -34,7 +32,6 @@ import (
|
||||||
"github.com/projectdiscovery/nuclei/v2/pkg/reporting"
|
"github.com/projectdiscovery/nuclei/v2/pkg/reporting"
|
||||||
"github.com/projectdiscovery/nuclei/v2/pkg/reporting/exporters/markdown"
|
"github.com/projectdiscovery/nuclei/v2/pkg/reporting/exporters/markdown"
|
||||||
"github.com/projectdiscovery/nuclei/v2/pkg/reporting/exporters/sarif"
|
"github.com/projectdiscovery/nuclei/v2/pkg/reporting/exporters/sarif"
|
||||||
"github.com/projectdiscovery/nuclei/v2/pkg/templates"
|
|
||||||
"github.com/projectdiscovery/nuclei/v2/pkg/types"
|
"github.com/projectdiscovery/nuclei/v2/pkg/types"
|
||||||
"github.com/projectdiscovery/nuclei/v2/pkg/utils"
|
"github.com/projectdiscovery/nuclei/v2/pkg/utils"
|
||||||
"github.com/projectdiscovery/nuclei/v2/pkg/utils/stats"
|
"github.com/projectdiscovery/nuclei/v2/pkg/utils/stats"
|
||||||
|
@ -52,7 +49,7 @@ type Runner struct {
|
||||||
colorizer aurora.Aurora
|
colorizer aurora.Aurora
|
||||||
issuesClient *reporting.Client
|
issuesClient *reporting.Client
|
||||||
addColor func(severity.Severity) string
|
addColor func(severity.Severity) string
|
||||||
hmapInputProvider *hmap.Input
|
hmapInputProvider *hybrid.Input
|
||||||
browser *engine.Browser
|
browser *engine.Browser
|
||||||
ratelimiter ratelimit.Limiter
|
ratelimiter ratelimit.Limiter
|
||||||
hostErrors *hosterrorscache.Cache
|
hostErrors *hosterrorscache.Cache
|
||||||
|
@ -112,7 +109,7 @@ func New(options *types.Options) (*Runner, error) {
|
||||||
}
|
}
|
||||||
|
|
||||||
// Initialize the input source
|
// Initialize the input source
|
||||||
hmapInput, err := hmap.New(options)
|
hmapInput, err := hybrid.New(options)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, errors.Wrap(err, "could not create input provider")
|
return nil, errors.Wrap(err, "could not create input provider")
|
||||||
}
|
}
|
||||||
|
@ -259,6 +256,7 @@ func (r *Runner) RunEnumeration() error {
|
||||||
ProjectFile: r.projectFile,
|
ProjectFile: r.projectFile,
|
||||||
Browser: r.browser,
|
Browser: r.browser,
|
||||||
HostErrorsCache: cache,
|
HostErrorsCache: cache,
|
||||||
|
Colorizer: r.colorizer,
|
||||||
}
|
}
|
||||||
engine := core.New(r.options)
|
engine := core.New(r.options)
|
||||||
engine.SetExecuterOptions(executerOpts)
|
engine.SetExecuterOptions(executerOpts)
|
||||||
|
@ -351,9 +349,6 @@ func (r *Runner) RunEnumeration() error {
|
||||||
gologger.Info().Msgf("Workflows loaded for scan: %d", len(store.Workflows()))
|
gologger.Info().Msgf("Workflows loaded for scan: %d", len(store.Workflows()))
|
||||||
}
|
}
|
||||||
|
|
||||||
// pre-parse all the templates, apply filters
|
|
||||||
finalTemplates := []*templates.Template{}
|
|
||||||
|
|
||||||
var unclusteredRequests int64
|
var unclusteredRequests int64
|
||||||
for _, template := range store.Templates() {
|
for _, template := range store.Templates() {
|
||||||
// workflows will dynamically adjust the totals while running, as
|
// workflows will dynamically adjust the totals while running, as
|
||||||
|
@ -373,6 +368,12 @@ func (r *Runner) RunEnumeration() error {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Cluster the templates first because we want info on how many
|
||||||
|
// templates did we cluster for showing to user in CLI
|
||||||
|
originalTemplatesCount := len(store.Templates())
|
||||||
|
finalTemplates, clusterCount := engine.ClusterTemplates(store.Templates())
|
||||||
|
finalTemplates = append(finalTemplates, store.Workflows()...)
|
||||||
|
|
||||||
var totalRequests int64
|
var totalRequests int64
|
||||||
for _, t := range finalTemplates {
|
for _, t := range finalTemplates {
|
||||||
if len(t.Workflows) > 0 {
|
if len(t.Workflows) > 0 {
|
||||||
|
@ -391,32 +392,10 @@ func (r *Runner) RunEnumeration() error {
|
||||||
return errors.New("no valid templates were found")
|
return errors.New("no valid templates were found")
|
||||||
}
|
}
|
||||||
|
|
||||||
/*
|
|
||||||
TODO does it make sense to run the logic below if there are no targets specified?
|
|
||||||
Can we safely assume the user is just experimenting with the template/workflow filters before running them?
|
|
||||||
*/
|
|
||||||
|
|
||||||
results := &atomic.Bool{}
|
|
||||||
wgtemplates := sizedwaitgroup.New(r.options.TemplateThreads)
|
|
||||||
|
|
||||||
// tracks global progress and captures stdout/stderr until p.Wait finishes
|
// tracks global progress and captures stdout/stderr until p.Wait finishes
|
||||||
r.progress.Init(r.hmapInputProvider.Count(), templateCount, totalRequests)
|
r.progress.Init(r.hmapInputProvider.Count(), templateCount, totalRequests)
|
||||||
|
|
||||||
for _, t := range finalTemplates {
|
results := engine.ExecuteWithOpts(finalTemplates, r.hmapInputProvider, true)
|
||||||
wgtemplates.Add()
|
|
||||||
go func(template *templates.Template) {
|
|
||||||
defer wgtemplates.Done()
|
|
||||||
|
|
||||||
if template.SelfContained {
|
|
||||||
results.CAS(false, r.processSelfContainedTemplates(template))
|
|
||||||
} else if len(template.Workflows) > 0 {
|
|
||||||
results.CAS(false, r.processWorkflowWithList(template))
|
|
||||||
} else {
|
|
||||||
results.CAS(false, r.processTemplateWithList(template))
|
|
||||||
}
|
|
||||||
}(t)
|
|
||||||
}
|
|
||||||
wgtemplates.Wait()
|
|
||||||
|
|
||||||
if r.interactsh != nil {
|
if r.interactsh != nil {
|
||||||
matched := r.interactsh.Close()
|
matched := r.interactsh.Close()
|
||||||
|
|
|
@ -24,6 +24,10 @@ type Engine struct {
|
||||||
//
|
//
|
||||||
// An example InputProvider is provided in form of hmap input provider.
|
// An example InputProvider is provided in form of hmap input provider.
|
||||||
type InputProvider interface {
|
type InputProvider interface {
|
||||||
|
// Count returns the number of items for input provider
|
||||||
|
Count() int64
|
||||||
|
// Scan calls a callback function till the input provider is exhausted
|
||||||
|
Scan(callback func(value string))
|
||||||
}
|
}
|
||||||
|
|
||||||
// New returns a new Engine instance
|
// New returns a new Engine instance
|
||||||
|
|
|
@ -3,9 +3,12 @@ package core
|
||||||
import (
|
import (
|
||||||
"fmt"
|
"fmt"
|
||||||
|
|
||||||
|
"github.com/projectdiscovery/gologger"
|
||||||
"github.com/projectdiscovery/nuclei/v2/pkg/protocols/common/clusterer"
|
"github.com/projectdiscovery/nuclei/v2/pkg/protocols/common/clusterer"
|
||||||
"github.com/projectdiscovery/nuclei/v2/pkg/templates"
|
"github.com/projectdiscovery/nuclei/v2/pkg/templates"
|
||||||
|
"github.com/remeh/sizedwaitgroup"
|
||||||
"github.com/rs/xid"
|
"github.com/rs/xid"
|
||||||
|
"go.uber.org/atomic"
|
||||||
)
|
)
|
||||||
|
|
||||||
// Execute takes a list of templates/workflows that have been compiled
|
// Execute takes a list of templates/workflows that have been compiled
|
||||||
|
@ -13,27 +16,86 @@ import (
|
||||||
//
|
//
|
||||||
// All the execution logic for the templates/workflows happens in this part
|
// All the execution logic for the templates/workflows happens in this part
|
||||||
// of the engine.
|
// of the engine.
|
||||||
func (e *Engine) Execute(templates []*templates.Template) {
|
func (e *Engine) Execute(templates []*templates.Template, input InputProvider) *atomic.Bool {
|
||||||
finalTemplates, clusterCount := e.clusterTemplates(templates)
|
return e.ExecuteWithOpts(templates, input, false)
|
||||||
|
}
|
||||||
|
|
||||||
|
// ExecuteWithOpts is execute with the full options
|
||||||
|
func (e *Engine) ExecuteWithOpts(templatesList []*templates.Template, input InputProvider, noCluster bool) *atomic.Bool {
|
||||||
|
var finalTemplates []*templates.Template
|
||||||
|
if !noCluster {
|
||||||
|
finalTemplates, _ = e.ClusterTemplates(templatesList)
|
||||||
|
} else {
|
||||||
|
finalTemplates = templatesList
|
||||||
|
}
|
||||||
|
|
||||||
|
results := &atomic.Bool{}
|
||||||
for _, template := range finalTemplates {
|
for _, template := range finalTemplates {
|
||||||
templateType := template.Type()
|
templateType := template.Type()
|
||||||
|
|
||||||
|
var wg *sizedwaitgroup.SizedWaitGroup
|
||||||
|
if templateType == "headless" {
|
||||||
|
wg = e.workPool.Headless
|
||||||
|
} else {
|
||||||
|
wg = e.workPool.Default
|
||||||
|
}
|
||||||
|
|
||||||
|
wg.Add()
|
||||||
switch {
|
switch {
|
||||||
case template.SelfContained:
|
case template.SelfContained:
|
||||||
// Self Contained requests are executed here
|
// Self Contained requests are executed here separately
|
||||||
|
e.executeSelfContainedTemplateWithInput(template, results)
|
||||||
case templateType == "workflow":
|
|
||||||
// Workflows requests are executed here
|
|
||||||
|
|
||||||
default:
|
default:
|
||||||
// All other request types are executed here
|
// All other request types are executed here
|
||||||
|
e.executeModelWithInput(templateType, template, input, results)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
e.workPool.Wait()
|
||||||
|
return results
|
||||||
}
|
}
|
||||||
|
|
||||||
// clusterTemplates performs identical http requests clustering for a list of templates
|
// processSelfContainedTemplates execute a self-contained template.
|
||||||
func (e *Engine) clusterTemplates(templatesList []*templates.Template) ([]*templates.Template, int) {
|
func (e *Engine) executeSelfContainedTemplateWithInput(template *templates.Template, results *atomic.Bool) {
|
||||||
|
match, err := template.Executer.Execute("")
|
||||||
|
if err != nil {
|
||||||
|
gologger.Warning().Msgf("[%s] Could not execute step: %s\n", e.executerOpts.Colorizer.BrightBlue(template.ID), err)
|
||||||
|
}
|
||||||
|
results.CAS(false, match)
|
||||||
|
}
|
||||||
|
|
||||||
|
// executeModelWithInput executes a type of template with input
|
||||||
|
func (e *Engine) executeModelWithInput(templateType string, template *templates.Template, input InputProvider, results *atomic.Bool) {
|
||||||
|
wg := e.workPool.InputPool(templateType)
|
||||||
|
|
||||||
|
input.Scan(func(scannedValue string) {
|
||||||
|
// Skip if the host has had errors
|
||||||
|
if e.executerOpts.HostErrorsCache != nil && e.executerOpts.HostErrorsCache.Check(scannedValue) {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
wg.Waitgroup.Add()
|
||||||
|
go func(value string) {
|
||||||
|
defer wg.Waitgroup.Done()
|
||||||
|
|
||||||
|
var match bool
|
||||||
|
var err error
|
||||||
|
switch templateType {
|
||||||
|
case "workflow":
|
||||||
|
match = e.executeWorkflow(value, template.CompiledWorkflow)
|
||||||
|
default:
|
||||||
|
match, err = template.Executer.Execute(value)
|
||||||
|
}
|
||||||
|
if err != nil {
|
||||||
|
gologger.Warning().Msgf("[%s] Could not execute step: %s\n", e.executerOpts.Colorizer.BrightBlue(template.ID), err)
|
||||||
|
}
|
||||||
|
results.CAS(false, match)
|
||||||
|
}(scannedValue)
|
||||||
|
})
|
||||||
|
wg.Waitgroup.Wait()
|
||||||
|
}
|
||||||
|
|
||||||
|
// ClusterTemplates performs identical http requests clustering for a list of templates
|
||||||
|
func (e *Engine) ClusterTemplates(templatesList []*templates.Template) ([]*templates.Template, int) {
|
||||||
if e.options.OfflineHTTP {
|
if e.options.OfflineHTTP {
|
||||||
return templatesList, 0
|
return templatesList, 0
|
||||||
}
|
}
|
||||||
|
@ -65,85 +127,3 @@ func (e *Engine) clusterTemplates(templatesList []*templates.Template) ([]*templ
|
||||||
}
|
}
|
||||||
return finalTemplatesList, clusterCount
|
return finalTemplatesList, clusterCount
|
||||||
}
|
}
|
||||||
|
|
||||||
/*
|
|
||||||
import (
|
|
||||||
"github.com/projectdiscovery/gologger"
|
|
||||||
"github.com/projectdiscovery/nuclei/v2/pkg/templates"
|
|
||||||
"github.com/remeh/sizedwaitgroup"
|
|
||||||
"go.uber.org/atomic"
|
|
||||||
)
|
|
||||||
|
|
||||||
// processSelfContainedTemplates execute a self-contained template.
|
|
||||||
func (r *Runner) processSelfContainedTemplates(template *templates.Template) bool {
|
|
||||||
match, err := template.Executer.Execute("")
|
|
||||||
if err != nil {
|
|
||||||
gologger.Warning().Msgf("[%s] Could not execute step: %s\n", r.colorizer.BrightBlue(template.ID), err)
|
|
||||||
}
|
|
||||||
return match
|
|
||||||
}
|
|
||||||
|
|
||||||
// processTemplateWithList execute a template against the list of user provided targets
|
|
||||||
func (r *Runner) processTemplateWithList(template *templates.Template) bool {
|
|
||||||
results := &atomic.Bool{}
|
|
||||||
wg := sizedwaitgroup.New(r.options.BulkSize)
|
|
||||||
processItem := func(k, _ []byte) error {
|
|
||||||
URL := string(k)
|
|
||||||
|
|
||||||
// Skip if the host has had errors
|
|
||||||
if r.hostErrors != nil && r.hostErrors.Check(URL) {
|
|
||||||
return nil
|
|
||||||
}
|
|
||||||
wg.Add()
|
|
||||||
go func(URL string) {
|
|
||||||
defer wg.Done()
|
|
||||||
|
|
||||||
match, err := template.Executer.Execute(URL)
|
|
||||||
if err != nil {
|
|
||||||
gologger.Warning().Msgf("[%s] Could not execute step: %s\n", r.colorizer.BrightBlue(template.ID), err)
|
|
||||||
}
|
|
||||||
results.CAS(false, match)
|
|
||||||
}(URL)
|
|
||||||
return nil
|
|
||||||
}
|
|
||||||
if r.options.Stream {
|
|
||||||
_ = r.hostMapStream.Scan(processItem)
|
|
||||||
} else {
|
|
||||||
r.hostMap.Scan(processItem)
|
|
||||||
}
|
|
||||||
|
|
||||||
wg.Wait()
|
|
||||||
return results.Load()
|
|
||||||
}
|
|
||||||
|
|
||||||
// processTemplateWithList process a template on the URL list
|
|
||||||
func (r *Runner) processWorkflowWithList(template *templates.Template) bool {
|
|
||||||
results := &atomic.Bool{}
|
|
||||||
wg := sizedwaitgroup.New(r.options.BulkSize)
|
|
||||||
|
|
||||||
processItem := func(k, _ []byte) error {
|
|
||||||
URL := string(k)
|
|
||||||
|
|
||||||
// Skip if the host has had errors
|
|
||||||
if r.hostErrors != nil && r.hostErrors.Check(URL) {
|
|
||||||
return nil
|
|
||||||
}
|
|
||||||
wg.Add()
|
|
||||||
go func(URL string) {
|
|
||||||
defer wg.Done()
|
|
||||||
match := template.CompiledWorkflow.RunWorkflow(URL)
|
|
||||||
results.CAS(false, match)
|
|
||||||
}(URL)
|
|
||||||
return nil
|
|
||||||
}
|
|
||||||
|
|
||||||
if r.options.Stream {
|
|
||||||
_ = r.hostMapStream.Scan(processItem)
|
|
||||||
} else {
|
|
||||||
r.hostMap.Scan(processItem)
|
|
||||||
}
|
|
||||||
|
|
||||||
wg.Wait()
|
|
||||||
return results.Load()
|
|
||||||
}
|
|
||||||
*/
|
|
||||||
|
|
|
@ -1,6 +1,6 @@
|
||||||
// Package hmap implements a hybrid hmap/filekv backed input provider
|
// Package hybrid implements a hybrid hmap/filekv backed input provider
|
||||||
// for nuclei that can either stream or store results using different kv stores.
|
// for nuclei that can either stream or store results using different kv stores.
|
||||||
package hmap
|
package hybrid
|
||||||
|
|
||||||
import (
|
import (
|
||||||
"bufio"
|
"bufio"
|
||||||
|
@ -118,3 +118,16 @@ func (i *Input) normalizeStoreInputValue(value string) {
|
||||||
func (i *Input) Count() int64 {
|
func (i *Input) Count() int64 {
|
||||||
return i.inputCount
|
return i.inputCount
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Scan calls an input provider till the callback is exhausted
|
||||||
|
func (i *Input) Scan(callback func(value string)) {
|
||||||
|
callbackFunc := func(k, _ []byte) error {
|
||||||
|
callback(string(k))
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
if i.hostMapStream != nil {
|
||||||
|
_ = i.hostMapStream.Scan(callbackFunc)
|
||||||
|
} else {
|
||||||
|
i.hostMap.Scan(callbackFunc)
|
||||||
|
}
|
||||||
|
}
|
|
@ -1,15 +1,22 @@
|
||||||
package core
|
package core
|
||||||
|
|
||||||
/*
|
import (
|
||||||
// RunWorkflow runs a workflow on an input and returns true or false
|
"github.com/projectdiscovery/gologger"
|
||||||
func (w *Workflow) RunWorkflow(input string) bool {
|
"github.com/projectdiscovery/nuclei/v2/pkg/output"
|
||||||
|
"github.com/projectdiscovery/nuclei/v2/pkg/workflows"
|
||||||
|
"github.com/remeh/sizedwaitgroup"
|
||||||
|
"go.uber.org/atomic"
|
||||||
|
)
|
||||||
|
|
||||||
|
// executeWorkflow runs a workflow on an input and returns true or false
|
||||||
|
func (e *Engine) executeWorkflow(input string, w *workflows.Workflow) bool {
|
||||||
results := &atomic.Bool{}
|
results := &atomic.Bool{}
|
||||||
|
|
||||||
swg := sizedwaitgroup.New(w.Options.Options.TemplateThreads)
|
swg := sizedwaitgroup.New(w.Options.Options.TemplateThreads)
|
||||||
for _, template := range w.Workflows {
|
for _, template := range w.Workflows {
|
||||||
swg.Add()
|
swg.Add()
|
||||||
func(template *WorkflowTemplate) {
|
func(template *workflows.WorkflowTemplate) {
|
||||||
if err := w.runWorkflowStep(template, input, results, &swg); err != nil {
|
if err := e.runWorkflowStep(template, input, results, &swg, w); err != nil {
|
||||||
gologger.Warning().Msgf("[%s] Could not execute workflow step: %s\n", template.Template, err)
|
gologger.Warning().Msgf("[%s] Could not execute workflow step: %s\n", template.Template, err)
|
||||||
}
|
}
|
||||||
swg.Done()
|
swg.Done()
|
||||||
|
@ -21,7 +28,7 @@ func (w *Workflow) RunWorkflow(input string) bool {
|
||||||
|
|
||||||
// runWorkflowStep runs a workflow step for the workflow. It executes the workflow
|
// runWorkflowStep runs a workflow step for the workflow. It executes the workflow
|
||||||
// in a recursive manner running all subtemplates and matchers.
|
// in a recursive manner running all subtemplates and matchers.
|
||||||
func (w *Workflow) runWorkflowStep(template *WorkflowTemplate, input string, results *atomic.Bool, swg *sizedwaitgroup.SizedWaitGroup) error {
|
func (e *Engine) runWorkflowStep(template *workflows.WorkflowTemplate, input string, results *atomic.Bool, swg *sizedwaitgroup.SizedWaitGroup, w *workflows.Workflow) error {
|
||||||
var firstMatched bool
|
var firstMatched bool
|
||||||
var err error
|
var err error
|
||||||
var mainErr error
|
var mainErr error
|
||||||
|
@ -84,8 +91,8 @@ func (w *Workflow) runWorkflowStep(template *WorkflowTemplate, input string, res
|
||||||
for _, subtemplate := range matcher.Subtemplates {
|
for _, subtemplate := range matcher.Subtemplates {
|
||||||
swg.Add()
|
swg.Add()
|
||||||
|
|
||||||
go func(subtemplate *WorkflowTemplate) {
|
go func(subtemplate *workflows.WorkflowTemplate) {
|
||||||
if err := w.runWorkflowStep(subtemplate, input, results, swg); err != nil {
|
if err := e.runWorkflowStep(subtemplate, input, results, swg, w); err != nil {
|
||||||
gologger.Warning().Msgf("[%s] Could not execute workflow step: %s\n", subtemplate.Template, err)
|
gologger.Warning().Msgf("[%s] Could not execute workflow step: %s\n", subtemplate.Template, err)
|
||||||
}
|
}
|
||||||
swg.Done()
|
swg.Done()
|
||||||
|
@ -108,8 +115,8 @@ func (w *Workflow) runWorkflowStep(template *WorkflowTemplate, input string, res
|
||||||
for _, subtemplate := range template.Subtemplates {
|
for _, subtemplate := range template.Subtemplates {
|
||||||
swg.Add()
|
swg.Add()
|
||||||
|
|
||||||
go func(template *WorkflowTemplate) {
|
go func(template *workflows.WorkflowTemplate) {
|
||||||
if err := w.runWorkflowStep(template, input, results, swg); err != nil {
|
if err := e.runWorkflowStep(template, input, results, swg, w); err != nil {
|
||||||
gologger.Warning().Msgf("[%s] Could not execute workflow step: %s\n", template.Template, err)
|
gologger.Warning().Msgf("[%s] Could not execute workflow step: %s\n", template.Template, err)
|
||||||
}
|
}
|
||||||
swg.Done()
|
swg.Done()
|
||||||
|
@ -117,4 +124,4 @@ func (w *Workflow) runWorkflowStep(template *WorkflowTemplate, input string, res
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
return mainErr
|
return mainErr
|
||||||
}*/
|
}
|
||||||
|
|
|
@ -1,6 +1,5 @@
|
||||||
package core
|
package core
|
||||||
|
|
||||||
/*
|
|
||||||
import (
|
import (
|
||||||
"testing"
|
"testing"
|
||||||
|
|
||||||
|
@ -17,13 +16,14 @@ import (
|
||||||
func TestWorkflowsSimple(t *testing.T) {
|
func TestWorkflowsSimple(t *testing.T) {
|
||||||
progressBar, _ := progress.NewStatsTicker(0, false, false, false, 0)
|
progressBar, _ := progress.NewStatsTicker(0, false, false, false, 0)
|
||||||
|
|
||||||
workflow := &workflows.Workflow{Options: &protocols.ExecuterOptions{Options: &types.Options{TemplateThreads: 10}}, Workflows: []*WorkflowTemplate{
|
workflow := &workflows.Workflow{Options: &protocols.ExecuterOptions{Options: &types.Options{TemplateThreads: 10}}, Workflows: []*workflows.WorkflowTemplate{
|
||||||
{Executers: []*workflows.ProtocolExecuterPair{{
|
{Executers: []*workflows.ProtocolExecuterPair{{
|
||||||
Executer: &mockExecuter{result: true}, Options: &protocols.ExecuterOptions{Progress: progressBar}},
|
Executer: &mockExecuter{result: true}, Options: &protocols.ExecuterOptions{Progress: progressBar}},
|
||||||
}},
|
}},
|
||||||
}}
|
}}
|
||||||
|
|
||||||
matched := workflow.RunWorkflow("https://test.com")
|
engine := &Engine{}
|
||||||
|
matched := engine.executeWorkflow("https://test.com", workflow)
|
||||||
require.True(t, matched, "could not get correct match value")
|
require.True(t, matched, "could not get correct match value")
|
||||||
}
|
}
|
||||||
|
|
||||||
|
@ -31,7 +31,7 @@ func TestWorkflowsSimpleMultiple(t *testing.T) {
|
||||||
progressBar, _ := progress.NewStatsTicker(0, false, false, false, 0)
|
progressBar, _ := progress.NewStatsTicker(0, false, false, false, 0)
|
||||||
|
|
||||||
var firstInput, secondInput string
|
var firstInput, secondInput string
|
||||||
workflow := &workflows.Workflow{Options: &protocols.ExecuterOptions{Options: &types.Options{TemplateThreads: 10}}, Workflows: []*WorkflowTemplate{
|
workflow := &workflows.Workflow{Options: &protocols.ExecuterOptions{Options: &types.Options{TemplateThreads: 10}}, Workflows: []*workflows.WorkflowTemplate{
|
||||||
{Executers: []*workflows.ProtocolExecuterPair{{
|
{Executers: []*workflows.ProtocolExecuterPair{{
|
||||||
Executer: &mockExecuter{result: true, executeHook: func(input string) {
|
Executer: &mockExecuter{result: true, executeHook: func(input string) {
|
||||||
firstInput = input
|
firstInput = input
|
||||||
|
@ -44,7 +44,8 @@ func TestWorkflowsSimpleMultiple(t *testing.T) {
|
||||||
}},
|
}},
|
||||||
}}
|
}}
|
||||||
|
|
||||||
matched := workflow.RunWorkflow("https://test.com")
|
engine := &Engine{}
|
||||||
|
matched := engine.executeWorkflow("https://test.com", workflow)
|
||||||
require.True(t, matched, "could not get correct match value")
|
require.True(t, matched, "could not get correct match value")
|
||||||
|
|
||||||
require.Equal(t, "https://test.com", firstInput, "could not get correct first input")
|
require.Equal(t, "https://test.com", firstInput, "could not get correct first input")
|
||||||
|
@ -55,7 +56,7 @@ func TestWorkflowsSubtemplates(t *testing.T) {
|
||||||
progressBar, _ := progress.NewStatsTicker(0, false, false, false, 0)
|
progressBar, _ := progress.NewStatsTicker(0, false, false, false, 0)
|
||||||
|
|
||||||
var firstInput, secondInput string
|
var firstInput, secondInput string
|
||||||
workflow := &workflows.Workflow{Options: &protocols.ExecuterOptions{Options: &types.Options{TemplateThreads: 10}}, Workflows: []*WorkflowTemplate{
|
workflow := &workflows.Workflow{Options: &protocols.ExecuterOptions{Options: &types.Options{TemplateThreads: 10}}, Workflows: []*workflows.WorkflowTemplate{
|
||||||
{Executers: []*workflows.ProtocolExecuterPair{{
|
{Executers: []*workflows.ProtocolExecuterPair{{
|
||||||
Executer: &mockExecuter{result: true, executeHook: func(input string) {
|
Executer: &mockExecuter{result: true, executeHook: func(input string) {
|
||||||
firstInput = input
|
firstInput = input
|
||||||
|
@ -69,7 +70,8 @@ func TestWorkflowsSubtemplates(t *testing.T) {
|
||||||
}}}},
|
}}}},
|
||||||
}}
|
}}
|
||||||
|
|
||||||
matched := workflow.RunWorkflow("https://test.com")
|
engine := &Engine{}
|
||||||
|
matched := engine.executeWorkflow("https://test.com", workflow)
|
||||||
require.True(t, matched, "could not get correct match value")
|
require.True(t, matched, "could not get correct match value")
|
||||||
|
|
||||||
require.Equal(t, "https://test.com", firstInput, "could not get correct first input")
|
require.Equal(t, "https://test.com", firstInput, "could not get correct first input")
|
||||||
|
@ -80,7 +82,7 @@ func TestWorkflowsSubtemplatesNoMatch(t *testing.T) {
|
||||||
progressBar, _ := progress.NewStatsTicker(0, false, false, false, 0)
|
progressBar, _ := progress.NewStatsTicker(0, false, false, false, 0)
|
||||||
|
|
||||||
var firstInput, secondInput string
|
var firstInput, secondInput string
|
||||||
workflow := &workflows.Workflow{Options: &protocols.ExecuterOptions{Options: &types.Options{TemplateThreads: 10}}, Workflows: []*WorkflowTemplate{
|
workflow := &workflows.Workflow{Options: &protocols.ExecuterOptions{Options: &types.Options{TemplateThreads: 10}}, Workflows: []*workflows.WorkflowTemplate{
|
||||||
{Executers: []*workflows.ProtocolExecuterPair{{
|
{Executers: []*workflows.ProtocolExecuterPair{{
|
||||||
Executer: &mockExecuter{result: false, executeHook: func(input string) {
|
Executer: &mockExecuter{result: false, executeHook: func(input string) {
|
||||||
firstInput = input
|
firstInput = input
|
||||||
|
@ -92,7 +94,8 @@ func TestWorkflowsSubtemplatesNoMatch(t *testing.T) {
|
||||||
}}}},
|
}}}},
|
||||||
}}
|
}}
|
||||||
|
|
||||||
matched := workflow.RunWorkflow("https://test.com")
|
engine := &Engine{}
|
||||||
|
matched := engine.executeWorkflow("https://test.com", workflow)
|
||||||
require.False(t, matched, "could not get correct match value")
|
require.False(t, matched, "could not get correct match value")
|
||||||
|
|
||||||
require.Equal(t, "https://test.com", firstInput, "could not get correct first input")
|
require.Equal(t, "https://test.com", firstInput, "could not get correct first input")
|
||||||
|
@ -103,7 +106,7 @@ func TestWorkflowsSubtemplatesWithMatcher(t *testing.T) {
|
||||||
progressBar, _ := progress.NewStatsTicker(0, false, false, false, 0)
|
progressBar, _ := progress.NewStatsTicker(0, false, false, false, 0)
|
||||||
|
|
||||||
var firstInput, secondInput string
|
var firstInput, secondInput string
|
||||||
workflow := &workflows.Workflow{Options: &protocols.ExecuterOptions{Options: &types.Options{TemplateThreads: 10}}, Workflows: []*WorkflowTemplate{
|
workflow := &workflows.Workflow{Options: &protocols.ExecuterOptions{Options: &types.Options{TemplateThreads: 10}}, Workflows: []*workflows.WorkflowTemplate{
|
||||||
{Executers: []*workflows.ProtocolExecuterPair{{
|
{Executers: []*workflows.ProtocolExecuterPair{{
|
||||||
Executer: &mockExecuter{result: true, executeHook: func(input string) {
|
Executer: &mockExecuter{result: true, executeHook: func(input string) {
|
||||||
firstInput = input
|
firstInput = input
|
||||||
|
@ -120,7 +123,8 @@ func TestWorkflowsSubtemplatesWithMatcher(t *testing.T) {
|
||||||
}}}}}},
|
}}}}}},
|
||||||
}}
|
}}
|
||||||
|
|
||||||
matched := workflow.RunWorkflow("https://test.com")
|
engine := &Engine{}
|
||||||
|
matched := engine.executeWorkflow("https://test.com", workflow)
|
||||||
require.True(t, matched, "could not get correct match value")
|
require.True(t, matched, "could not get correct match value")
|
||||||
|
|
||||||
require.Equal(t, "https://test.com", firstInput, "could not get correct first input")
|
require.Equal(t, "https://test.com", firstInput, "could not get correct first input")
|
||||||
|
@ -148,7 +152,8 @@ func TestWorkflowsSubtemplatesWithMatcherNoMatch(t *testing.T) {
|
||||||
}}}}}},
|
}}}}}},
|
||||||
}}
|
}}
|
||||||
|
|
||||||
matched := workflow.RunWorkflow("https://test.com")
|
engine := &Engine{}
|
||||||
|
matched := engine.executeWorkflow("https://test.com", workflow)
|
||||||
require.False(t, matched, "could not get correct match value")
|
require.False(t, matched, "could not get correct match value")
|
||||||
|
|
||||||
require.Equal(t, "https://test.com", firstInput, "could not get correct first input")
|
require.Equal(t, "https://test.com", firstInput, "could not get correct first input")
|
||||||
|
@ -189,4 +194,3 @@ func (m *mockExecuter) ExecuteWithResults(input string, callback protocols.Outpu
|
||||||
}
|
}
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
*/
|
|
||||||
|
|
|
@ -1,7 +1,7 @@
|
||||||
package core
|
package core
|
||||||
|
|
||||||
import (
|
import (
|
||||||
"github.com/projectdiscovery/nuclei/v2/pkg/templates"
|
"github.com/remeh/sizedwaitgroup"
|
||||||
)
|
)
|
||||||
|
|
||||||
// WorkPool implements an execution pool for executing different
|
// WorkPool implements an execution pool for executing different
|
||||||
|
@ -10,6 +10,8 @@ import (
|
||||||
// It also allows Configuration of such requirements. This is used
|
// It also allows Configuration of such requirements. This is used
|
||||||
// for per-module like separate headless concurrency etc.
|
// for per-module like separate headless concurrency etc.
|
||||||
type WorkPool struct {
|
type WorkPool struct {
|
||||||
|
Headless *sizedwaitgroup.SizedWaitGroup
|
||||||
|
Default *sizedwaitgroup.SizedWaitGroup
|
||||||
config WorkPoolConfig
|
config WorkPoolConfig
|
||||||
}
|
}
|
||||||
|
|
||||||
|
@ -27,9 +29,35 @@ type WorkPoolConfig struct {
|
||||||
|
|
||||||
// NewWorkPool returns a new WorkPool instance
|
// NewWorkPool returns a new WorkPool instance
|
||||||
func NewWorkPool(config WorkPoolConfig) *WorkPool {
|
func NewWorkPool(config WorkPoolConfig) *WorkPool {
|
||||||
return &WorkPool{config: config}
|
headlessWg := sizedwaitgroup.New(config.HeadlessTypeConcurrency)
|
||||||
|
defaultWg := sizedwaitgroup.New(config.TypeConcurrency)
|
||||||
|
|
||||||
|
return &WorkPool{
|
||||||
|
config: config,
|
||||||
|
Headless: &headlessWg,
|
||||||
|
Default: &defaultWg,
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func (w *WorkPool) Execute(templates []*templates.Template) {
|
// Wait waits for all the workpool waitgroups to finish
|
||||||
|
func (w *WorkPool) Wait() {
|
||||||
|
w.Default.Wait()
|
||||||
|
w.Headless.Wait()
|
||||||
|
}
|
||||||
|
|
||||||
|
// InputWorkPool is a workpool per-input
|
||||||
|
type InputWorkPool struct {
|
||||||
|
Waitgroup *sizedwaitgroup.SizedWaitGroup
|
||||||
|
}
|
||||||
|
|
||||||
|
// InputPool returns a workpool for an input type
|
||||||
|
func (w *WorkPool) InputPool(templateType string) *InputWorkPool {
|
||||||
|
var count int
|
||||||
|
if templateType == "headless" {
|
||||||
|
count = w.config.HeadlessInputConcurrency
|
||||||
|
} else {
|
||||||
|
count = w.config.InputConcurrency
|
||||||
|
}
|
||||||
|
swg := sizedwaitgroup.New(count)
|
||||||
|
return &InputWorkPool{Waitgroup: &swg}
|
||||||
}
|
}
|
||||||
|
|
|
@ -3,6 +3,7 @@ package protocols
|
||||||
import (
|
import (
|
||||||
"go.uber.org/ratelimit"
|
"go.uber.org/ratelimit"
|
||||||
|
|
||||||
|
"github.com/logrusorgru/aurora"
|
||||||
"github.com/projectdiscovery/nuclei/v2/pkg/catalog"
|
"github.com/projectdiscovery/nuclei/v2/pkg/catalog"
|
||||||
"github.com/projectdiscovery/nuclei/v2/pkg/model"
|
"github.com/projectdiscovery/nuclei/v2/pkg/model"
|
||||||
"github.com/projectdiscovery/nuclei/v2/pkg/operators"
|
"github.com/projectdiscovery/nuclei/v2/pkg/operators"
|
||||||
|
@ -61,6 +62,7 @@ type ExecuterOptions struct {
|
||||||
|
|
||||||
Operators []*operators.Operators // only used by offlinehttp module
|
Operators []*operators.Operators // only used by offlinehttp module
|
||||||
|
|
||||||
|
Colorizer aurora.Aurora
|
||||||
WorkflowLoader model.WorkflowLoader
|
WorkflowLoader model.WorkflowLoader
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
Loading…
Reference in New Issue