Skip to content

Commit

Permalink
support stop-at-first-match for network templates
Browse files Browse the repository at this point in the history
  • Loading branch information
RamanaReddy0M committed Aug 22, 2024
1 parent 8e1072e commit bb4a644
Show file tree
Hide file tree
Showing 2 changed files with 33 additions and 4 deletions.
4 changes: 4 additions & 0 deletions pkg/protocols/network/network.go
Original file line number Diff line number Diff line change
Expand Up @@ -85,6 +85,10 @@ type Request struct {
// SelfContained specifies if the request is self-contained.
SelfContained bool `yaml:"-" json:"-"`

// description: |
// StopAtFirstMatch stops the execution of the requests and template as soon as a match is found.
StopAtFirstMatch bool `yaml:"stop-at-first-match,omitempty" json:"stop-at-first-match,omitempty" jsonschema:"title=stop at first match,description=Stop the execution after a match is found"`

// description: |
// ports is post processed list of ports to scan (obtained from Port)
ports []string `yaml:"-" json:"-"`
Expand Down
33 changes: 29 additions & 4 deletions pkg/protocols/network/request.go
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ import (
"os"
"strings"
"sync"
"sync/atomic"
"time"

"github.com/pkg/errors"
Expand Down Expand Up @@ -99,6 +100,16 @@ func (request *Request) ExecuteWithResults(target *contextargs.Context, metadata
gologger.Verbose().Msgf("[%v] got errors while checking open ports: %s\n", request.options.TemplateID, err)
}

// stop at first match if requested
atomicBool := &atomic.Bool{}
shouldStopAtFirstMatch := request.StopAtFirstMatch || request.options.StopAtFirstMatch || request.options.Options.StopAtFirstMatch
wrappedCallback := func(event *output.InternalWrappedEvent) {
if event != nil && event.HasOperatorResult() {
atomicBool.Store(true)
}
callback(event)
}

for _, port := range ports {
input := target.Clone()
// use network port updates input with new port requested in template file
Expand All @@ -107,9 +118,12 @@ func (request *Request) ExecuteWithResults(target *contextargs.Context, metadata
if err := input.UseNetworkPort(port, request.ExcludePorts); err != nil {
gologger.Debug().Msgf("Could not network port from constants: %s\n", err)
}
if err := request.executeOnTarget(input, visitedAddresses, metadata, previous, callback); err != nil {
if err := request.executeOnTarget(input, visitedAddresses, metadata, previous, wrappedCallback); err != nil {
return err
}
if shouldStopAtFirstMatch && atomicBool.Load() {
break
}
}

return nil
Expand Down Expand Up @@ -141,6 +155,16 @@ func (request *Request) executeOnTarget(input *contextargs.Context, visited maps
variablesMap := request.options.Variables.Evaluate(variables)
variables = generators.MergeMaps(variablesMap, variables, request.options.Constants)

// stop at first match if requested
atomicBool := &atomic.Bool{}
shouldStopAtFirstMatch := request.StopAtFirstMatch || request.options.StopAtFirstMatch || request.options.Options.StopAtFirstMatch
wrappedCallback := func(event *output.InternalWrappedEvent) {
if event != nil && event.HasOperatorResult() {
atomicBool.Store(true)
}
callback(event)
}

for _, kv := range request.addresses {
select {
case <-input.Context().Done():
Expand All @@ -154,12 +178,13 @@ func (request *Request) executeOnTarget(input *contextargs.Context, visited maps
continue
}
visited.Set(actualAddress, struct{}{})

if err := request.executeAddress(variables, actualAddress, address, input, kv.tls, previous, callback); err != nil {
if err := request.executeAddress(variables, actualAddress, address, input, kv.tls, previous, wrappedCallback); err != nil {
outputEvent := request.responseToDSLMap("", "", "", address, "")
callback(&output.InternalWrappedEvent{InternalEvent: outputEvent})
gologger.Warning().Msgf("[%v] Could not make network request for (%s) : %s\n", request.options.TemplateID, actualAddress, err)
continue
}
if shouldStopAtFirstMatch && atomicBool.Load() {
break
}
}
return nil
Expand Down

0 comments on commit bb4a644

Please sign in to comment.