From 758a27468479b782cb6d92edf3b9857691cba193 Mon Sep 17 00:00:00 2001 From: M09Ic Date: Thu, 12 Jan 2023 17:41:44 +0800 Subject: [PATCH] =?UTF-8?q?=E6=96=B0=E5=A2=9E--rate-limit,=20=E7=94=A8?= =?UTF-8?q?=E6=9D=A5=E9=99=90=E5=88=B6=E5=8D=95=E4=B8=AApool=E8=AF=B7?= =?UTF-8?q?=E6=B1=82=E9=80=9F=E7=8E=87?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- go.mod | 1 + go.sum | 2 ++ internal/option.go | 40 +++++++++++++++------------ internal/pool.go | 69 +++++++++++++++++++++++++--------------------- internal/runner.go | 2 ++ pkg/config.go | 1 + 6 files changed, 67 insertions(+), 48 deletions(-) diff --git a/go.mod b/go.mod index a8cb9a7..05790b4 100644 --- a/go.mod +++ b/go.mod @@ -32,5 +32,6 @@ require ( github.com/twmb/murmur3 v1.1.6 // indirect github.com/valyala/bytebufferpool v1.0.0 // indirect golang.org/x/sys v0.2.0 // indirect + golang.org/x/time v0.3.0 // indirect gopkg.in/yaml.v2 v2.4.0 // indirect ) diff --git a/go.sum b/go.sum index 6a14b1e..d9793b2 100644 --- a/go.sum +++ b/go.sum @@ -103,6 +103,8 @@ golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= golang.org/x/text v0.3.2/go.mod h1:bEr9sfX3Q8Zfm5fL9x+3itogRgK3+ptLWKqgva+5dAk= golang.org/x/text v0.3.6/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ= +golang.org/x/time v0.3.0 h1:rg5rLMjNzMS1RkNLzCG38eapWhnYLFYXDXj2gOlr8j4= +golang.org/x/time v0.3.0/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ= golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= diff --git a/internal/option.go b/internal/option.go index 0c32482..9d2f039 100644 --- a/internal/option.go +++ b/internal/option.go @@ -21,6 +21,7 @@ type Option struct { InputOptions `group:"Input Options"` FunctionOptions `group:"Function Options"` OutputOptions `group:"Output Options"` + PluginOptions `group:"Plugin Options"` RequestOptions `group:"Request Options"` ModeOptions `group:"Modify Options"` MiscOptions `group:"Miscellaneous Options"` @@ -28,15 +29,15 @@ type Option struct { type InputOptions struct { ResumeFrom string `long:"resume"` - URL []string `short:"u" long:"url" description:"String, Multi, input baseurl, e.g.: http://google.com"` + URL []string `short:"u" long:"url" description:"Strings, input baseurl, e.g.: http://google.com"` URLFile string `short:"l" long:"list" description:"File, input filename"` Raw string `long:"raw" description:"File, input raw request filename"` Offset int `long:"offset" description:"Int, wordlist offset"` Limit int `long:"limit" description:"Int, wordlist limit, start with offset. e.g.: --offset 1000 --limit 100"` Dictionaries []string `short:"d" long:"dict" description:"Files, Multi,dict files, e.g.: -d 1.txt -d 2.txt"` Word string `short:"w" long:"word" description:"String, word generate dsl, e.g.: -w test{?ld#4}"` - Rules []string `short:"r" long:"rules" description:"Files, Multi, rule files, e.g.: -r rule1.txt -r rule2.txt"` - AppendRule []string `long:"append-rule" description:"File, when found valid path , use append rule generator new word with current path"` + Rules []string `short:"r" long:"rules" description:"Files, rule files, e.g.: -r rule1.txt -r rule2.txt"` + AppendRule []string `long:"append-rule" description:"Files, when found valid path , use append rule generator new word with current path"` FilterRule string `long:"filter-rule" description:"String, filter rule, e.g.: --rule-filter '>8 <4'"` } @@ -46,15 +47,15 @@ type FunctionOptions struct { RemoveExtensions string `long:"remove-extension" description:"String, remove extensions (separated by commas), e.g.: --remove-extension jsp,jspx"` Uppercase bool `short:"U" long:"uppercase" desvcription:"Bool, upper wordlist, e.g.: --uppercase"` Lowercase bool `short:"L" long:"lowercase" description:"Bool, lower wordlist, e.g.: --lowercase"` - Prefixes []string `long:"prefix" description:"Strings, Multi, add prefix, e.g.: --prefix aaa --prefix bbb"` - Suffixes []string `long:"suffix" description:"Strings, Multi, add suffix, e.g.: --suffix aaa --suffix bbb"` - Replaces map[string]string `long:"replace" description:"Strings, Multi, replace string, e.g.: --replace aaa:bbb --replace ccc:ddd"` + Prefixes []string `long:"prefix" description:"Strings, add prefix, e.g.: --prefix aaa --prefix bbb"` + Suffixes []string `long:"suffix" description:"Strings, add suffix, e.g.: --suffix aaa --suffix bbb"` + Replaces map[string]string `long:"replace" description:"Strings, replace string, e.g.: --replace aaa:bbb --replace ccc:ddd"` } type OutputOptions struct { Match string `long:"match" description:"String, custom match function, e.g.: --match current.Status != 200" json:"match,omitempty"` Filter string `long:"filter" description:"String, custom filter function, e.g.: --filter current.Body contains 'hello'" json:"filter,omitempty"` - Extracts []string `long:"extract" description:"String, Multi, extract response, e.g.: --extract js --extract ip --extract version:(.*?)" json:"extracts,omitempty"` + Extracts []string `long:"extract" description:"Strings, extract response, e.g.: --extract js --extract ip --extract version:(.*?)" json:"extracts,omitempty"` OutputFile string `short:"f" long:"file" description:"String, output filename" json:"output_file,omitempty"` Format string `short:"F" long:"format" description:"String, output format, e.g.: --format 1.json"` FuzzyFile string `long:"fuzzy-file" description:"String, fuzzy output filename" json:"fuzzy_file,omitempty"` @@ -66,27 +67,31 @@ type OutputOptions struct { } type RequestOptions struct { - Headers []string `long:"header" description:"String, Multi, custom headers, e.g.: --headers 'Auth: example_auth'"` + Headers []string `long:"header" description:"Strings, custom headers, e.g.: --headers 'Auth: example_auth'"` UserAgent string `long:"user-agent" description:"String, custom user-agent, e.g.: --user-agent Custom"` RandomUserAgent bool `long:"random-agent" description:"Bool, use random with default user-agent"` - Cookie []string `long:"cookie" description:"String, Multi, custom cookie"` + Cookie []string `long:"cookie" description:"Strings, custom cookie"` ReadAll bool `long:"read-all" description:"Bool, read all response body"` MaxBodyLength int `long:"max-length" default:"100" description:"Int, max response body length (kb), default 100k, e.g. -max-length 1000"` } +type PluginOptions struct { + Advance bool `short:"a" long:"advance" description:"Bool, enable crawl and active"` + Active bool `long:"active" description:"Bool, enable active finger detect"` + Bak bool `long:"bak" description:"Bool, enable bak found"` + FileBak bool `long:"file-bak" description:"Bool, enable valid result bak found, equal --append-rule rule/filebak.txt"` + Common bool `long:"common" description:"Bool, enable common file found"` + Crawl bool `long:"crawl" description:"Bool, enable crawl"` + CrawlDepth int `long:"crawl-depth" default:"3" description:"Int, crawl depth"` + CrawlScope string `long:"crawl-scope" description:"Int, crawl scope (todo)"` +} + type ModeOptions struct { - Advance bool `short:"a" long:"advance" description:"Bool, enable crawl and active"` - Active bool `long:"active" description:"Bool, enable active finger detect"` - Crawl bool `long:"crawl" description:"Bool, enable crawl"` - Bak bool `long:"bak" description:"Bool, enable bak found"` - FileBak bool `long:"file-bak" description:"Bool, enable valid result bak found, equal --append-rule rule/filebak.txt"` - Common bool `long:"common" description:"Bool, enable common file found"` + RateLimit int `long:"rate-limit" default:"0" description:"Int, request rate limit (rate/s), e.g.: --rate-limit 100"` Force bool `long:"force" description:"Bool, skip error break"` CheckOnly bool `long:"check-only" description:"Bool, check only"` Recursive string `long:"recursive" default:"current.IsDir()" description:"String,custom recursive rule, e.g.: --recursive current.IsDir()"` Depth int `long:"depth" default:"0" description:"Int, recursive depth"` - CrawlDepth int `long:"crawl-depth" default:"3" description:"Int, crawl depth"` - CrawlScope string `long:"crawl-scope" description:"Int, crawl scope (todo)"` CheckPeriod int `long:"check-period" default:"200" description:"Int, check period when request"` ErrPeriod int `long:"error-period" default:"10" description:"Int, check period when error"` BreakThreshold int `long:"error-threshold" default:"20" description:"Int, break when the error exceeds the threshold "` @@ -121,6 +126,7 @@ func (opt *Option) PrepareRunner() (*Runner, error) { PoolSize: opt.PoolSize, Mod: opt.Mod, Timeout: opt.Timeout, + RateLimit: opt.RateLimit, Deadline: opt.Deadline, Headers: make(map[string]string), Offset: opt.Offset, diff --git a/internal/pool.go b/internal/pool.go index c09b5b6..ebf1b26 100644 --- a/internal/pool.go +++ b/internal/pool.go @@ -13,6 +13,7 @@ import ( "github.com/chainreactors/words/rule" "github.com/panjf2000/ants/v2" "github.com/valyala/fasthttp" + "golang.org/x/time/rate" "net/url" "path" "strconv" @@ -50,8 +51,9 @@ func NewPool(ctx context.Context, config *pkg.Config) (*Pool, error) { tempCh: make(chan *pkg.Baseline, config.Thread), checkCh: make(chan int), additionCh: make(chan *Unit, 100), - wg: sync.WaitGroup{}, + waiter: sync.WaitGroup{}, initwg: sync.WaitGroup{}, + limiter: rate.NewLimiter(rate.Limit(config.RateLimit), 1), reqCount: 1, failedCount: 1, } @@ -127,7 +129,7 @@ func NewPool(ctx context.Context, config *pkg.Config) (*Pool, error) { // 如果要进行递归判断, 要满足 bl有效, mod为path-spray, 当前深度小于最大递归深度 if bl.IsValid { - pool.wg.Add(2) + pool.waiter.Add(2) pool.doCrawl(bl) pool.doRule(bl) if bl.RecuDepth < MaxRecursion { @@ -137,7 +139,7 @@ func NewPool(ctx context.Context, config *pkg.Config) (*Pool, error) { } } pool.OutputCh <- bl - pool.wg.Done() + pool.waiter.Done() } pool.analyzeDone = true @@ -170,8 +172,9 @@ type Pool struct { urls map[string]struct{} analyzeDone bool worder *words.Worder + limiter *rate.Limiter locker sync.Mutex - wg sync.WaitGroup + waiter sync.WaitGroup initwg sync.WaitGroup // 初始化用, 之后改成锁 } @@ -237,17 +240,17 @@ func (pool *Pool) genReq(s string) (*ihttp.Request, error) { func (pool *Pool) Run(ctx context.Context, offset, limit int) { pool.worder.RunWithRules() if pool.Active { - pool.wg.Add(1) + pool.waiter.Add(1) go pool.doActive() } if pool.Bak { - pool.wg.Add(1) + pool.waiter.Add(1) go pool.doBak() } if pool.Common { - pool.wg.Add(1) + pool.waiter.Add(1) go pool.doCommonFile() } @@ -256,7 +259,7 @@ func (pool *Pool) Run(ctx context.Context, offset, limit int) { wait := func() { if !worderDone { worderDone = true - pool.wg.Wait() + pool.waiter.Wait() close(closeCh) } } @@ -280,7 +283,7 @@ Loop: continue } - pool.wg.Add(1) + pool.waiter.Add(1) pool.urls[u] = struct{}{} pool.reqPool.Invoke(newUnit(pool.safePath(u), WordSource)) // 原样的目录拼接, 输入了几个"/"就是几个, 适配java的目录解析 case source := <-pool.checkCh: @@ -296,7 +299,7 @@ Loop: } if _, ok := pool.urls[unit.path]; ok { logs.Log.Debugf("[%s] duplicate path: %s, skipped", pkg.GetSourceName(unit.source), pool.base+unit.path) - pool.wg.Done() + pool.waiter.Done() } else { pool.urls[unit.path] = struct{}{} pool.reqPool.Invoke(unit) @@ -310,12 +313,16 @@ Loop: } } - pool.wg.Wait() + pool.waiter.Wait() pool.Statistor.EndTime = time.Now().Unix() pool.Close() } func (pool *Pool) Invoke(v interface{}) { + if pool.RateLimit != 0 { + pool.limiter.Wait(pool.ctx) + } + atomic.AddInt32(&pool.Statistor.ReqTotal, 1) unit := v.(*Unit) req, err := pool.genReq(unit.path) @@ -355,7 +362,7 @@ func (pool *Pool) Invoke(v interface{}) { // 手动处理重定向 if bl.IsValid && unit.source != CheckSource && bl.RedirectURL != "" { - pool.wg.Add(1) + pool.waiter.Add(1) pool.doRedirect(bl, unit.depth) } @@ -378,7 +385,7 @@ func (pool *Pool) Invoke(v interface{}) { pool.locker.Lock() pool.index = bl pool.locker.Unlock() - pool.wg.Add(1) + pool.waiter.Add(1) pool.doCrawl(bl) if bl.Status == 200 || (bl.Status/100) == 3 { pool.OutputCh <- bl @@ -524,13 +531,13 @@ func (pool *Pool) Upgrade(bl *pkg.Baseline) error { } func (pool *Pool) doRedirect(bl *pkg.Baseline, depth int) { - defer pool.wg.Done() + defer pool.waiter.Done() if depth >= MaxRedirect { return } reURL := FormatURL(bl.Url.Path, bl.RedirectURL) - pool.wg.Add(1) + pool.waiter.Add(1) go pool.addAddition(&Unit{ path: reURL, source: RedirectSource, @@ -541,24 +548,24 @@ func (pool *Pool) doRedirect(bl *pkg.Baseline, depth int) { func (pool *Pool) doCrawl(bl *pkg.Baseline) { if !pool.Crawl || bl.ReqDepth >= MaxCrawl { - pool.wg.Done() + pool.waiter.Done() return } bl.CollectURL() if bl.URLs == nil { - pool.wg.Done() + pool.waiter.Done() return } go func() { - defer pool.wg.Done() + defer pool.waiter.Done() for _, u := range bl.URLs { if u = FormatURL(bl.Url.Path, u); u == "" { continue } // 通过map去重, 只有新的url才会进入到该逻辑 - pool.wg.Add(1) + pool.waiter.Add(1) pool.addAddition(&Unit{ path: u, source: CrawlSource, @@ -571,18 +578,18 @@ func (pool *Pool) doCrawl(bl *pkg.Baseline) { func (pool *Pool) doRule(bl *pkg.Baseline) { if pool.AppendRule == nil { - pool.wg.Done() + pool.waiter.Done() return } if bl.Source == int(RuleSource) { - pool.wg.Done() + pool.waiter.Done() return } go func() { - defer pool.wg.Done() + defer pool.waiter.Done() for u := range rule.RunAsStream(pool.AppendRule.Expressions, path.Base(bl.Path)) { - pool.wg.Add(1) + pool.waiter.Add(1) pool.addAddition(&Unit{ path: Dir(bl.Url.Path) + u, source: RuleSource, @@ -592,9 +599,9 @@ func (pool *Pool) doRule(bl *pkg.Baseline) { } func (pool *Pool) doActive() { - defer pool.wg.Done() + defer pool.waiter.Done() for _, u := range pkg.ActivePath { - pool.wg.Add(1) + pool.waiter.Add(1) pool.addAddition(&Unit{ path: pool.dir + u[1:], source: ActiveSource, @@ -603,14 +610,14 @@ func (pool *Pool) doActive() { } func (pool *Pool) doBak() { - defer pool.wg.Done() + defer pool.waiter.Done() worder, err := words.NewWorderWithDsl("{?0}.{@bak_ext}", [][]string{pkg.BakGenerator(pool.url.Host)}, nil) if err != nil { return } worder.Run() for w := range worder.C { - pool.wg.Add(1) + pool.waiter.Add(1) pool.addAddition(&Unit{ path: pool.dir + w, source: BakSource, @@ -623,7 +630,7 @@ func (pool *Pool) doBak() { } worder.Run() for w := range worder.C { - pool.wg.Add(1) + pool.waiter.Add(1) pool.addAddition(&Unit{ path: pool.dir + w, source: BakSource, @@ -632,9 +639,9 @@ func (pool *Pool) doBak() { } func (pool *Pool) doCommonFile() { - defer pool.wg.Done() + defer pool.waiter.Done() for _, u := range mask.SpecialWords["common_file"] { - pool.wg.Add(1) + pool.waiter.Add(1) pool.addAddition(&Unit{ path: pool.dir + u, source: CommonFileSource, @@ -665,7 +672,7 @@ func (pool *Pool) addAddition(u *Unit) { func (pool *Pool) addFuzzyBaseline(bl *pkg.Baseline) { if _, ok := pool.baselines[bl.Status]; !ok && pkg.IntsContains(FuzzyStatus, bl.Status) { bl.Collect() - pool.wg.Add(1) + pool.waiter.Add(1) pool.doCrawl(bl) pool.baselines[bl.Status] = bl logs.Log.Infof("[baseline.%dinit] %s", bl.Status, bl.Format([]string{"status", "length", "spend", "title", "frame", "redirect"})) diff --git a/internal/runner.go b/internal/runner.go index 8c49882..198ebca 100644 --- a/internal/runner.go +++ b/internal/runner.go @@ -64,6 +64,7 @@ type Runner struct { Progress *uiprogress.Progress Offset int Limit int + RateLimit int Total int Deadline int CheckPeriod int @@ -83,6 +84,7 @@ func (r *Runner) PrepareConfig() *pkg.Config { config := &pkg.Config{ Thread: r.Threads, Timeout: r.Timeout, + RateLimit: r.RateLimit, Headers: r.Headers, Mod: pkg.ModMap[r.Mod], OutputCh: r.OutputCh, diff --git a/pkg/config.go b/pkg/config.go index 72d7615..271609b 100644 --- a/pkg/config.go +++ b/pkg/config.go @@ -24,6 +24,7 @@ type Config struct { Thread int Wordlist []string Timeout int + RateLimit int CheckPeriod int ErrPeriod int BreakThreshold int