Skip to content

Commit

Permalink
remove use of iterate() in flow (#4688)
Browse files Browse the repository at this point in the history
  • Loading branch information
tarunKoyalwar committed Jan 28, 2024
1 parent e102cae commit 0371846
Show file tree
Hide file tree
Showing 5 changed files with 79 additions and 4 deletions.
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
}
}

0 comments on commit 0371846

Please sign in to comment.