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/pkg/errors"
|
||||
"github.com/remeh/sizedwaitgroup"
|
||||
"go.uber.org/atomic"
|
||||
"go.uber.org/ratelimit"
|
||||
"gopkg.in/yaml.v2"
|
||||
|
||||
|
@ -20,7 +18,7 @@ import (
|
|||
"github.com/projectdiscovery/nuclei/v2/pkg/catalog/config"
|
||||
"github.com/projectdiscovery/nuclei/v2/pkg/catalog/loader"
|
||||
"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/output"
|
||||
"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/exporters/markdown"
|
||||
"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/utils"
|
||||
"github.com/projectdiscovery/nuclei/v2/pkg/utils/stats"
|
||||
|
@ -52,7 +49,7 @@ type Runner struct {
|
|||
colorizer aurora.Aurora
|
||||
issuesClient *reporting.Client
|
||||
addColor func(severity.Severity) string
|
||||
hmapInputProvider *hmap.Input
|
||||
hmapInputProvider *hybrid.Input
|
||||
browser *engine.Browser
|
||||
ratelimiter ratelimit.Limiter
|
||||
hostErrors *hosterrorscache.Cache
|
||||
|
@ -112,7 +109,7 @@ func New(options *types.Options) (*Runner, error) {
|
|||
}
|
||||
|
||||
// Initialize the input source
|
||||
hmapInput, err := hmap.New(options)
|
||||
hmapInput, err := hybrid.New(options)
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "could not create input provider")
|
||||
}
|
||||
|
@ -259,6 +256,7 @@ func (r *Runner) RunEnumeration() error {
|
|||
ProjectFile: r.projectFile,
|
||||
Browser: r.browser,
|
||||
HostErrorsCache: cache,
|
||||
Colorizer: r.colorizer,
|
||||
}
|
||||
engine := core.New(r.options)
|
||||
engine.SetExecuterOptions(executerOpts)
|
||||
|
@ -351,9 +349,6 @@ func (r *Runner) RunEnumeration() error {
|
|||
gologger.Info().Msgf("Workflows loaded for scan: %d", len(store.Workflows()))
|
||||
}
|
||||
|
||||
// pre-parse all the templates, apply filters
|
||||
finalTemplates := []*templates.Template{}
|
||||
|
||||
var unclusteredRequests int64
|
||||
for _, template := range store.Templates() {
|
||||
// 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
|
||||
for _, t := range finalTemplates {
|
||||
if len(t.Workflows) > 0 {
|
||||
|
@ -391,32 +392,10 @@ func (r *Runner) RunEnumeration() error {
|
|||
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
|
||||
r.progress.Init(r.hmapInputProvider.Count(), templateCount, totalRequests)
|
||||
|
||||
for _, t := range finalTemplates {
|
||||
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()
|
||||
results := engine.ExecuteWithOpts(finalTemplates, r.hmapInputProvider, true)
|
||||
|
||||
if r.interactsh != nil {
|
||||
matched := r.interactsh.Close()
|
||||
|
|
|
@ -24,6 +24,10 @@ type Engine struct {
|
|||
//
|
||||
// An example InputProvider is provided in form of hmap input provider.
|
||||
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
|
||||
|
|
|
@ -3,9 +3,12 @@ package core
|
|||
import (
|
||||
"fmt"
|
||||
|
||||
"github.com/projectdiscovery/gologger"
|
||||
"github.com/projectdiscovery/nuclei/v2/pkg/protocols/common/clusterer"
|
||||
"github.com/projectdiscovery/nuclei/v2/pkg/templates"
|
||||
"github.com/remeh/sizedwaitgroup"
|
||||
"github.com/rs/xid"
|
||||
"go.uber.org/atomic"
|
||||
)
|
||||
|
||||
// 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
|
||||
// of the engine.
|
||||
func (e *Engine) Execute(templates []*templates.Template) {
|
||||
finalTemplates, clusterCount := e.clusterTemplates(templates)
|
||||
func (e *Engine) Execute(templates []*templates.Template, input InputProvider) *atomic.Bool {
|
||||
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 {
|
||||
templateType := template.Type()
|
||||
|
||||
var wg *sizedwaitgroup.SizedWaitGroup
|
||||
if templateType == "headless" {
|
||||
wg = e.workPool.Headless
|
||||
} else {
|
||||
wg = e.workPool.Default
|
||||
}
|
||||
|
||||
wg.Add()
|
||||
switch {
|
||||
case template.SelfContained:
|
||||
// Self Contained requests are executed here
|
||||
|
||||
case templateType == "workflow":
|
||||
// Workflows requests are executed here
|
||||
|
||||
// Self Contained requests are executed here separately
|
||||
e.executeSelfContainedTemplateWithInput(template, results)
|
||||
default:
|
||||
// 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
|
||||
func (e *Engine) clusterTemplates(templatesList []*templates.Template) ([]*templates.Template, int) {
|
||||
// processSelfContainedTemplates execute a self-contained template.
|
||||
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 {
|
||||
return templatesList, 0
|
||||
}
|
||||
|
@ -65,85 +127,3 @@ func (e *Engine) clusterTemplates(templatesList []*templates.Template) ([]*templ
|
|||
}
|
||||
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.
|
||||
package hmap
|
||||
package hybrid
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
|
@ -118,3 +118,16 @@ func (i *Input) normalizeStoreInputValue(value string) {
|
|||
func (i *Input) Count() int64 {
|
||||
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
|
||||
|
||||
/*
|
||||
// RunWorkflow runs a workflow on an input and returns true or false
|
||||
func (w *Workflow) RunWorkflow(input string) bool {
|
||||
import (
|
||||
"github.com/projectdiscovery/gologger"
|
||||
"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{}
|
||||
|
||||
swg := sizedwaitgroup.New(w.Options.Options.TemplateThreads)
|
||||
for _, template := range w.Workflows {
|
||||
swg.Add()
|
||||
func(template *WorkflowTemplate) {
|
||||
if err := w.runWorkflowStep(template, input, results, &swg); err != nil {
|
||||
func(template *workflows.WorkflowTemplate) {
|
||||
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)
|
||||
}
|
||||
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
|
||||
// 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 err error
|
||||
var mainErr error
|
||||
|
@ -84,8 +91,8 @@ func (w *Workflow) runWorkflowStep(template *WorkflowTemplate, input string, res
|
|||
for _, subtemplate := range matcher.Subtemplates {
|
||||
swg.Add()
|
||||
|
||||
go func(subtemplate *WorkflowTemplate) {
|
||||
if err := w.runWorkflowStep(subtemplate, input, results, swg); err != nil {
|
||||
go func(subtemplate *workflows.WorkflowTemplate) {
|
||||
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)
|
||||
}
|
||||
swg.Done()
|
||||
|
@ -108,8 +115,8 @@ func (w *Workflow) runWorkflowStep(template *WorkflowTemplate, input string, res
|
|||
for _, subtemplate := range template.Subtemplates {
|
||||
swg.Add()
|
||||
|
||||
go func(template *WorkflowTemplate) {
|
||||
if err := w.runWorkflowStep(template, input, results, swg); err != nil {
|
||||
go func(template *workflows.WorkflowTemplate) {
|
||||
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)
|
||||
}
|
||||
swg.Done()
|
||||
|
@ -117,4 +124,4 @@ func (w *Workflow) runWorkflowStep(template *WorkflowTemplate, input string, res
|
|||
}
|
||||
}
|
||||
return mainErr
|
||||
}*/
|
||||
}
|
||||
|
|
|
@ -1,6 +1,5 @@
|
|||
package core
|
||||
|
||||
/*
|
||||
import (
|
||||
"testing"
|
||||
|
||||
|
@ -17,13 +16,14 @@ import (
|
|||
func TestWorkflowsSimple(t *testing.T) {
|
||||
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{{
|
||||
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")
|
||||
}
|
||||
|
||||
|
@ -31,7 +31,7 @@ func TestWorkflowsSimpleMultiple(t *testing.T) {
|
|||
progressBar, _ := progress.NewStatsTicker(0, false, false, false, 0)
|
||||
|
||||
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{{
|
||||
Executer: &mockExecuter{result: true, executeHook: func(input string) {
|
||||
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.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)
|
||||
|
||||
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{{
|
||||
Executer: &mockExecuter{result: true, executeHook: func(input string) {
|
||||
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.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)
|
||||
|
||||
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{{
|
||||
Executer: &mockExecuter{result: false, executeHook: func(input string) {
|
||||
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.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)
|
||||
|
||||
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{{
|
||||
Executer: &mockExecuter{result: true, executeHook: func(input string) {
|
||||
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.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.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
|
||||
}
|
||||
*/
|
||||
|
|
|
@ -1,7 +1,7 @@
|
|||
package core
|
||||
|
||||
import (
|
||||
"github.com/projectdiscovery/nuclei/v2/pkg/templates"
|
||||
"github.com/remeh/sizedwaitgroup"
|
||||
)
|
||||
|
||||
// WorkPool implements an execution pool for executing different
|
||||
|
@ -10,6 +10,8 @@ import (
|
|||
// It also allows Configuration of such requirements. This is used
|
||||
// for per-module like separate headless concurrency etc.
|
||||
type WorkPool struct {
|
||||
Headless *sizedwaitgroup.SizedWaitGroup
|
||||
Default *sizedwaitgroup.SizedWaitGroup
|
||||
config WorkPoolConfig
|
||||
}
|
||||
|
||||
|
@ -27,9 +29,35 @@ type WorkPoolConfig struct {
|
|||
|
||||
// NewWorkPool returns a new WorkPool instance
|
||||
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 (
|
||||
"go.uber.org/ratelimit"
|
||||
|
||||
"github.com/logrusorgru/aurora"
|
||||
"github.com/projectdiscovery/nuclei/v2/pkg/catalog"
|
||||
"github.com/projectdiscovery/nuclei/v2/pkg/model"
|
||||
"github.com/projectdiscovery/nuclei/v2/pkg/operators"
|
||||
|
@ -61,6 +62,7 @@ type ExecuterOptions struct {
|
|||
|
||||
Operators []*operators.Operators // only used by offlinehttp module
|
||||
|
||||
Colorizer aurora.Aurora
|
||||
WorkflowLoader model.WorkflowLoader
|
||||
}
|
||||
|
||||
|
|
Loading…
Reference in New Issue