diff --git a/pkg/protocols/network/network.go b/pkg/protocols/network/network.go index 70d618dcb5..c90f5e7019 100644 --- a/pkg/protocols/network/network.go +++ b/pkg/protocols/network/network.go @@ -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:"-"` diff --git a/pkg/protocols/network/request.go b/pkg/protocols/network/request.go index 5fa8609d51..3a0d2af22a 100644 --- a/pkg/protocols/network/request.go +++ b/pkg/protocols/network/request.go @@ -8,6 +8,7 @@ import ( "os" "strings" "sync" + "sync/atomic" "time" "github.com/pkg/errors" @@ -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 @@ -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 @@ -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(): @@ -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 err