Skip to content

Commit

Permalink
introduce scan context (#4373)
Browse files Browse the repository at this point in the history 
* introduce scan context

* minor

* add joined errors to resultevents

* change `executor` funcs' signature

* fix tests

* join errors in `LogError` func

* change func signature

* add guard
  • Loading branch information
@dogancanbakir
dogancanbakir authored Nov 27, 2023
1 parent 469598c commit ce5df9c
Show file tree
Hide file tree
Showing 27 changed files with 257 additions and 130 deletions.
45 changes: 25 additions & 20 deletions pkg/core/executors.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,8 @@ import (
"sync/atomic"

"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/templates"
"github.com/projectdiscovery/nuclei/v3/pkg/templates/types"
generalTypes "github.com/projectdiscovery/nuclei/v3/pkg/types"
Expand All @@ -23,15 +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 {
err = template.Executer.ExecuteWithResults(contextargs.New(), func(event *output.InternalWrappedEvent) {
for _, result := range event.Results {
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 {
err = template.Executer.ExecuteWithResults(ctxArgs, func(event *output.InternalWrappedEvent) {
for _, result := range event.Results {
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 {
err = template.Executer.ExecuteWithResults(ctxArgs, func(event *output.InternalWrappedEvent) {
for _, result := range event.Results {
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
35 changes: 19 additions & 16 deletions pkg/core/workflow_execute.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,19 +10,20 @@ import (
"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 @@ -39,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 @@ -50,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 @@ -61,7 +62,7 @@ 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) {
ctx.OnResult = func(result *output.InternalWrappedEvent) {
if result.OperatorsResult == nil {
return
}
Expand All @@ -75,29 +76,30 @@ func (e *Engine) runWorkflowStep(template *workflows.WorkflowTemplate, input *co
switch len(v) {
case 0, 1:
// - key:[item] => key: item
input.Set(k, v[0])
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)
input.Set(normalizedKIdx, vVal)
ctx.Input.Set(normalizedKIdx, vVal)
}
// also add the original name with full slice
input.Set(k, v)
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 @@ -115,14 +117,14 @@ 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) {
ctx.OnResult = 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)
ctx.Input.Set(k, v)
}
}

Expand All @@ -137,13 +139,14 @@ func (e *Engine) runWorkflowStep(template *workflows.WorkflowTemplate, input *co
go func(subtemplate *workflows.WorkflowTemplate) {
defer swg.Done()

if err := e.runWorkflowStep(subtemplate, input, results, swg, w); err != nil {
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 @@ -160,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
37 changes: 25 additions & 12 deletions pkg/core/workflow_execute_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ import (
"github.com/projectdiscovery/nuclei/v3/pkg/progress"
"github.com/projectdiscovery/nuclei/v3/pkg/protocols"
"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/contextargs"
"github.com/projectdiscovery/nuclei/v3/pkg/scan"
"github.com/projectdiscovery/nuclei/v3/pkg/types"
"github.com/projectdiscovery/nuclei/v3/pkg/workflows"
"github.com/stretchr/testify/require"
Expand All @@ -24,7 +25,9 @@ func TestWorkflowsSimple(t *testing.T) {
}}

engine := &Engine{}
matched := engine.executeWorkflow(&contextargs.MetaInput{Input: "https://test.com"}, workflow)
input := contextargs.NewWithInput("https://test.com")
ctx := scan.NewScanContext(input)
matched := engine.executeWorkflow(ctx, workflow)
require.True(t, matched, "could not get correct match value")
}

Expand All @@ -46,7 +49,9 @@ func TestWorkflowsSimpleMultiple(t *testing.T) {
}}

engine := &Engine{}
matched := engine.executeWorkflow(&contextargs.MetaInput{Input: "https://test.com"}, workflow)
input := contextargs.NewWithInput("https://test.com")
ctx := scan.NewScanContext(input)
matched := engine.executeWorkflow(ctx, workflow)
require.True(t, matched, "could not get correct match value")

require.Equal(t, "https://test.com", firstInput, "could not get correct first input")
Expand All @@ -72,7 +77,9 @@ func TestWorkflowsSubtemplates(t *testing.T) {
}}

engine := &Engine{}
matched := engine.executeWorkflow(&contextargs.MetaInput{Input: "https://test.com"}, workflow)
input := contextargs.NewWithInput("https://test.com")
ctx := scan.NewScanContext(input)
matched := engine.executeWorkflow(ctx, workflow)
require.True(t, matched, "could not get correct match value")

require.Equal(t, "https://test.com", firstInput, "could not get correct first input")
Expand All @@ -96,7 +103,9 @@ func TestWorkflowsSubtemplatesNoMatch(t *testing.T) {
}}

engine := &Engine{}
matched := engine.executeWorkflow(&contextargs.MetaInput{Input: "https://test.com"}, workflow)
input := contextargs.NewWithInput("https://test.com")
ctx := scan.NewScanContext(input)
matched := engine.executeWorkflow(ctx, workflow)
require.False(t, matched, "could not get correct match value")

require.Equal(t, "https://test.com", firstInput, "could not get correct first input")
Expand Down Expand Up @@ -125,7 +134,9 @@ func TestWorkflowsSubtemplatesWithMatcher(t *testing.T) {
}}

engine := &Engine{}
matched := engine.executeWorkflow(&contextargs.MetaInput{Input: "https://test.com"}, workflow)
input := contextargs.NewWithInput("https://test.com")
ctx := scan.NewScanContext(input)
matched := engine.executeWorkflow(ctx, workflow)
require.True(t, matched, "could not get correct match value")

require.Equal(t, "https://test.com", firstInput, "could not get correct first input")
Expand Down Expand Up @@ -154,7 +165,9 @@ func TestWorkflowsSubtemplatesWithMatcherNoMatch(t *testing.T) {
}}

engine := &Engine{}
matched := engine.executeWorkflow(&contextargs.MetaInput{Input: "https://test.com"}, workflow)
input := contextargs.NewWithInput("https://test.com")
ctx := scan.NewScanContext(input)
matched := engine.executeWorkflow(ctx, workflow)
require.False(t, matched, "could not get correct match value")

require.Equal(t, "https://test.com", firstInput, "could not get correct first input")
Expand All @@ -178,20 +191,20 @@ func (m *mockExecuter) Requests() int {
}

// Execute executes the protocol group and returns true or false if results were found.
func (m *mockExecuter) Execute(input *contextargs.Context) (bool, error) {
func (m *mockExecuter) Execute(ctx *scan.ScanContext) (bool, error) {
if m.executeHook != nil {
m.executeHook(input.MetaInput)
m.executeHook(ctx.Input.MetaInput)
}
return m.result, nil
}

// ExecuteWithResults executes the protocol requests and returns results instead of writing them.
func (m *mockExecuter) ExecuteWithResults(input *contextargs.Context, callback protocols.OutputEventCallback) error {
func (m *mockExecuter) ExecuteWithResults(ctx *scan.ScanContext) ([]*output.ResultEvent, error) {
if m.executeHook != nil {
m.executeHook(input.MetaInput)
m.executeHook(ctx.Input.MetaInput)
}
for _, output := range m.outputs {
callback(output)
ctx.LogEvent(output)
}
return nil
return ctx.GenerateResult(), nil
}
2 changes: 2 additions & 0 deletions pkg/output/output.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -159,6 +159,7 @@ type ResultEvent struct {
Lines []int `json:"matched-line,omitempty"`

FileToIndexPosition map[string]int `json:"-"`
Error string `json:"error,omitempty"`
}

// NewStandardWriter creates a new output writer based on user configurations
Expand Down Expand Up @@ -360,6 +361,7 @@ func (w *StandardWriter) WriteFailure(wrappedEvent *InternalWrappedEvent) error
Timestamp: time.Now(),
//FIXME: this is workaround to encode the template when no results were found
TemplateEncoded: w.encodeTemplate(types.ToString(event["template-path"])),
Error: types.ToString(event["error"]),
}
return w.Write(data)
}
Expand Down
1 change: 1 addition & 0 deletions pkg/protocols/code/code.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -255,6 +255,7 @@ func (request *Request) MakeResultEventItem(wrapped *output.InternalWrappedEvent
Timestamp: time.Now(),
MatcherStatus: true,
TemplateEncoded: request.options.EncodeTemplate(),
Error: types.ToString(wrapped.InternalEvent["error"]),
}
return data
}
Expand Down
1 change: 1 addition & 0 deletions pkg/protocols/dns/operators.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -123,6 +123,7 @@ func (request *Request) MakeResultEventItem(wrapped *output.InternalWrappedEvent
Request: types.ToString(wrapped.InternalEvent["request"]),
Response: types.ToString(wrapped.InternalEvent["raw"]),
TemplateEncoded: request.options.EncodeTemplate(),
Error: types.ToString(wrapped.InternalEvent["error"]),
}
return data
}
Expand Down
1 change: 1 addition & 0 deletions pkg/protocols/file/operators.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -108,6 +108,7 @@ func (request *Request) MakeResultEventItem(wrapped *output.InternalWrappedEvent
Response: types.ToString(wrapped.InternalEvent["raw"]),
Timestamp: time.Now(),
TemplateEncoded: request.options.EncodeTemplate(),
Error: types.ToString(wrapped.InternalEvent["error"]),
}
return data
}
1 change: 1 addition & 0 deletions pkg/protocols/headless/operators.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -139,6 +139,7 @@ func (request *Request) MakeResultEventItem(wrapped *output.InternalWrappedEvent
Request: types.ToString(wrapped.InternalEvent["request"]),
Response: types.ToString(wrapped.InternalEvent["data"]),
TemplateEncoded: request.options.EncodeTemplate(),
Error: types.ToString(wrapped.InternalEvent["error"]),
}
return data
}
1 change: 1 addition & 0 deletions pkg/protocols/http/operators.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -165,6 +165,7 @@ func (request *Request) MakeResultEventItem(wrapped *output.InternalWrappedEvent
Response: request.truncateResponse(wrapped.InternalEvent["response"]),
CURLCommand: types.ToString(wrapped.InternalEvent["curl-command"]),
TemplateEncoded: request.options.EncodeTemplate(),
Error: types.ToString(wrapped.InternalEvent["error"]),
}
return data
}
Expand Down
Loading

0 comments on commit ce5df9c

Please sign in to comment.