diff --git a/internal/option.go b/internal/option.go index 3ec5191..fe1d897 100644 --- a/internal/option.go +++ b/internal/option.go @@ -271,7 +271,7 @@ func (opt *Option) PrepareRunner() (*Runner, error) { } taskfrom = "resume " + opt.ResumeFrom for _, stat := range stats { - tasks = append(tasks, &Task{baseUrl: stat.BaseUrl, offset: stat.Offset + stat.ReqNumber, total: r.Total}) + tasks = append(tasks, &Task{baseUrl: stat.BaseUrl, offset: stat.Offset + stat.End, total: r.Total}) } } else { var file *os.File diff --git a/internal/pool.go b/internal/pool.go index f9a080d..ef735b2 100644 --- a/internal/pool.go +++ b/internal/pool.go @@ -15,6 +15,7 @@ import ( "strconv" "strings" "sync" + "sync/atomic" "time" ) @@ -37,6 +38,7 @@ func NewPool(ctx context.Context, config *pkg.Config) (*Pool, error) { worder: words.NewWorder(config.Wordlist, config.Fns), baselines: make(map[int]*pkg.Baseline), tempCh: make(chan *pkg.Baseline, config.Thread), + checkCh: make(chan *Unit), wg: sync.WaitGroup{}, initwg: sync.WaitGroup{}, reqCount: 1, @@ -45,9 +47,8 @@ func NewPool(ctx context.Context, config *pkg.Config) (*Pool, error) { pool.worder.Rules = pool.Rules pool.worder.RunWithRules() - p, _ := ants.NewPoolWithFunc(config.Thread, func(i interface{}) { - pool.Statistor.Total++ + atomic.AddInt32(&pool.Statistor.ReqTotal, 1) unit := i.(*Unit) req, err := pool.genReq(unit.path) if err != nil { @@ -65,7 +66,7 @@ func NewPool(ctx context.Context, config *pkg.Config) (*Pool, error) { var bl *pkg.Baseline if reqerr != nil && reqerr != fasthttp.ErrBodyTooLarge { pool.failedCount++ - pool.Statistor.FailedNumber++ + atomic.AddInt32(&pool.Statistor.FailedNumber, 1) bl = &pkg.Baseline{UrlString: pool.BaseURL + unit.path, IsValid: false, ErrString: reqerr.Error(), Reason: ErrRequestFailed.Error()} pool.failedBaselines = append(pool.failedBaselines, bl) } else { @@ -313,7 +314,7 @@ Loop: break Loop } pool.Statistor.End++ - if pool.reqCount < offset { + if int(pool.reqCount) < offset { pool.reqCount++ continue } @@ -340,7 +341,6 @@ Loop: } } pool.wg.Wait() - pool.Statistor.ReqNumber = pool.reqCount pool.Statistor.EndTime = time.Now().Unix() pool.Close() } diff --git a/internal/runner.go b/internal/runner.go index 364a8ab..0e0b25a 100644 --- a/internal/runner.go +++ b/internal/runner.go @@ -136,7 +136,7 @@ func (r *Runner) Prepare(ctx context.Context) error { r.Done() return } - + pool.Statistor.Total = r.Total pool.bar = pkg.NewBar(config.BaseURL, t.total-t.offset, r.Progress) err = pool.Init() if err != nil { diff --git a/pkg/statistor.go b/pkg/statistor.go index fc8b39e..e484774 100644 --- a/pkg/statistor.go +++ b/pkg/statistor.go @@ -23,26 +23,27 @@ func NewStatistor(url string) *Statistor { type Statistor struct { BaseUrl string `json:"url"` Counts map[int]int `json:"counts"` - ReqNumber int `json:"req"` - FailedNumber int `json:"failed"` + FailedNumber int32 `json:"failed"` + ReqTotal int32 `json:"req_total"` CheckNumber int `json:"check"` FoundNumber int `json:"found"` FilteredNumber int `json:"filtered"` FuzzyNumber int `json:"fuzzy"` WafedNumber int `json:"wafed"` - End int `json:"end"` - Offset int `json:"offset"` - Total int `json:"total"` - StartTime int64 `json:"start_time"` - EndTime int64 `json:"end_time"` - WordCount int `json:"word_count"` - Word string `json:"word"` - Dictionaries []string `json:"dictionaries"` + + End int `json:"end"` + Offset int `json:"offset"` + Total int `json:"total"` + StartTime int64 `json:"start_time"` + EndTime int64 `json:"end_time"` + WordCount int `json:"word_count"` + Word string `json:"word"` + Dictionaries []string `json:"dictionaries"` } func (stat *Statistor) String() string { var s strings.Builder - s.WriteString(fmt.Sprintf("[stat] %s took %d s, request total: %d, found: %d, check: %d, failed: %d", stat.BaseUrl, stat.EndTime-stat.StartTime, stat.ReqNumber, stat.FoundNumber, stat.CheckNumber, stat.FailedNumber)) + s.WriteString(fmt.Sprintf("[stat] %s took %d s, request total: %d, finish: %d/%d, found: %d, check: %d, failed: %d", stat.BaseUrl, stat.EndTime-stat.StartTime, stat.ReqTotal, stat.End, stat.Total, stat.FoundNumber, stat.CheckNumber, stat.FailedNumber)) if stat.FuzzyNumber != 0 { s.WriteString(", fuzzy: " + strconv.Itoa(stat.FuzzyNumber))