diff --git a/internal/baseline.go b/internal/baseline.go index d382ea4..86aabf4 100644 --- a/internal/baseline.go +++ b/internal/baseline.go @@ -49,6 +49,8 @@ func NewInvalidBaseline(u, host string, resp *ihttp.Response) *baseline { bl.Host = host } + bl.Body = resp.Body() + bl.BodyLength = resp.ContentLength() bl.RedirectURL = string(resp.GetHeader("Location")) return bl diff --git a/internal/option.go b/internal/option.go index 33cecd3..f515b22 100644 --- a/internal/option.go +++ b/internal/option.go @@ -48,6 +48,7 @@ type RequestOptions struct { Headers []string `long:"header"` Method string `long:"method"` Cookie string `long:"cookie"` + Force bool `long:"force"` } type MiscOptions struct { @@ -73,7 +74,6 @@ func (opt *Option) PrepareRunner() (*Runner, error) { PoolSize: opt.PoolSize, Mod: opt.Mod, Timeout: opt.Timeout, - Probes: strings.Split(opt.OutputProbe, ","), Deadline: opt.Deadline, Offset: opt.Offset, Limit: opt.Limit, @@ -94,6 +94,10 @@ func (opt *Option) PrepareRunner() (*Runner, error) { logs.Log.Writer = r.Progress.Bypass() } + if opt.Force { + breakThreshold = 999999 + } + // prepare url var urls []string var file *os.File @@ -217,6 +221,9 @@ func (opt *Option) PrepareRunner() (*Runner, error) { } } + if opt.OutputProbe != "" { + r.Probes = strings.Split(opt.OutputProbe, ",") + } return r, nil } diff --git a/internal/pool.go b/internal/pool.go index 96b5ba2..5c72dd4 100644 --- a/internal/pool.go +++ b/internal/pool.go @@ -19,7 +19,7 @@ var ( CheckWaf func([]byte) bool ) -var breakThreshold int = 10 +var breakThreshold int = 20 func NewPool(ctx context.Context, config *pkg.Config, outputCh chan *baseline) (*Pool, error) { pctx, cancel := context.WithCancel(ctx) @@ -35,6 +35,8 @@ func NewPool(ctx context.Context, config *pkg.Config, outputCh chan *baseline) ( initwg: sync.WaitGroup{}, checkPeriod: 100, errPeriod: 10, + reqCount: 1, + failedCount: 1, } switch config.Mod { @@ -87,6 +89,7 @@ func NewPool(ctx context.Context, config *pkg.Config, outputCh chan *baseline) ( if reqerr != nil && reqerr != fasthttp.ErrBodyTooLarge { pool.failedCount++ bl = &baseline{Url: pool.BaseURL + unit.path, Err: reqerr} + pool.failedBaselines = append(pool.failedBaselines, bl) } else { if err = pool.PreCompare(resp); err == nil || unit.source == CheckSource || unit.source == InitSource { // 通过预对比跳过一些无用数据, 减少性能消耗 @@ -271,8 +274,8 @@ func (p *Pool) ResetFailed() { func (p *Pool) Recover() { logs.Log.Errorf("failed request exceeds the threshold , task will exit. Breakpoint %d", p.reqCount) logs.Log.Error("collecting failed check") - for _, bl := range p.failedBaselines { - logs.Log.Error(bl.String()) + for i, bl := range p.failedBaselines { + logs.Log.Errorf("[failed.%d] %s", i, bl.String()) } } diff --git a/internal/runner.go b/internal/runner.go index f99a4de..73e1bfd 100644 --- a/internal/runner.go +++ b/internal/runner.go @@ -113,15 +113,22 @@ Loop: } func (r *Runner) Outputting() { + var outFunc func(baseline2 *baseline) + if len(r.Probes) > 0 { + outFunc = func(bl *baseline) { + logs.Log.Console("[+] " + bl.Format(r.Probes) + "\n") + } + } else { + outFunc = func(bl *baseline) { + logs.Log.Console("[+] " + bl.String() + "\n") + } + } + for { select { case bl := <-r.OutputCh: if bl.IsValid { - if len(r.Probes) > 0 { - logs.Log.Console("[+]" + bl.Format(r.Probes) + "\n") - } else { - logs.Log.Console("[+] " + bl.String() + "\n") - } + outFunc(bl) } else { logs.Log.Debug(bl.String()) }