From 8eea7345bd0e62291dac345aad60292e8b64ffad Mon Sep 17 00:00:00 2001 From: Mzack9999 Date: Mon, 16 Nov 2020 00:40:32 +0100 Subject: [PATCH] targets hmap + global rate limit + clistats windows compatibility --- v2/go.mod | 6 +- v2/go.sum | 12 +++ v2/internal/progress/progress.go | 55 ++++++---- v2/internal/runner/options.go | 2 +- v2/internal/runner/processor.go | 20 ++-- v2/internal/runner/runner.go | 141 +++++++++++++------------- v2/pkg/executer/executer_dns.go | 4 + v2/pkg/executer/executer_http.go | 9 +- v2/pkg/globalratelimiter/ratelimit.go | 82 --------------- 9 files changed, 144 insertions(+), 187 deletions(-) delete mode 100644 v2/pkg/globalratelimiter/ratelimit.go diff --git a/v2/go.mod b/v2/go.mod index a466229e..fd863e52 100644 --- a/v2/go.mod +++ b/v2/go.mod @@ -6,6 +6,7 @@ require ( github.com/Knetic/govaluate v3.0.0+incompatible github.com/blang/semver v3.5.1+incompatible github.com/d5/tengo/v2 v2.6.2 + github.com/eiannone/keyboard v0.0.0-20200508000154-caf4b762e807 // indirect github.com/golang/snappy v0.0.2 // indirect github.com/google/go-github/v32 v32.1.0 github.com/json-iterator/go v1.1.10 @@ -15,7 +16,7 @@ require ( github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect github.com/modern-go/reflect2 v1.0.1 // indirect github.com/pkg/errors v0.9.1 - github.com/projectdiscovery/clistats v0.0.4 + github.com/projectdiscovery/clistats v0.0.5-0.20201115190457-4a189e365b54 github.com/projectdiscovery/collaborator v0.0.0-20201023080839-2aa1290ed09d github.com/projectdiscovery/gologger v1.0.1 github.com/projectdiscovery/hmap v0.0.0-20201026185329-db41b5717bcb @@ -24,11 +25,10 @@ require ( github.com/projectdiscovery/retryabledns v1.0.4 github.com/projectdiscovery/retryablehttp-go v1.0.1 github.com/remeh/sizedwaitgroup v1.0.0 + github.com/segmentio/ksuid v1.0.3 github.com/spaolacci/murmur3 v1.1.0 github.com/stretchr/testify v1.5.1 go.uber.org/ratelimit v0.1.0 - golang.org/x/crypto v0.0.0-20201016220609-9e8e0b390897 // indirect golang.org/x/net v0.0.0-20201029221708-28c70e62bb1d - golang.org/x/sys v0.0.0-20201029080932-201ba4db2418 // indirect gopkg.in/yaml.v2 v2.3.0 ) diff --git a/v2/go.sum b/v2/go.sum index 0a417498..9f9e1360 100644 --- a/v2/go.sum +++ b/v2/go.sum @@ -19,6 +19,7 @@ github.com/d5/tengo/v2 v2.6.2/go.mod h1:XRGjEs5I9jYIKTxly6HCF8oiiilk5E/RYXOZ5b0D github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/eiannone/keyboard v0.0.0-20200508000154-caf4b762e807/go.mod h1:Xoiu5VdKMvbRgHuY7+z64lhu/7lvax/22nzASF6GrO8= github.com/fsnotify/fsnotify v1.4.7 h1:IXs+QLmnXW2CcXuY+8Mzv/fWEsPGWxqefPtCP5CnV9I= github.com/fsnotify/fsnotify v1.4.7/go.mod h1:jwhsz4b93w/PPRr/qN1Yymfu8t87LnFCMoQvtojpjFo= github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U= @@ -68,6 +69,9 @@ github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4= github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/projectdiscovery/clistats v0.0.4/go.mod h1:+RpHbYUY4FJ6yQiDlL4Ux44jK3p4AwFKUsr6neu43RA= +github.com/projectdiscovery/clistats v0.0.5-0.20201115190457-4a189e365b54 h1:7Yz219bGdbH5oxwFs4dTpqcm48HG5DhMxOqw2kIaiX4= +github.com/projectdiscovery/clistats v0.0.5-0.20201115190457-4a189e365b54/go.mod h1:lV6jUHAv2bYWqrQstqW8iVIydKJhWlVaLl3Xo9ioVGg= github.com/projectdiscovery/collaborator v0.0.0-20201023080839-2aa1290ed09d h1:iHb2v6VX1Fjl9IX8UaJFLKLej+KyJJ59W1pWIRCug4Q= github.com/projectdiscovery/collaborator v0.0.0-20201023080839-2aa1290ed09d/go.mod h1:M7Csn+hQVDOLCEEFkj6dazmtgG1tIqJpbuPHlRlpYGQ= github.com/projectdiscovery/collaborator v0.0.0-20201023192422-4a25cca69447 h1:1+NNwnPD0Rm7r4NJaX6IQwi+DiTJUc3JC6k0Zuc0/Qc= @@ -81,6 +85,7 @@ github.com/projectdiscovery/hmap v0.0.0-20201026185329-db41b5717bcb h1:y+O2ZCGIC github.com/projectdiscovery/hmap v0.0.0-20201026185329-db41b5717bcb/go.mod h1:VDEfgzkKQdq7iGTKz8Ooul0NuYHQ8qiDs6r8bPD1Sb0= github.com/projectdiscovery/httpx v1.0.2 h1:g7EeRAPckZgWcHkcAH2Qzv9MkRACVRLF+T2LJcM7SCk= github.com/projectdiscovery/httpx v1.0.2/go.mod h1:OwvMc5ogx69xukKXY6kIrDP6dgOYr4VtEWyr6o573Xs= +github.com/projectdiscovery/nuclei v1.1.7 h1:5Z1fBHcjyAuuI89xcCzv8tYK7b6ucqLxs+mCC/nJjno= github.com/projectdiscovery/rawhttp v0.0.4 h1:O5IreNGk83d4xTD9e6SpkKbX0sHTs8K1Q33Bz4eYl2E= github.com/projectdiscovery/rawhttp v0.0.4/go.mod h1:PQERZAhAv7yxI/hR6hdDPgK1WTU56l204BweXrBec+0= github.com/projectdiscovery/rawhttp v0.0.5-0.20201030111302-fd184be37926 h1:aAh1EqrurioC8OeOJAzZG2gYN7MEmMNc2UWLJnQNV78= @@ -92,6 +97,8 @@ github.com/projectdiscovery/retryablehttp-go v1.0.1/go.mod h1:SrN6iLZilNG1X4neq1 github.com/remeh/sizedwaitgroup v1.0.0 h1:VNGGFwNo/R5+MJBf6yrsr110p0m4/OX4S3DCy7Kyl5E= github.com/remeh/sizedwaitgroup v1.0.0/go.mod h1:3j2R4OIe/SeS6YDhICBy22RWjJC5eNCJ1V+9+NVNYlo= github.com/rs/xid v1.2.1/go.mod h1:+uKXf+4Djp6Md1KODXJxgGQPKngRmWyn10oCKFzNHOQ= +github.com/segmentio/ksuid v1.0.3 h1:FoResxvleQwYiPAVKe1tMUlEirodZqlqglIuFsdDntY= +github.com/segmentio/ksuid v1.0.3/go.mod h1:/XUiZBD3kVx5SmUOl55voK5yeAbBNNIed+2O73XgrPE= github.com/spaolacci/murmur3 v0.0.0-20180118202830-f09979ecbc72/go.mod h1:JwIasOWyU6f++ZhiEuf87xNszmSA2myDM2Kzu9HwQUA= github.com/spaolacci/murmur3 v1.1.0 h1:7c1g84S4BPRrfL5Xrdp6fOJ206sU9y293DDHaoy0bLI= github.com/spaolacci/murmur3 v1.1.0/go.mod h1:JwIasOWyU6f++ZhiEuf87xNszmSA2myDM2Kzu9HwQUA= @@ -112,6 +119,8 @@ golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8U golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto= golang.org/x/crypto v0.0.0-20201016220609-9e8e0b390897 h1:pLI5jrR7OSLijeIDcmRxNmw2api+jEfxLoykJVice/E= golang.org/x/crypto v0.0.0-20201016220609-9e8e0b390897/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto= +golang.org/x/crypto v0.0.0-20201112155050-0c6587e931a9 h1:umElSU9WZirRdgu2yFHY0ayQkEnKiOC1TtM3fWXFnoU= +golang.org/x/crypto v0.0.0-20201112155050-0c6587e931a9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto= golang.org/x/mod v0.1.1-0.20191105210325-c90efee705ee/go.mod h1:QqPTAvyqsEbceGzBzNggFXnrqF1CaUcvgkdR5Ot7KZg= golang.org/x/net v0.0.0-20180906233101-161cd47e91fd/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= golang.org/x/net v0.0.0-20181220203305-927f97764cc3/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= @@ -131,6 +140,7 @@ golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJ golang.org/x/sys v0.0.0-20180909124046-d0be0721c37e/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20190626150813-e07cf5db2756/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20190924154521-2837fb4f24fe/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20200323222414-85ca7c5b95cd/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20200810151505-1b9f1253b3ed/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= @@ -139,6 +149,8 @@ golang.org/x/sys v0.0.0-20201022201747-fb209a7c41cd h1:WgqgiQvkiZWz7XLhphjt2GI2G golang.org/x/sys v0.0.0-20201022201747-fb209a7c41cd/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20201029080932-201ba4db2418 h1:HlFl4V6pEMziuLXyRkm5BIYq1y1GAbb02pRlWvI54OM= golang.org/x/sys v0.0.0-20201029080932-201ba4db2418/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20201113233024-12cec1faf1ba h1:xmhUJGQGbxlod18iJGqVEp9cHIPLl7QiX2aA3to708s= +golang.org/x/sys v0.0.0-20201113233024-12cec1faf1ba/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= golang.org/x/text v0.3.3 h1:cokOdA+Jmi5PJGXLlLllQSgYigAEfHXJAERHVMaCc2k= golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= diff --git a/v2/internal/progress/progress.go b/v2/internal/progress/progress.go index d157961e..a30fbf32 100644 --- a/v2/internal/progress/progress.go +++ b/v2/internal/progress/progress.go @@ -11,6 +11,7 @@ import ( // Progress is a progress instance for showing program stats type Progress struct { + active bool stats clistats.StatisticsClient tickDuration time.Duration } @@ -24,40 +25,56 @@ func NewProgress(active bool) *Progress { tickDuration = -1 } - progress := &Progress{ - tickDuration: tickDuration, - stats: clistats.New(), + var progress Progress + if active { + stats, err := clistats.New() + if err != nil { + gologger.Warningf("Couldn't create progress engine: %s\n", err) + } + progress.active = active + progress.stats = stats + progress.tickDuration = tickDuration } - return progress + + return &progress } // Init initializes the progress display mechanism by setting counters, etc. func (p *Progress) Init(hostCount int64, rulesCount int, requestCount int64) { - p.stats.AddStatic("templates", rulesCount) - p.stats.AddStatic("hosts", hostCount) - p.stats.AddStatic("startedAt", time.Now()) - p.stats.AddCounter("requests", uint64(0)) - p.stats.AddCounter("errors", uint64(0)) - p.stats.AddCounter("total", uint64(requestCount)) - - _ = p.stats.Start(makePrintCallback(), p.tickDuration) + if p.active { + p.stats.AddStatic("templates", rulesCount) + p.stats.AddStatic("hosts", hostCount) + p.stats.AddStatic("startedAt", time.Now()) + p.stats.AddCounter("requests", uint64(0)) + p.stats.AddCounter("errors", uint64(0)) + p.stats.AddCounter("total", uint64(requestCount)) + if err := p.stats.Start(makePrintCallback(), p.tickDuration); err != nil { + gologger.Warningf("Couldn't start statistics: %s\n", err) + } + } } // AddToTotal adds a value to the total request count func (p *Progress) AddToTotal(delta int64) { - p.stats.IncrementCounter("total", int(delta)) + if p.active { + p.stats.IncrementCounter("total", int(delta)) + } } // Update progress tracking information and increments the request counter by one unit. func (p *Progress) Update() { - p.stats.IncrementCounter("requests", 1) + if p.active { + p.stats.IncrementCounter("requests", 1) + } } // Drop drops the specified number of requests from the progress bar total. // This may be the case when uncompleted requests are encountered and shouldn't be part of the total count. func (p *Progress) Drop(count int64) { - // mimic dropping by incrementing the completed requests - p.stats.IncrementCounter("errors", int(count)) + if p.active { + // mimic dropping by incrementing the completed requests + p.stats.IncrementCounter("errors", int(count)) + } } const bufferSize = 128 @@ -120,5 +137,9 @@ func fmtDuration(d time.Duration) string { // Stop stops the progress bar execution func (p *Progress) Stop() { - _ = p.stats.Stop() + if p.active { + if err := p.stats.Stop(); err != nil { + gologger.Warningf("Couldn't stop statistics: %s\n", err) + } + } } diff --git a/v2/internal/runner/options.go b/v2/internal/runner/options.go index 723d5a25..b7d4b377 100644 --- a/v2/internal/runner/options.go +++ b/v2/internal/runner/options.go @@ -88,7 +88,7 @@ func ParseOptions() *Options { flag.BoolVar(&options.JSONRequests, "json-requests", false, "Write requests/responses for matches in JSON output") flag.BoolVar(&options.EnableProgressBar, "pbar", false, "Enable the progress bar") flag.BoolVar(&options.TemplateList, "tl", false, "List available templates") - flag.IntVar(&options.RateLimit, "rate-limit", 150, "Rate-Limit Per Target (maximum requests/second") + flag.IntVar(&options.RateLimit, "rate-limit", 150, "Rate-Limit (maximum requests/second") flag.BoolVar(&options.StopAtFirstMatch, "stop-at-first-match", false, "Stop processing http requests at first match (this may break template/workflow logic)") flag.IntVar(&options.BulkSize, "bulk-size", 25, "Maximum Number of hosts analyzed in parallel per template") flag.IntVar(&options.TemplateThreads, "c", 10, "Maximum Number of templates executed in parallel") diff --git a/v2/internal/runner/processor.go b/v2/internal/runner/processor.go index d1f846b8..80184327 100644 --- a/v2/internal/runner/processor.go +++ b/v2/internal/runner/processor.go @@ -1,7 +1,6 @@ package runner import ( - "bufio" "context" "fmt" "net/http/cookiejar" @@ -50,6 +49,7 @@ func (r *Runner) processTemplateWithList(p *progress.Progress, template *templat ColoredOutput: !r.options.NoColor, Colorizer: r.colorizer, Decolorizer: r.decolorizer, + RateLimiter: r.ratelimiter, }) case *requests.BulkHTTPRequest: httpExecuter, err = executer.NewHTTPExecuter(&executer.HTTPOptions{ @@ -73,6 +73,7 @@ func (r *Runner) processTemplateWithList(p *progress.Progress, template *templat StopAtFirstMatch: r.options.StopAtFirstMatch, PF: r.pf, Dialer: r.dialer, + RateLimiter: r.ratelimiter, }) } @@ -87,9 +88,8 @@ func (r *Runner) processTemplateWithList(p *progress.Progress, template *templat wg := sizedwaitgroup.New(r.options.BulkSize) - scanner := bufio.NewScanner(strings.NewReader(r.input)) - for scanner.Scan() { - URL := scanner.Text() + r.hm.Scan(func(k, _ []byte) error { + URL := string(k) wg.Add() go func(URL string) { defer wg.Done() @@ -110,7 +110,9 @@ func (r *Runner) processTemplateWithList(p *progress.Progress, template *templat gologger.Warningf("[%s] Could not execute step: %s\n", r.colorizer.Colorizer.BrightBlue(template.ID), result.Error) } }(URL) - } + + return nil + }) wg.Wait() @@ -132,9 +134,8 @@ func (r *Runner) processWorkflowWithList(p *progress.Progress, workflow *workflo wg := sizedwaitgroup.New(r.options.BulkSize) - scanner := bufio.NewScanner(strings.NewReader(r.input)) - for scanner.Scan() { - targetURL := scanner.Text() + r.hm.Scan(func(k, _ []byte) error { + targetURL := string(k) wg.Add() go func(targetURL string) { @@ -168,7 +169,8 @@ func (r *Runner) processWorkflowWithList(p *progress.Progress, workflow *workflo } } }(targetURL) - } + return nil + }) wg.Wait() diff --git a/v2/internal/runner/runner.go b/v2/internal/runner/runner.go index d0e52234..49f18181 100644 --- a/v2/internal/runner/runner.go +++ b/v2/internal/runner/runner.go @@ -2,9 +2,6 @@ package runner import ( "bufio" - "fmt" - "io" - "io/ioutil" "os" "regexp" "strings" @@ -12,6 +9,7 @@ import ( "github.com/logrusorgru/aurora" "github.com/pkg/errors" "github.com/projectdiscovery/gologger" + "github.com/projectdiscovery/hmap/store/hybrid" "github.com/projectdiscovery/httpx/common/cache" "github.com/projectdiscovery/nuclei/v2/internal/bufwriter" "github.com/projectdiscovery/nuclei/v2/internal/progress" @@ -19,18 +17,16 @@ import ( "github.com/projectdiscovery/nuclei/v2/pkg/atomicboolean" "github.com/projectdiscovery/nuclei/v2/pkg/collaborator" "github.com/projectdiscovery/nuclei/v2/pkg/colorizer" - "github.com/projectdiscovery/nuclei/v2/pkg/globalratelimiter" "github.com/projectdiscovery/nuclei/v2/pkg/projectfile" "github.com/projectdiscovery/nuclei/v2/pkg/templates" "github.com/projectdiscovery/nuclei/v2/pkg/workflows" "github.com/remeh/sizedwaitgroup" + "go.uber.org/ratelimit" ) // Runner is a client for running the enumeration process. type Runner struct { - input string inputCount int64 - tempFile string traceLog tracelog.Log @@ -52,6 +48,12 @@ type Runner struct { // http dialer dialer cache.DialerFunc + + // rate limiter + ratelimiter ratelimit.Limiter + + // input deduplication + hm *hybrid.HybridMap } // New creates a new client for running enumeration process. @@ -94,79 +96,67 @@ func New(options *Options) (*Runner, error) { runner.readNucleiIgnoreFile() } - // If we have stdin, write it to a new file - if options.Stdin { - tempInput, err := ioutil.TempFile("", "stdin-input-*") - - if err != nil { - return nil, err - } - - if _, err := io.Copy(tempInput, os.Stdin); err != nil { - return nil, err - } - - runner.tempFile = tempInput.Name() - tempInput.Close() - } - // If we have single target, write it to a new file - if options.Target != "" { - tempInput, err := ioutil.TempFile("", "stdin-input-*") - if err != nil { - return nil, err - } - - fmt.Fprintf(tempInput, "%s\n", options.Target) - runner.tempFile = tempInput.Name() - tempInput.Close() + if hm, err := hybrid.New(hybrid.DefaultDiskOptions); err != nil { + gologger.Fatalf("Could not create temporary input file: %s\n", err) + } else { + runner.hm = hm } - // Setup input, handle a list of hosts as argument - var err error - - var input *os.File - - if options.Targets != "" { - input, err = os.Open(options.Targets) - } else if options.Stdin || options.Target != "" { - input, err = os.Open(runner.tempFile) - } - - if err != nil { - gologger.Fatalf("Could not open targets file '%s': %s\n", options.Targets, err) - } - - // Sanitize input and pre-compute total number of targets - var usedInput = make(map[string]struct{}) - - dupeCount := 0 - sb := strings.Builder{} - scanner := bufio.NewScanner(input) runner.inputCount = 0 + dupeCount := 0 + + // Handle single target + if options.Target != "" { + runner.inputCount++ + runner.hm.Set(options.Target, nil) + } + + // Handle stdin + if options.Stdin { + scanner := bufio.NewScanner(os.Stdin) + for scanner.Scan() { + url := strings.TrimSpace(scanner.Text()) + // skip empty lines + if url == "" { + continue + } + + // skip dupes + if _, ok := runner.hm.Get(url); ok { + dupeCount++ + continue + } - for scanner.Scan() { - url := scanner.Text() - // skip empty lines - if url == "" { - continue - } - // deduplication - if _, ok := usedInput[url]; !ok { - usedInput[url] = struct{}{} runner.inputCount++ - - // allocate global rate limiters - globalratelimiter.Add(url, options.RateLimit) - - sb.WriteString(url) - sb.WriteString("\n") - } else { - dupeCount++ + runner.hm.Set(url, nil) } } - input.Close() - runner.input = sb.String() + // Handle taget file + if options.Targets != "" { + input, err := os.Open(options.Targets) + if err != nil { + gologger.Fatalf("Could not open targets file '%s': %s\n", options.Targets, err) + } + scanner := bufio.NewScanner(input) + for scanner.Scan() { + url := strings.TrimSpace(scanner.Text()) + // skip empty lines + if url == "" { + continue + } + + // skip dupes + if _, ok := runner.hm.Get(url); ok { + dupeCount++ + continue + } + + runner.inputCount++ + runner.hm.Set(url, nil) + } + input.Close() + } if dupeCount > 0 { gologger.Labelf("Supplied input was automatically deduplicated (%d removed).", dupeCount) @@ -199,11 +189,18 @@ func New(options *Options) (*Runner, error) { } // Create Dialer + var err error runner.dialer, err = cache.NewDialer(cache.DefaultOptions) if err != nil { return nil, err } + if options.RateLimit > 0 { + runner.ratelimiter = ratelimit.New(options.RateLimit) + } else { + runner.ratelimiter = ratelimit.NewUnlimited() + } + return runner, nil } @@ -212,7 +209,7 @@ func (r *Runner) Close() { if r.output != nil { r.output.Close() } - os.Remove(r.tempFile) + r.hm.Close() if r.pf != nil { r.pf.Close() } diff --git a/v2/pkg/executer/executer_dns.go b/v2/pkg/executer/executer_dns.go index 983fb377..65f19635 100644 --- a/v2/pkg/executer/executer_dns.go +++ b/v2/pkg/executer/executer_dns.go @@ -15,6 +15,7 @@ import ( "github.com/projectdiscovery/nuclei/v2/pkg/requests" "github.com/projectdiscovery/nuclei/v2/pkg/templates" retryabledns "github.com/projectdiscovery/retryabledns" + "go.uber.org/ratelimit" ) // DNSExecuter is a client for performing a DNS request @@ -32,6 +33,7 @@ type DNSExecuter struct { template *templates.Template dnsRequest *requests.DNSRequest writer *bufwriter.Writer + ratelimiter ratelimit.Limiter colorizer colorizer.NucleiColorizer decolorizer *regexp.Regexp @@ -59,6 +61,7 @@ type DNSOptions struct { Colorizer colorizer.NucleiColorizer Decolorizer *regexp.Regexp + RateLimiter ratelimit.Limiter } // NewDNSExecuter creates a new DNS executer from a template @@ -79,6 +82,7 @@ func NewDNSExecuter(options *DNSOptions) *DNSExecuter { coloredOutput: options.ColoredOutput, colorizer: options.Colorizer, decolorizer: options.Decolorizer, + ratelimiter: options.RateLimiter, } return executer diff --git a/v2/pkg/executer/executer_http.go b/v2/pkg/executer/executer_http.go index a65d36da..f224dc07 100644 --- a/v2/pkg/executer/executer_http.go +++ b/v2/pkg/executer/executer_http.go @@ -26,7 +26,6 @@ import ( "github.com/projectdiscovery/nuclei/v2/internal/tracelog" "github.com/projectdiscovery/nuclei/v2/pkg/colorizer" "github.com/projectdiscovery/nuclei/v2/pkg/generators" - "github.com/projectdiscovery/nuclei/v2/pkg/globalratelimiter" "github.com/projectdiscovery/nuclei/v2/pkg/matchers" projetctfile "github.com/projectdiscovery/nuclei/v2/pkg/projectfile" "github.com/projectdiscovery/nuclei/v2/pkg/requests" @@ -34,6 +33,7 @@ import ( "github.com/projectdiscovery/rawhttp" "github.com/projectdiscovery/retryablehttp-go" "github.com/remeh/sizedwaitgroup" + "go.uber.org/ratelimit" "golang.org/x/net/proxy" ) @@ -65,6 +65,7 @@ type HTTPExecuter struct { jsonRequest bool noMeta bool stopAtFirstMatch bool + ratelimiter ratelimit.Limiter } // HTTPOptions contains configuration options for the HTTP executer. @@ -90,6 +91,7 @@ type HTTPOptions struct { StopAtFirstMatch bool PF *projetctfile.ProjectFile Dialer cache.DialerFunc + RateLimiter ratelimit.Limiter } // NewHTTPExecuter creates a new HTTP executer from a template @@ -144,6 +146,7 @@ func NewHTTPExecuter(options *HTTPOptions) (*HTTPExecuter, error) { decolorizer: options.Decolorizer, stopAtFirstMatch: options.StopAtFirstMatch, pf: options.PF, + ratelimiter: options.RateLimiter, } return executer, nil @@ -220,7 +223,7 @@ func (e *HTTPExecuter) ExecuteParallelHTTP(p *progress.Progress, reqURL string) go func(httpRequest *requests.HTTPRequest) { defer swg.Done() - globalratelimiter.Take(reqURL) + e.ratelimiter.Take() // If the request was built correctly then execute it err = e.handleHTTP(reqURL, httpRequest, dynamicvalues, result, "") @@ -351,7 +354,7 @@ func (e *HTTPExecuter) ExecuteHTTP(p *progress.Progress, reqURL string) *Result result.Error = err p.Drop(remaining) } else { - globalratelimiter.Take(reqURL) + e.ratelimiter.Take() // If the request was built correctly then execute it format := "%s_" + strconv.Itoa(requestNumber) err = e.handleHTTP(reqURL, httpRequest, dynamicvalues, result, format) diff --git a/v2/pkg/globalratelimiter/ratelimit.go b/v2/pkg/globalratelimiter/ratelimit.go deleted file mode 100644 index d79f82a8..00000000 --- a/v2/pkg/globalratelimiter/ratelimit.go +++ /dev/null @@ -1,82 +0,0 @@ -package globalratelimiter - -import ( - "sync" - - "go.uber.org/ratelimit" -) - -var defaultrwmutex sync.RWMutex -var defaultGlobalRateLimiter GlobalRateLimiter = GlobalRateLimiter{ratesLimiters: make(map[string]ratelimit.Limiter)} - -type GlobalRateLimiter struct { - sync.RWMutex - ratesLimiters map[string]ratelimit.Limiter -} - -func Add(k string, rateLimit int) { - defaultrwmutex.Lock() - defer defaultrwmutex.Unlock() - - if rateLimit > 0 { - defaultGlobalRateLimiter.ratesLimiters[k] = ratelimit.New(rateLimit) - } else { - defaultGlobalRateLimiter.ratesLimiters[k] = ratelimit.NewUnlimited() - } -} - -func Take(k string) { - rl := take(k) - - rl.Take() -} - -func take(k string) ratelimit.Limiter { - defaultrwmutex.RLock() - defer defaultrwmutex.RUnlock() //nolint - - return defaultGlobalRateLimiter.ratesLimiters[k] -} - -func Del(k string, rateLimit int) { - defaultrwmutex.Lock() - defer defaultrwmutex.Unlock() //nolint - - delete(defaultGlobalRateLimiter.ratesLimiters, k) -} - -func New() *GlobalRateLimiter { - var globalRateLimiter GlobalRateLimiter - globalRateLimiter.ratesLimiters = make(map[string]ratelimit.Limiter) - return &globalRateLimiter -} - -func (grl *GlobalRateLimiter) Add(k string, rateLimit int) { - grl.Lock() - defer grl.Unlock() - - if rateLimit > 0 { - grl.ratesLimiters[k] = ratelimit.New(rateLimit) - } else { - grl.ratesLimiters[k] = ratelimit.NewUnlimited() - } -} - -func (grl *GlobalRateLimiter) take(k string) ratelimit.Limiter { - grl.RLock() - defer grl.RUnlock() //nolint - - return grl.ratesLimiters[k] -} - -func (grl *GlobalRateLimiter) Take(k string) { - rl := grl.take(k) - rl.Take() -} - -func (grl *GlobalRateLimiter) Del(k string, rateLimit int) { - grl.Lock() - defer grl.Unlock() - - delete(grl.ratesLimiters, k) -}