From de168e0be9f0afceedba21b536005c0a42ef9d1c Mon Sep 17 00:00:00 2001 From: M09Ic Date: Sun, 25 Aug 2024 23:06:10 +0800 Subject: [PATCH] enhance basepool and brutepool structure --- internal/option.go | 3 +- internal/pool/brutepool.go | 82 ++++++++++++++++++++++++++++++++------ internal/pool/config.go | 2 +- internal/pool/pool.go | 67 +------------------------------ internal/runner.go | 8 ++-- 5 files changed, 78 insertions(+), 84 deletions(-) diff --git a/internal/option.go b/internal/option.go index 627a461..bfe6ada 100644 --- a/internal/option.go +++ b/internal/option.go @@ -265,6 +265,7 @@ func (opt *Option) NewRunner() (*Runner, error) { Option: opt, taskCh: make(chan *Task), outputCh: make(chan *pkg.Baseline, 256), + poolwg: &sync.WaitGroup{}, outwg: &sync.WaitGroup{}, fuzzyCh: make(chan *pkg.Baseline, 256), Headers: make(map[string]string), @@ -541,7 +542,7 @@ func (opt *Option) BuildWords(r *Runner) error { logs.Log.Logf(pkg.LogVerbose, "Loaded %d word from %s", len(dicts[i]), f) } - if len(dicts) == 0 && opt.Word == "" { + if len(dicts) == 0 && opt.Word == "" && len(opt.Rules) == 0 && len(opt.AppendRule) == 0 { r.IsCheck = true } diff --git a/internal/pool/brutepool.go b/internal/pool/brutepool.go index 1e93db8..27b8391 100644 --- a/internal/pool/brutepool.go +++ b/internal/pool/brutepool.go @@ -10,11 +10,14 @@ import ( "github.com/chainreactors/spray/pkg" "github.com/chainreactors/utils/iutils" "github.com/chainreactors/words" + "github.com/chainreactors/words/mask" + "github.com/chainreactors/words/rule" "github.com/panjf2000/ants/v2" "github.com/valyala/fasthttp" "golang.org/x/time/rate" "math/rand" "net/url" + "path" "strings" "sync" "sync/atomic" @@ -378,12 +381,9 @@ func (pool *BrutePool) Invoke(v interface{}) { pool.locker.Lock() pool.index = bl pool.locker.Unlock() - if bl.Status == 200 || (bl.Status/100) == 3 { - // 保留index输出结果 - pool.wg.Add(1) - pool.doCrawl(bl) - pool.putToOutput(bl) - } + pool.wg.Add(1) + pool.doCrawl(bl) + pool.putToOutput(bl) pool.initwg.Done() case parsers.CheckSource: if bl.ErrString != "" { @@ -521,13 +521,10 @@ func (pool *BrutePool) Handler() { } if bl.IsValid || bl.IsFuzzy { - pool.wg.Add(2) + pool.wg.Add(3) pool.doCrawl(bl) - pool.doRule(bl) - if iutils.IntsContains(pkg.WhiteStatus, bl.Status) || iutils.IntsContains(pkg.UniqueStatus, bl.Status) { - pool.wg.Add(1) - pool.doAppendWords(bl) - } + pool.doAppendRule(bl) + pool.doAppendWords(bl) } // 如果要进行递归判断, 要满足 bl有效, mod为path-spray, 当前深度小于最大递归深度 @@ -549,6 +546,67 @@ func (pool *BrutePool) Handler() { pool.analyzeDone = true } +func (pool *BrutePool) doAppendRule(bl *pkg.Baseline) { + if pool.AppendRule == nil || bl.Source == parsers.RuleSource { + pool.wg.Done() + return + } + + go func() { + defer pool.wg.Done() + for u := range rule.RunAsStream(pool.AppendRule.Expressions, path.Base(bl.Path)) { + pool.addAddition(&Unit{ + path: pkg.Dir(bl.Url.Path) + u, + source: parsers.RuleSource, + }) + } + }() +} + +func (pool *BrutePool) doAppendWords(bl *pkg.Baseline) { + if pool.AppendWords == nil || bl.Source == parsers.AppendSource || bl.Source == parsers.RuleSource { + // 防止自身递归 + pool.wg.Done() + return + } + + go func() { + defer pool.wg.Done() + for _, u := range pool.AppendWords { + pool.addAddition(&Unit{ + path: pkg.SafePath(bl.Path, u), + source: parsers.AppendSource, + }) + } + }() +} + +func (pool *BrutePool) doAppend(bl *pkg.Baseline) { + pool.wg.Add(2) + pool.doAppendWords(bl) + pool.doAppendRule(bl) +} + +func (pool *BrutePool) doActive() { + defer pool.wg.Done() + for _, u := range pkg.ActivePath { + pool.addAddition(&Unit{ + path: pool.dir + u[1:], + source: parsers.FingerSource, + }) + } +} + +func (pool *BrutePool) doCommonFile() { + defer pool.wg.Done() + for _, u := range mask.SpecialWords["common_file"] { + pool.addAddition(&Unit{ + path: pool.dir + u, + source: parsers.CommonFileSource, + }) + } +} + func (pool *BrutePool) PreCompare(resp *ihttp.Response) error { status := resp.StatusCode() if iutils.IntsContains(pkg.WhiteStatus, status) { diff --git a/internal/pool/config.go b/internal/pool/config.go index 447c1ac..f3cae85 100644 --- a/internal/pool/config.go +++ b/internal/pool/config.go @@ -51,7 +51,7 @@ type Config struct { Active bool Bak bool Common bool - Retry int + RetryLimit int RandomUserAgent bool Random string Index string diff --git a/internal/pool/pool.go b/internal/pool/pool.go index 3ba1ddf..c9e9621 100644 --- a/internal/pool/pool.go +++ b/internal/pool/pool.go @@ -7,10 +7,7 @@ import ( "github.com/chainreactors/spray/internal/ihttp" "github.com/chainreactors/spray/pkg" "github.com/chainreactors/words" - "github.com/chainreactors/words/mask" - "github.com/chainreactors/words/rule" "github.com/panjf2000/ants/v2" - "path" "sync" ) @@ -49,50 +46,8 @@ func (pool *BasePool) doRedirect(bl *pkg.Baseline, depth int) { }() } -func (pool *BasePool) doRule(bl *pkg.Baseline) { - if pool.AppendRule == nil { - pool.wg.Done() - return - } - if bl.Source == parsers.RuleSource { - pool.wg.Done() - return - } - - go func() { - defer pool.wg.Done() - for u := range rule.RunAsStream(pool.AppendRule.Expressions, path.Base(bl.Path)) { - pool.addAddition(&Unit{ - path: pkg.Dir(bl.Url.Path) + u, - source: parsers.RuleSource, - }) - } - }() -} - -func (pool *BasePool) doAppendWords(bl *pkg.Baseline) { - if pool.AppendWords == nil { - pool.wg.Done() - return - } - if bl.Source == parsers.AppendSource { - pool.wg.Done() - return - } - - go func() { - defer pool.wg.Done() - for _, u := range pool.AppendWords { - pool.addAddition(&Unit{ - path: pkg.SafePath(bl.Path, u), - source: parsers.AppendSource, - }) - } - }() -} - func (pool *BasePool) doRetry(bl *pkg.Baseline) { - if bl.Retry >= pool.Retry { + if bl.Retry >= pool.RetryLimit { return } pool.wg.Add(1) @@ -106,26 +61,6 @@ func (pool *BasePool) doRetry(bl *pkg.Baseline) { }() } -func (pool *BasePool) doActive() { - defer pool.wg.Done() - for _, u := range pkg.ActivePath { - pool.addAddition(&Unit{ - path: pool.dir + u[1:], - source: parsers.FingerSource, - }) - } -} - -func (pool *BasePool) doCommonFile() { - defer pool.wg.Done() - for _, u := range mask.SpecialWords["common_file"] { - pool.addAddition(&Unit{ - path: pool.dir + u, - source: parsers.CommonFileSource, - }) - } -} - func (pool *BasePool) addAddition(u *Unit) { // 强行屏蔽报错, 防止goroutine泄露 pool.wg.Add(1) diff --git a/internal/runner.go b/internal/runner.go index b70087f..5cd2d9f 100644 --- a/internal/runner.go +++ b/internal/runner.go @@ -30,7 +30,7 @@ type Runner struct { *Option taskCh chan *Task - poolwg sync.WaitGroup + poolwg *sync.WaitGroup outwg *sync.WaitGroup outputCh chan *pkg.Baseline fuzzyCh chan *pkg.Baseline @@ -81,15 +81,15 @@ func (r *Runner) PrepareConfig() *pool.Config { MatchExpr: r.MatchExpr, FilterExpr: r.FilterExpr, RecuExpr: r.RecursiveExpr, - AppendRule: r.AppendRules, - AppendWords: r.AppendWords, + AppendRule: r.AppendRules, // 对有效目录追加规则, 根据rule生成 + AppendWords: r.AppendWords, // 对有效目录追加字典 //IgnoreWaf: r.IgnoreWaf, Crawl: r.Crawl, Scope: r.Scope, Active: r.Finger, Bak: r.Bak, Common: r.Common, - Retry: r.RetryCount, + RetryLimit: r.RetryCount, ClientType: r.ClientType, RandomUserAgent: r.RandomUserAgent, Random: r.Random,