Add support for multiple uncover queries (#5132)

* Add support for multiple uncover queries

* Fix issue with multiple engines
This commit is contained in:
Ramana Reddy 2024-05-10 22:16:01 +05:30 committed by GitHub
parent 6067b78bcf
commit 0f983d8204
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -11,6 +11,7 @@ import (
"github.com/projectdiscovery/uncover" "github.com/projectdiscovery/uncover"
"github.com/projectdiscovery/uncover/sources" "github.com/projectdiscovery/uncover/sources"
mapsutil "github.com/projectdiscovery/utils/maps" mapsutil "github.com/projectdiscovery/utils/maps"
sliceutil "github.com/projectdiscovery/utils/slice"
stringsutil "github.com/projectdiscovery/utils/strings" stringsutil "github.com/projectdiscovery/utils/strings"
) )
@ -84,9 +85,21 @@ func GetUncoverTargetsFromMetadata(ctx context.Context, templates []*templates.T
if queriesMap[engine] == nil { if queriesMap[engine] == nil {
queriesMap[engine] = []string{} queriesMap[engine] = []string{}
} }
switch v := v.(type) {
case []interface{}:
qs := queriesMap[engine]
for _, vv := range v {
qs = append(qs, fmt.Sprint(vv))
}
queriesMap[engine] = qs
default:
queriesMap[engine] = append(queriesMap[engine], fmt.Sprint(v)) queriesMap[engine] = append(queriesMap[engine], fmt.Sprint(v))
} }
} }
}
for engine, queries := range queriesMap {
queriesMap[engine] = sliceutil.Dedupe(queries)
}
keys := mapsutil.GetKeys(queriesMap) keys := mapsutil.GetKeys(queriesMap)
gologger.Info().Msgf("Running uncover queries from template against: %s", strings.Join(keys, ",")) gologger.Info().Msgf("Running uncover queries from template against: %s", strings.Join(keys, ","))
result := make(chan string, runtime.NumCPU()) result := make(chan string, runtime.NumCPU())
@ -97,8 +110,11 @@ func GetUncoverTargetsFromMetadata(ctx context.Context, templates []*templates.T
// TODO: add support for map[engine]queries in uncover // TODO: add support for map[engine]queries in uncover
// Note below implementation is intentionally sequential to avoid burning all the API keys // Note below implementation is intentionally sequential to avoid burning all the API keys
counter := 0 counter := 0
outerLoop:
for eng, queries := range queriesMap { for eng, queries := range queriesMap {
if opts.Limit > 0 && counter >= opts.Limit {
break
}
// create new uncover options for each engine // create new uncover options for each engine
uncoverOpts := &uncover.Options{ uncoverOpts := &uncover.Options{
Agents: []string{eng}, Agents: []string{eng},
@ -114,18 +130,20 @@ func GetUncoverTargetsFromMetadata(ctx context.Context, templates []*templates.T
gologger.Error().Msgf("Could not get targets using %v engine from uncover: %s", eng, err) gologger.Error().Msgf("Could not get targets using %v engine from uncover: %s", eng, err)
return return
} }
innerLoop:
for { for {
select { select {
case <-ctx.Done(): case <-ctx.Done():
return return
case res, ok := <-ch: case res, ok := <-ch:
if !ok { if !ok {
return continue outerLoop
} }
result <- res result <- res
counter++ counter++
if opts.Limit > 0 && counter >= opts.Limit { if opts.Limit > 0 && counter >= opts.Limit {
return break innerLoop
} }
} }
} }