diff --git a/internal/pool.go b/internal/pool.go index a88aadc..87941e8 100644 --- a/internal/pool.go +++ b/internal/pool.go @@ -30,8 +30,6 @@ func NewPool(ctx context.Context, config *pkg.Config) (*Pool, error) { cancel: cancel, client: ihttp.NewClient(config.Thread, 2, config.ClientType), worder: words.NewWorder(config.Wordlist), - outputCh: config.OutputCh, - fuzzyCh: config.FuzzyCh, baselines: make(map[int]*pkg.Baseline), tempCh: make(chan *pkg.Baseline, config.Thread), wg: sync.WaitGroup{}, @@ -131,8 +129,10 @@ func NewPool(ctx context.Context, config *pkg.Config) (*Pool, error) { pool.tempCh <- bl if pool.reqCount%pool.checkPeriod == 0 { + pool.reqCount++ go pool.check() } else if pool.failedCount%pool.errPeriod == 0 { + pool.failedCount++ go pool.check() } pool.bar.Done() @@ -142,7 +142,19 @@ func NewPool(ctx context.Context, config *pkg.Config) (*Pool, error) { }) pool.pool = p - go pool.comparing() + go func() { + for bl := range pool.tempCh { + if pool.customCompare != nil { + if pool.customCompare(bl) { + pool.OutputCh <- bl + } + } else { + pool.BaseCompare(bl) + } + } + + pool.analyzeDone = true + }() return pool, nil } @@ -153,8 +165,6 @@ type Pool struct { bar *pkg.Bar ctx context.Context cancel context.CancelFunc - outputCh chan *pkg.Baseline // 输出的chan, 全局统一 - fuzzyCh chan *pkg.Baseline tempCh chan *pkg.Baseline // 待处理的baseline reqCount int failedCount int @@ -166,6 +176,7 @@ type Pool struct { analyzeDone bool genReq func(s string) (*ihttp.Request, error) check func() + customCompare func(*pkg.Baseline) bool worder *words.Worder wg sync.WaitGroup initwg sync.WaitGroup // 初始化用, 之后改成锁 @@ -259,42 +270,37 @@ func (p *Pool) PreCompare(resp *ihttp.Response) error { return nil } -func (p *Pool) comparing() { -Loop: - for bl := range p.tempCh { - if !bl.IsValid { - // precompare 确认无效数据直接送入管道 - p.outputCh <- bl - continue +func (p *Pool) BaseCompare(bl *pkg.Baseline) { + if !bl.IsValid { + // precompare 确认无效数据直接送入管道 + p.OutputCh <- bl + return + } + var status int + base, ok := p.baselines[bl.Status] + if ok { + // 挑选对应状态码的baseline进行compare + if status = base.Compare(bl); status == 1 { + p.PutToInvalid(bl, "compare failed") + return } - var status int - base, ok := p.baselines[bl.Status] - if ok { - // 挑选对应状态码的baseline进行compare - if status = base.Compare(bl); status == 1 { - p.PutToInvalid(bl, "compare failed") - continue - } - } - - bl.Collect() - for _, f := range bl.Frameworks { - if f.Tag == "waf/cdn" { - p.PutToInvalid(bl, "waf") - continue Loop - } - } - - if status == 0 && ok && base.FuzzyCompare(bl) { - p.PutToInvalid(bl, "fuzzy compare failed") - p.PutToFuzzy(bl) - continue - } - - p.outputCh <- bl } - p.analyzeDone = true + bl.Collect() + for _, f := range bl.Frameworks { + if f.Tag == "waf/cdn" { + p.PutToInvalid(bl, "waf") + return + } + } + + if status == 0 && ok && base.FuzzyCompare(bl) { + p.PutToInvalid(bl, "fuzzy compare failed") + p.PutToFuzzy(bl) + return + } + + p.OutputCh <- bl } func (p *Pool) addFuzzyBaseline(bl *pkg.Baseline) { @@ -308,12 +314,12 @@ func (p *Pool) addFuzzyBaseline(bl *pkg.Baseline) { func (p *Pool) PutToInvalid(bl *pkg.Baseline, reason string) { bl.IsValid = false bl.Reason = reason - p.outputCh <- bl + p.OutputCh <- bl } func (p *Pool) PutToFuzzy(bl *pkg.Baseline) { bl.IsFuzzy = true - p.fuzzyCh <- bl + p.FuzzyCh <- bl } func (p *Pool) resetFailed() { diff --git a/pkg/baseline.go b/pkg/baseline.go index a32a301..8c72ffb 100644 --- a/pkg/baseline.go +++ b/pkg/baseline.go @@ -4,7 +4,6 @@ import ( "encoding/json" "github.com/chainreactors/parsers" "github.com/chainreactors/spray/pkg/ihttp" - "math" "net/url" "strconv" "strings" @@ -101,7 +100,7 @@ func (bl *Baseline) Compare(other *Baseline) int { return 1 } - if math.Abs(float64(bl.BodyLength-other.BodyLength)) < 16 { + if i := bl.BodyLength - other.BodyLength; i < 16 || i > -16 { // 如果body length相等且md5相等, 则说明是同一个页面 if bl.BodyMd5 == parsers.Md5Hash(other.Body) { // 如果length相等, md5也相等, 则判断为全同 @@ -121,10 +120,9 @@ func (bl *Baseline) Compare(other *Baseline) int { return -1 } -var Distance uint8 +var Distance uint8 = 5 func (bl *Baseline) FuzzyCompare(other *Baseline) bool { - // todo 模糊匹配 if parsers.SimhashCompare(other.BodySimhash, bl.BodySimhash) < Distance { return true } @@ -241,3 +239,9 @@ func (bl *Baseline) Jsonify() string { } return string(bs) } + +func (bl *Baseline) ToMap() map[string]interface{} { + return map[string]interface{}{ + "status": bl.Status, + } +} diff --git a/pkg/config.go b/pkg/config.go index 818ec69..52238d9 100644 --- a/pkg/config.go +++ b/pkg/config.go @@ -19,17 +19,15 @@ var ModMap = map[string]SprayMod{ } type Config struct { - BaseURL string - Wordlist []string - Thread int - Timeout int - BaseReq *http.Request - Method string - Mod SprayMod - Headers http.Header - EnableFuzzy bool - ClientType int - Fns []func(string) string - OutputCh chan *Baseline - FuzzyCh chan *Baseline + BaseURL string + Wordlist []string + Thread int + Timeout int + Method string + Mod SprayMod + Headers http.Header + ClientType int + Fns []func(string) string + OutputCh chan *Baseline + FuzzyCh chan *Baseline }