github.com/projectdiscovery/nuclei/v2@v2.9.15/internal/runner/inputs.go (about)

     1  package runner
     2  
     3  import (
     4  	"sync/atomic"
     5  	"time"
     6  
     7  	"github.com/pkg/errors"
     8  	"github.com/projectdiscovery/gologger"
     9  	"github.com/projectdiscovery/hmap/store/hybrid"
    10  	"github.com/projectdiscovery/httpx/common/httpx"
    11  	"github.com/projectdiscovery/nuclei/v2/pkg/protocols/common/contextargs"
    12  	"github.com/projectdiscovery/nuclei/v2/pkg/utils"
    13  	stringsutil "github.com/projectdiscovery/utils/strings"
    14  	"github.com/remeh/sizedwaitgroup"
    15  )
    16  
    17  const probeBulkSize = 50
    18  
    19  // initializeTemplatesHTTPInput initializes the http form of input
    20  // for any loaded http templates if input is in non-standard format.
    21  func (r *Runner) initializeTemplatesHTTPInput() (*hybrid.HybridMap, error) {
    22  	hm, err := hybrid.New(hybrid.DefaultDiskOptions)
    23  	if err != nil {
    24  		return nil, errors.Wrap(err, "could not create temporary input file")
    25  	}
    26  	gologger.Info().Msgf("Running httpx on input host")
    27  
    28  	var bulkSize = probeBulkSize
    29  	if r.options.BulkSize > probeBulkSize {
    30  		bulkSize = r.options.BulkSize
    31  	}
    32  
    33  	httpxOptions := httpx.DefaultOptions
    34  	httpxOptions.RetryMax = r.options.Retries
    35  	httpxOptions.Timeout = time.Duration(r.options.Timeout) * time.Second
    36  	httpxClient, err := httpx.New(&httpxOptions)
    37  	if err != nil {
    38  		return nil, errors.Wrap(err, "could not create httpx client")
    39  	}
    40  
    41  	// Probe the non-standard URLs and store them in cache
    42  	swg := sizedwaitgroup.New(bulkSize)
    43  	count := int32(0)
    44  	r.hmapInputProvider.Scan(func(value *contextargs.MetaInput) bool {
    45  		if stringsutil.HasPrefixAny(value.Input, "http://", "https://") {
    46  			return true
    47  		}
    48  
    49  		swg.Add()
    50  		go func(input *contextargs.MetaInput) {
    51  			defer swg.Done()
    52  
    53  			if result := utils.ProbeURL(input.Input, httpxClient); result != "" {
    54  				atomic.AddInt32(&count, 1)
    55  				_ = hm.Set(input.Input, []byte(result))
    56  			}
    57  		}(value)
    58  		return true
    59  	})
    60  	swg.Wait()
    61  
    62  	gologger.Info().Msgf("Found %d URL from httpx", atomic.LoadInt32(&count))
    63  	return hm, nil
    64  }