github.com/projectdiscovery/nuclei/v2@v2.9.15/pkg/protocols/common/uncover/uncover.go (about)

     1  package uncover
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"runtime"
     7  	"strings"
     8  
     9  	"github.com/projectdiscovery/gologger"
    10  	"github.com/projectdiscovery/nuclei/v2/pkg/templates"
    11  	"github.com/projectdiscovery/uncover"
    12  	"github.com/projectdiscovery/uncover/sources"
    13  	mapsutil "github.com/projectdiscovery/utils/maps"
    14  	stringsutil "github.com/projectdiscovery/utils/strings"
    15  )
    16  
    17  // returns csv string of uncover supported agents
    18  func GetUncoverSupportedAgents() string {
    19  	u, _ := uncover.New(&uncover.Options{})
    20  	return strings.Join(u.AllAgents(), ",")
    21  }
    22  
    23  // GetTargetsFromUncover returns targets from uncover
    24  func GetTargetsFromUncover(ctx context.Context, outputFormat string, opts *uncover.Options) (chan string, error) {
    25  	u, err := uncover.New(opts)
    26  	if err != nil {
    27  		return nil, err
    28  	}
    29  	resChan, err := u.Execute(ctx)
    30  	if err != nil {
    31  		return nil, err
    32  	}
    33  	outputChan := make(chan string) // buffered channel
    34  	go func() {
    35  		defer close(outputChan)
    36  		for {
    37  			select {
    38  			case <-ctx.Done():
    39  				return
    40  			case res, ok := <-resChan:
    41  				if !ok {
    42  					return
    43  				}
    44  				if res.Error != nil {
    45  					// only log in verbose mode
    46  					gologger.Verbose().Msgf("uncover: %v", res.Error)
    47  					continue
    48  				}
    49  				outputChan <- processUncoverOutput(res, outputFormat)
    50  			}
    51  		}
    52  	}()
    53  	return outputChan, nil
    54  }
    55  
    56  // processUncoverOutput returns output string depending on uncover field
    57  func processUncoverOutput(result sources.Result, outputFormat string) string {
    58  	if (result.IP == "" || result.Port == 0) && stringsutil.ContainsAny(outputFormat, "ip", "port") {
    59  		// if ip or port is not present, fallback to using host
    60  		outputFormat = "host"
    61  	}
    62  	replacer := strings.NewReplacer(
    63  		"ip", result.IP,
    64  		"host", result.Host,
    65  		"port", fmt.Sprint(result.Port),
    66  		"url", result.Url,
    67  	)
    68  	return replacer.Replace(outputFormat)
    69  }
    70  
    71  // GetUncoverTargetsFromMetadata returns targets from uncover metadata
    72  func GetUncoverTargetsFromMetadata(ctx context.Context, templates []*templates.Template, outputFormat string, opts *uncover.Options) chan string {
    73  	// contains map[engine]queries
    74  	queriesMap := make(map[string][]string)
    75  	for _, template := range templates {
    76  	innerLoop:
    77  		for k, v := range template.Info.Metadata {
    78  			if !strings.HasSuffix(k, "-query") {
    79  				// this is not a query
    80  				// query keys are like shodan-query, fofa-query, etc
    81  				continue innerLoop
    82  			}
    83  			engine := strings.TrimSuffix(k, "-query")
    84  			if queriesMap[engine] == nil {
    85  				queriesMap[engine] = []string{}
    86  			}
    87  			queriesMap[engine] = append(queriesMap[engine], fmt.Sprint(v))
    88  		}
    89  	}
    90  	keys := mapsutil.GetKeys(queriesMap)
    91  	gologger.Info().Msgf("Running uncover queries from template against: %s", strings.Join(keys, ","))
    92  	result := make(chan string, runtime.NumCPU())
    93  	go func() {
    94  		defer close(result)
    95  		// unfortunately uncover doesn't support execution of map[engine]queries
    96  		// if queries are given they are executed against all engines which is not what we want
    97  		// TODO: add support for map[engine]queries in uncover
    98  		// Note below implementation is intentionally sequential to avoid burning all the API keys
    99  		counter := 0
   100  
   101  		for eng, queries := range queriesMap {
   102  			// create new uncover options for each engine
   103  			uncoverOpts := &uncover.Options{
   104  				Agents:        []string{eng},
   105  				Queries:       queries,
   106  				Limit:         opts.Limit,
   107  				MaxRetry:      opts.MaxRetry,
   108  				Timeout:       opts.Timeout,
   109  				RateLimit:     opts.RateLimit,
   110  				RateLimitUnit: opts.RateLimitUnit,
   111  			}
   112  			ch, err := GetTargetsFromUncover(ctx, outputFormat, uncoverOpts)
   113  			if err != nil {
   114  				gologger.Error().Msgf("Could not get targets using %v engine from uncover: %s", eng, err)
   115  				return
   116  			}
   117  			for {
   118  				select {
   119  				case <-ctx.Done():
   120  					return
   121  				case res, ok := <-ch:
   122  					if !ok {
   123  						return
   124  					}
   125  					result <- res
   126  					counter++
   127  					if opts.Limit > 0 && counter >= opts.Limit {
   128  						return
   129  					}
   130  				}
   131  			}
   132  		}
   133  	}()
   134  	return result
   135  }