support stop-at-first-match for network templates (#5554)

This commit is contained in:
Ramana Reddy 2024-10-14 20:54:58 +05:30 committed by GitHub
parent 2c832f5590
commit 98948d0266
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 33 additions and 4 deletions

View File

@ -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:"-"`

View File

@ -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