Merge branch 'CasperGN-rate-limit'

dev
Mzack9999 2020-09-24 01:09:55 +02:00
commit fcd231aed0
4 changed files with 17 additions and 0 deletions

View File

@ -17,5 +17,9 @@ require (
github.com/projectdiscovery/retryablehttp-go v1.0.1
github.com/vbauerster/mpb/v5 v5.3.0
golang.org/x/net v0.0.0-20200923182212-328152dc79b1
github.com/stretchr/testify v1.5.1
github.com/vbauerster/mpb/v5 v5.2.4
golang.org/x/net v0.0.0-20200707034311-ab3426394381
golang.org/x/time v0.0.0-20200630173020-3af7569d3a1e
gopkg.in/yaml.v2 v2.3.0
)

View File

@ -72,6 +72,8 @@ golang.org/x/sys v0.0.0-20200323222414-85ca7c5b95cd/go.mod h1:h1NjWce9XRLGQEsW7w
golang.org/x/sys v0.0.0-20200810151505-1b9f1253b3ed h1:WBkVNH1zd9jg/dK4HCM4lNANnmd12EHC9z+LmcCG4ns=
golang.org/x/sys v0.0.0-20200810151505-1b9f1253b3ed/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
golang.org/x/time v0.0.0-20200630173020-3af7569d3a1e h1:EHBhcS0mlXEAVwNyO2dLfjToGsyY4j24pTs2ScHnX7s=
golang.org/x/time v0.0.0-20200630173020-3af7569d3a1e/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
golang.org/x/tools v0.0.0-20191216052735-49a3e744a425/go.mod h1:TB2adYChydJhpapKDTa4BR/hXlZSLoq2Wpct/0txZ28=
golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
google.golang.org/appengine v1.1.0/go.mod h1:EbEs0AVv82hx2wNQdGPgUI5lhzA/G0D9YwlJXL52JkM=

View File

@ -38,6 +38,7 @@ type Options struct {
ProxySocksURL string // ProxySocksURL is the URL for the proxy socks server
CustomHeaders requests.CustomHeaders // Custom global headers
TemplatesDirectory string // TemplatesDirectory is the directory to use for storing templates
RateLimit int // Rate-Limit of requests per specified target
}
type multiStringFlag []string
@ -78,6 +79,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, "rl", 9999999, "Rate-Limit of requests per specified target") // 9999999 to avoid limiting
flag.Parse()

View File

@ -10,6 +10,9 @@ import (
"regexp"
"strings"
"sync"
"time"
"golang.org/x/time/rate"
"github.com/logrusorgru/aurora"
"github.com/projectdiscovery/gologger"
@ -238,6 +241,8 @@ func (r *Runner) RunEnumeration() {
gologger.Errorf("Could not find any valid input URLs.")
} else if totalRequests > 0 || hasWorkflows {
ctx := context.Background()
// Limiter that will add to the tokenbucket every second and set the max size to -rl flag
rateLimit := rate.NewLimiter(rate.Every(1*time.Second), r.options.RateLimit)
// tracks global progress and captures stdout/stderr until p.Wait finishes
p := r.progress
p.InitProgressbar(r.inputCount, templateCount, totalRequests)
@ -246,6 +251,10 @@ func (r *Runner) RunEnumeration() {
wgtemplates.Add(1)
go func(template interface{}) {
defer wgtemplates.Done()
err := rateLimit.Wait(ctx)
if err != nil {
gologger.Errorf("Issue with rate-limit")
}
switch tt := template.(type) {
case *templates.Template:
for _, request := range tt.RequestsDNS {