Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

remove use of iterate() in flow #4688

Merged
merged 1 commit into from
Jan 28, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
11 changes: 11 additions & 0 deletions cmd/integration-test/flow.go
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ var flowTestcases = []TestCaseInfo{
{Path: "flow/conditional-flow.yaml", TestCase: &conditionalFlow{}},
{Path: "flow/conditional-flow-negative.yaml", TestCase: &conditionalFlowNegative{}},
{Path: "flow/iterate-values-flow.yaml", TestCase: &iterateValuesFlow{}},
{Path: "flow/iterate-one-value-flow.yaml", TestCase: &iterateOneValueFlow{}},
{Path: "flow/dns-ns-probe.yaml", TestCase: &dnsNsProbe{}},
{Path: "flow/flow-hide-matcher.yaml", TestCase: &flowHideMatcher{}},
}
Expand Down Expand Up @@ -70,6 +71,16 @@ func (t *iterateValuesFlow) Execute(filePath string) error {
return expectResultsCount(results, 2)
}

type iterateOneValueFlow struct{}

func (t *iterateOneValueFlow) Execute(filePath string) error {
results, err := testutils.RunNucleiTemplateAndGetResults(filePath, "https://scanme.sh", debug)
if err != nil {
return err
}
return expectResultsCount(results, 1)
}

type dnsNsProbe struct{}

func (t *dnsNsProbe) Execute(filePath string) error {
Expand Down
34 changes: 34 additions & 0 deletions integration_tests/flow/iterate-one-value-flow.yaml
Original file line number Diff line number Diff line change
@@ -0,0 +1,34 @@
id: flow-iterate-one-value-flow

info:
name: Test Flow Iterate One Value Flow
author: pdteam
severity: info

flow: |
http(1)
for(let value of template.extracted){
set("value", value)
http(2)
}

http:
- method: GET
path:
- "{{BaseURL}}"

extractors:
- type: regex
name: extracted
internal: true
regex:
- "[ok]+"

- method: GET
path:
- "{{BaseURL}}/{{value}}"

matchers:
- type: word
words:
- "ok"
12 changes: 12 additions & 0 deletions pkg/tmplexec/flow/flow_executor.go
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,9 @@ type FlowExecutor struct {
// logic related variables
results *atomic.Bool
allErrs mapsutil.SyncLockMap[string, error]
// these are keys whose values are meant to be flatten before executing
// a request ex: if dynamic extractor returns ["value"] it will be converted to "value"
flattenKeys []string
}

// NewFlowExecutor creates a new flow executor from a list of requests
Expand Down Expand Up @@ -157,6 +160,15 @@ func (f *FlowExecutor) Compile() error {
opts.reqIDS = append(opts.reqIDS, types.ToString(value))
}
}
// before executing any protocol function flatten tracked values
if len(f.flattenKeys) > 0 {
ctx := f.options.GetTemplateCtx(f.ctx.Input.MetaInput)
for _, key := range f.flattenKeys {
if value, ok := ctx.Get(key); ok {
ctx.Set(key, flatten(value))
}
}
}
return f.jsVM.ToValue(f.requestExecutor(reqMap, opts))
}
}
Expand Down
9 changes: 5 additions & 4 deletions pkg/tmplexec/flow/flow_internal.go
Original file line number Diff line number Diff line change
Expand Up @@ -100,11 +100,12 @@ func (f *FlowExecutor) protocolResultCallback(req protocols.Request, matcherStat
for k, v := range result.OperatorsResult.DynamicValues {
// if length of v is 1 then remove slice and convert it to single value
if len(v) == 1 {
f.options.GetTemplateCtx(f.ctx.Input.MetaInput).Set(k, v[0])
} else {
// if not let user handle it in flow ex: `for(let val of template.extracted)`
f.options.GetTemplateCtx(f.ctx.Input.MetaInput).Set(k, v)
// add it to flatten keys list so it will be flattened to a string later
f.flattenKeys = append(f.flattenKeys, k)
}
// always preserve extracted value type
f.options.GetTemplateCtx(f.ctx.Input.MetaInput).Set(k, v)

}
}
} else if !result.HasOperatorResult() && !hasOperators(req.GetCompiledOperators()) {
Expand Down
17 changes: 17 additions & 0 deletions pkg/tmplexec/flow/util.go
Original file line number Diff line number Diff line change
Expand Up @@ -21,3 +21,20 @@ func hasOperators(all []*operators.Operators) bool {
}
return false
}

func flatten(v interface{}) interface{} {
switch v := v.(type) {
case []interface{}:
if len(v) == 1 {
return v[0]
}
return v
case []string:
if len(v) == 1 {
return v[0]
}
return v
default:
return v
}
}