Skip to content

Commit

Permalink
change executor funcs' signature
Browse files Browse the repository at this point in the history
  • Loading branch information
@dogancanbakir
dogancanbakir committed Nov 16, 2023
1 parent 722da28 commit ead7ff9
Show file tree
Hide file tree
Showing 5 changed files with 110 additions and 101 deletions.
31 changes: 18 additions & 13 deletions pkg/core/executors.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ import (

"github.com/projectdiscovery/gologger"
"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/contextargs"
"github.com/projectdiscovery/nuclei/v3/pkg/scan"
"github.com/projectdiscovery/nuclei/v3/pkg/templates"
"github.com/projectdiscovery/nuclei/v3/pkg/templates/types"
generalTypes "github.com/projectdiscovery/nuclei/v3/pkg/types"
Expand All @@ -22,16 +23,17 @@ func (e *Engine) executeAllSelfContained(alltemplates []*templates.Template, res
defer sg.Done()
var err error
var match bool
ctx := scan.NewScanContext(contextargs.New())
if e.Callback != nil {
if results, err := template.Executer.ExecuteWithResults(contextargs.New()); err != nil {
if results, err := template.Executer.ExecuteWithResults(ctx); err != nil {
for _, result := range results {
e.Callback(result)
}
}

match = true
} else {
match, err = template.Executer.Execute(contextargs.New())
match, err = template.Executer.Execute(ctx)
}
if err != nil {
gologger.Warning().Msgf("[%s] Could not execute step: %s\n", e.executerOpts.Colorizer.BrightBlue(template.ID), err)
Expand Down Expand Up @@ -111,21 +113,22 @@ func (e *Engine) executeTemplateWithTargets(template *templates.Template, target

var match bool
var err error
ctxArgs := contextargs.New()
ctxArgs.MetaInput = value
ctx := scan.NewScanContext(ctxArgs)
switch template.Type() {
case types.WorkflowProtocol:
match = e.executeWorkflow(value, template.CompiledWorkflow)
match = e.executeWorkflow(ctx, template.CompiledWorkflow)
default:
ctxArgs := contextargs.New()
ctxArgs.MetaInput = value
if e.Callback != nil {
if results, err := template.Executer.ExecuteWithResults(ctxArgs); err != nil {
if results, err := template.Executer.ExecuteWithResults(ctx); err != nil {
for _, result := range results {
e.Callback(result)
}
}
match = true
} else {
match, err = template.Executer.Execute(ctxArgs)
match, err = template.Executer.Execute(ctx)
}
}
if err != nil {
Expand Down Expand Up @@ -166,21 +169,22 @@ func (e *Engine) executeTemplatesOnTarget(alltemplates []*templates.Template, ta

var match bool
var err error
ctxArgs := contextargs.New()
ctxArgs.MetaInput = value
ctx := scan.NewScanContext(ctxArgs)
switch template.Type() {
case types.WorkflowProtocol:
match = e.executeWorkflow(value, template.CompiledWorkflow)
match = e.executeWorkflow(ctx, template.CompiledWorkflow)
default:
ctxArgs := contextargs.New()
ctxArgs.MetaInput = value
if e.Callback != nil {
if results, err := template.Executer.ExecuteWithResults(ctxArgs); err != nil {
if results, err := template.Executer.ExecuteWithResults(ctx); err != nil {
for _, result := range results {
e.Callback(result)
}
}
match = true
} else {
match, err = template.Executer.Execute(ctxArgs)
match, err = template.Executer.Execute(ctx)
}
}
if err != nil {
Expand Down Expand Up @@ -221,7 +225,8 @@ func (e *ChildExecuter) Execute(template *templates.Template, value *contextargs

ctxArgs := contextargs.New()
ctxArgs.MetaInput = value
match, err := template.Executer.Execute(ctxArgs)
ctx := scan.NewScanContext(ctxArgs)
match, err := template.Executer.Execute(ctx)
if err != nil {
gologger.Warning().Msgf("[%s] Could not execute step: %s\n", e.e.executerOpts.Colorizer.BrightBlue(template.ID), err)
}
Expand Down
131 changes: 68 additions & 63 deletions pkg/core/workflow_execute.go
View file
Original file line number Diff line number Diff line change
@@ -1,26 +1,29 @@
package core

import (
"fmt"
"net/http/cookiejar"
"sync/atomic"

"github.com/remeh/sizedwaitgroup"

"github.com/projectdiscovery/gologger"
"github.com/projectdiscovery/nuclei/v3/pkg/output"
"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/contextargs"
"github.com/projectdiscovery/nuclei/v3/pkg/scan"
"github.com/projectdiscovery/nuclei/v3/pkg/workflows"
)

const workflowStepExecutionError = "[%s] Could not execute workflow step: %s\n"

// executeWorkflow runs a workflow on an input and returns true or false
func (e *Engine) executeWorkflow(input *contextargs.MetaInput, w *workflows.Workflow) bool {
func (e *Engine) executeWorkflow(ctx *scan.ScanContext, w *workflows.Workflow) bool {
results := &atomic.Bool{}

// at this point we should be at the start root execution of a workflow tree, hence we create global shared instances
workflowCookieJar, _ := cookiejar.New(nil)
ctxArgs := contextargs.New()
ctxArgs.MetaInput = input
ctxArgs.MetaInput = ctx.Input.MetaInput
ctxArgs.CookieJar = workflowCookieJar

// we can know the nesting level only at runtime, so the best we can do here is increase template threads by one unit in case it's equal to 1 to allow
Expand All @@ -37,7 +40,7 @@ func (e *Engine) executeWorkflow(input *contextargs.MetaInput, w *workflows.Work
func(template *workflows.WorkflowTemplate) {
defer swg.Done()

if err := e.runWorkflowStep(template, ctxArgs, results, &swg, w); err != nil {
if err := e.runWorkflowStep(template, ctx, results, &swg, w); err != nil {
gologger.Warning().Msgf(workflowStepExecutionError, template.Template, err)
}
}(template)
Expand All @@ -48,7 +51,7 @@ func (e *Engine) executeWorkflow(input *contextargs.MetaInput, w *workflows.Work

// runWorkflowStep runs a workflow step for the workflow. It executes the workflow
// in a recursive manner running all subtemplates and matchers.
func (e *Engine) runWorkflowStep(template *workflows.WorkflowTemplate, input *contextargs.Context, results *atomic.Bool, swg *sizedwaitgroup.SizedWaitGroup, w *workflows.Workflow) error {
func (e *Engine) runWorkflowStep(template *workflows.WorkflowTemplate, ctx *scan.ScanContext, results *atomic.Bool, swg *sizedwaitgroup.SizedWaitGroup, w *workflows.Workflow) error {
var firstMatched bool
var err error
var mainErr error
Expand All @@ -59,43 +62,44 @@ func (e *Engine) runWorkflowStep(template *workflows.WorkflowTemplate, input *co

// Don't print results with subtemplates, only print results on template.
if len(template.Subtemplates) > 0 {
// err = executer.Executer.ExecuteWithResults(input, func(result *output.InternalWrappedEvent) {
// if result.OperatorsResult == nil {
// return
// }
// if len(result.Results) > 0 {
// firstMatched = true
// }

// if result.OperatorsResult != nil && result.OperatorsResult.Extracts != nil {
// for k, v := range result.OperatorsResult.Extracts {
// // normalize items:
// switch len(v) {
// case 0, 1:
// // - key:[item] => key: item
// input.Set(k, v[0])
// default:
// // - key:[item_0, ..., item_n] => key0:item_0, keyn:item_n
// for vIdx, vVal := range v {
// normalizedKIdx := fmt.Sprintf("%s%d", k, vIdx)
// input.Set(normalizedKIdx, vVal)
// }
// // also add the original name with full slice
// input.Set(k, v)
// }
// }
// }
// })
ctx.OnResult = func(result *output.InternalWrappedEvent) {
if result.OperatorsResult == nil {
return
}
if len(result.Results) > 0 {
firstMatched = true
}

if result.OperatorsResult != nil && result.OperatorsResult.Extracts != nil {
for k, v := range result.OperatorsResult.Extracts {
// normalize items:
switch len(v) {
case 0, 1:
// - key:[item] => key: item
ctx.Input.Set(k, v[0])
default:
// - key:[item_0, ..., item_n] => key0:item_0, keyn:item_n
for vIdx, vVal := range v {
normalizedKIdx := fmt.Sprintf("%s%d", k, vIdx)
ctx.Input.Set(normalizedKIdx, vVal)
}
// also add the original name with full slice
ctx.Input.Set(k, v)
}
}
}
}
_, err = executer.Executer.ExecuteWithResults(ctx)
} else {
var matched bool
matched, err = executer.Executer.Execute(input)
matched, err = executer.Executer.Execute(ctx)
if matched {
firstMatched = true
}
}
if err != nil {
if w.Options.HostErrorsCache != nil {
w.Options.HostErrorsCache.MarkFailed(input.MetaInput.ID(), err)
w.Options.HostErrorsCache.MarkFailed(ctx.Input.MetaInput.ID(), err)
}
if len(template.Executers) == 1 {
mainErr = err
Expand All @@ -113,35 +117,36 @@ func (e *Engine) runWorkflowStep(template *workflows.WorkflowTemplate, input *co
for _, executer := range template.Executers {
executer.Options.Progress.AddToTotal(int64(executer.Executer.Requests()))

// err := executer.Executer.ExecuteWithResults(input, func(event *output.InternalWrappedEvent) {
// if event.OperatorsResult == nil {
// return
// }

// if event.OperatorsResult.Extracts != nil {
// for k, v := range event.OperatorsResult.Extracts {
// input.Set(k, v)
// }
// }

// for _, matcher := range template.Matchers {
// if !matcher.Match(event.OperatorsResult) {
// continue
// }

// for _, subtemplate := range matcher.Subtemplates {
// swg.Add()

// go func(subtemplate *workflows.WorkflowTemplate) {
// defer swg.Done()

// if err := e.runWorkflowStep(subtemplate, input, results, swg, w); err != nil {
// gologger.Warning().Msgf(workflowStepExecutionError, subtemplate.Template, err)
// }
// }(subtemplate)
// }
// }
// })
ctx.OnResult = func(event *output.InternalWrappedEvent) {
if event.OperatorsResult == nil {
return
}

if event.OperatorsResult.Extracts != nil {
for k, v := range event.OperatorsResult.Extracts {
ctx.Input.Set(k, v)
}
}

for _, matcher := range template.Matchers {
if !matcher.Match(event.OperatorsResult) {
continue
}

for _, subtemplate := range matcher.Subtemplates {
swg.Add()

go func(subtemplate *workflows.WorkflowTemplate) {
defer swg.Done()

if err := e.runWorkflowStep(subtemplate, ctx, results, swg, w); err != nil {
gologger.Warning().Msgf(workflowStepExecutionError, subtemplate.Template, err)
}
}(subtemplate)
}
}
}
_, err := executer.Executer.ExecuteWithResults(ctx)
if err != nil {
if len(template.Executers) == 1 {
mainErr = err
Expand All @@ -158,7 +163,7 @@ func (e *Engine) runWorkflowStep(template *workflows.WorkflowTemplate, input *co
swg.Add()

go func(template *workflows.WorkflowTemplate) {
if err := e.runWorkflowStep(template, input, results, swg, w); err != nil {
if err := e.runWorkflowStep(template, ctx, results, swg, w); err != nil {
gologger.Warning().Msgf(workflowStepExecutionError, template.Template, err)
}
swg.Done()
Expand Down
5 changes: 3 additions & 2 deletions pkg/protocols/protocols.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@ import (
"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/variables"
"github.com/projectdiscovery/nuclei/v3/pkg/protocols/headless/engine"
"github.com/projectdiscovery/nuclei/v3/pkg/reporting"
"github.com/projectdiscovery/nuclei/v3/pkg/scan"
templateTypes "github.com/projectdiscovery/nuclei/v3/pkg/templates/types"
"github.com/projectdiscovery/nuclei/v3/pkg/types"
)
Expand All @@ -40,9 +41,9 @@ type Executer interface {
// Requests returns the total number of requests the rule will perform
Requests() int
// Execute executes the protocol group and returns true or false if results were found.
Execute(input *contextargs.Context) (bool, error)
Execute(ctx *scan.ScanContext) (bool, error)
// ExecuteWithResults executes the protocol requests and returns results instead of writing them.
ExecuteWithResults(input *contextargs.Context) ([]*output.ResultEvent, error)
ExecuteWithResults(ctx *scan.ScanContext) ([]*output.ResultEvent, error)
}

// ExecutorOptions contains the configuration options for executer clients
Expand Down
23 changes: 11 additions & 12 deletions pkg/templates/cluster.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,6 @@ import (
"github.com/projectdiscovery/nuclei/v3/pkg/operators"
"github.com/projectdiscovery/nuclei/v3/pkg/output"
"github.com/projectdiscovery/nuclei/v3/pkg/protocols"
"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/contextargs"
"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/helpers/writer"
"github.com/projectdiscovery/nuclei/v3/pkg/scan"
"github.com/projectdiscovery/nuclei/v3/pkg/templates/types"
Expand Down Expand Up @@ -241,12 +240,12 @@ func (e *ClusterExecuter) Requests() int {
}

// Execute executes the protocol group and returns true or false if results were found.
func (e *ClusterExecuter) Execute(input *contextargs.Context) (bool, error) {
func (e *ClusterExecuter) Execute(ctx *scan.ScanContext) (bool, error) {
var results bool

inputItem := input.Clone()
if e.options.InputHelper != nil && input.MetaInput.Input != "" {
if inputItem.MetaInput.Input = e.options.InputHelper.Transform(input.MetaInput.Input, e.templateType); input.MetaInput.Input == "" {
inputItem := ctx.Input.Clone()
if e.options.InputHelper != nil && ctx.Input.MetaInput.Input != "" {
if inputItem.MetaInput.Input = e.options.InputHelper.Transform(ctx.Input.MetaInput.Input, e.templateType); ctx.Input.MetaInput.Input == "" {
return false, nil
}
}
Expand Down Expand Up @@ -275,19 +274,19 @@ func (e *ClusterExecuter) Execute(input *contextargs.Context) (bool, error) {
}
})
if err != nil && e.options.HostErrorsCache != nil {
e.options.HostErrorsCache.MarkFailed(input.MetaInput.Input, err)
e.options.HostErrorsCache.MarkFailed(ctx.Input.MetaInput.Input, err)
}
return results, err
}

// ExecuteWithResults executes the protocol requests and returns results instead of writing them.
func (e *ClusterExecuter) ExecuteWithResults(input *contextargs.Context) ([]*output.ResultEvent, error) {
scanCtx := scan.NewScanContext(input)
func (e *ClusterExecuter) ExecuteWithResults(ctx *scan.ScanContext) ([]*output.ResultEvent, error) {
scanCtx := scan.NewScanContext(ctx.Input)
dynamicValues := make(map[string]interface{})

inputItem := input.Clone()
if e.options.InputHelper != nil && input.MetaInput.Input != "" {
if inputItem.MetaInput.Input = e.options.InputHelper.Transform(input.MetaInput.Input, e.templateType); input.MetaInput.Input == "" {
inputItem := ctx.Input.Clone()
if e.options.InputHelper != nil && ctx.Input.MetaInput.Input != "" {
if inputItem.MetaInput.Input = e.options.InputHelper.Transform(ctx.Input.MetaInput.Input, e.templateType); ctx.Input.MetaInput.Input == "" {
return nil, nil
}
}
Expand All @@ -305,7 +304,7 @@ func (e *ClusterExecuter) ExecuteWithResults(input *contextargs.Context) ([]*out
}
})
if err != nil && e.options.HostErrorsCache != nil {
e.options.HostErrorsCache.MarkFailed(input.MetaInput.Input, err)
e.options.HostErrorsCache.MarkFailed(ctx.Input.MetaInput.Input, err)
}
return scanCtx.GenerateResult(), err
}
Loading

0 comments on commit ead7ff9

Please sign in to comment.