优化--force的逻辑, 开启--force时将关闭check

internal/option.go

@@ -16,6 +16,7 @@ type Option struct {
 	InputOptions
 	OutputOptions
 	RequestOptions
+	ModeOptions
 	MiscOptions
 }
 
@@ -50,9 +51,15 @@ type RequestOptions struct {
 	Headers         []string `long:"header"`
 	Method          string   `long:"method"`
 	Cookie          string   `long:"cookie"`
-	Force           bool     `long:"force"`
+	SimhashDistance int      `long:"distance" default:"5"`
-	SimhashDistance int      `long:"distance"`
+}
-	CheckOnly       bool     `long:"--check-only"`
+
+type ModeOptions struct {
+	Force          bool `long:"force"`
+	CheckOnly      bool `long:"check-only"`
+	CheckPeriod    int  `long:"check-period" default:"100"`
+	ErrPeriod      int  `long:"error-period" default:"10"`
+	BreakThreshold int  `long:"error-threshold" default:"20"`
 }
 
 type MiscOptions struct {
@@ -73,20 +80,23 @@ func (opt *Option) PrepareRunner() (*Runner, error) {
 	}
 	var err error
 	r := &Runner{
-		Progress:  uiprogress.New(),
+		Progress:       uiprogress.New(),
-		Threads:   opt.Threads,
+		Threads:        opt.Threads,
-		PoolSize:  opt.PoolSize,
+		PoolSize:       opt.PoolSize,
-		Mod:       opt.Mod,
+		Mod:            opt.Mod,
-		Timeout:   opt.Timeout,
+		Timeout:        opt.Timeout,
-		Deadline:  opt.Deadline,
+		Deadline:       opt.Deadline,
-		Offset:    opt.Offset,
+		Offset:         opt.Offset,
-		Limit:     opt.Limit,
+		Limit:          opt.Limit,
-		URLList:   make(chan string),
+		URLList:        make(chan string),
-		OutputCh:  make(chan *pkg.Baseline, 100),
+		OutputCh:       make(chan *pkg.Baseline, 100),
-		FuzzyCh:   make(chan *pkg.Baseline, 100),
+		FuzzyCh:        make(chan *pkg.Baseline, 100),
-		Fuzzy:     opt.Fuzzy,
+		Fuzzy:          opt.Fuzzy,
-		Force:     opt.Force,
+		Force:          opt.Force,
-		CheckOnly: opt.CheckOnly,
+		CheckOnly:      opt.CheckOnly,
+		CheckPeriod:    opt.CheckPeriod,
+		ErrPeriod:      opt.ErrPeriod,
+		BreakThreshold: opt.BreakThreshold,
 	}
 
 	err = pkg.LoadTemplates()
@@ -108,7 +118,10 @@ func (opt *Option) PrepareRunner() (*Runner, error) {
 	}
 
 	if opt.Force {
-		breakThreshold = 999999
+		// 如果开启了force模式, 将关闭check机制, err积累到一定数量自动退出机制
+		r.BreakThreshold = max
+		r.CheckPeriod = max
+		r.ErrPeriod = max
 	}
 
 	// prepare url

internal/pool.go

@@ -19,8 +19,7 @@ var (
 	CheckRedirect   func(string) bool
 	CheckWaf        func([]byte) bool
 )
-
+var max = 2147483647
-var breakThreshold int = 20
 
 func NewPool(ctx context.Context, config *pkg.Config) (*Pool, error) {
 	pctx, cancel := context.WithCancel(ctx)
@@ -34,8 +33,6 @@ func NewPool(ctx context.Context, config *pkg.Config) (*Pool, error) {
 		tempCh:      make(chan *pkg.Baseline, config.Thread),
 		wg:          sync.WaitGroup{},
 		initwg:      sync.WaitGroup{},
-		checkPeriod: 100,
-		errPeriod:   10,
 		reqCount:    1,
 		failedCount: 1,
 	}
@@ -49,7 +46,7 @@ func NewPool(ctx context.Context, config *pkg.Config) (*Pool, error) {
 			pool.wg.Add(1)
 			_ = pool.pool.Invoke(newUnit(pkg.RandPath(), CheckSource))
 
-			if pool.failedCount > breakThreshold {
+			if pool.failedCount > pool.BreakThreshold {
 				// 当报错次数超过上限是, 结束任务
 				pool.recover()
 				pool.cancel()
@@ -64,7 +61,7 @@ func NewPool(ctx context.Context, config *pkg.Config) (*Pool, error) {
 			pool.wg.Add(1)
 			_ = pool.pool.Invoke(newUnit(pkg.RandHost(), CheckSource))
 
-			if pool.failedCount > breakThreshold {
+			if pool.failedCount > pool.BreakThreshold {
 				// 当报错次数超过上限是, 结束任务
 				pool.recover()
 				pool.cancel()
@@ -109,7 +106,7 @@ func NewPool(ctx context.Context, config *pkg.Config) (*Pool, error) {
 			return
 		case CheckSource:
 			if bl.Err != "" {
-				logs.Log.Warnf("[check.error] maybe ip had banned by waf, break (%d/%d), error: %s", pool.failedCount, breakThreshold, bl.Err)
+				logs.Log.Warnf("[check.error] maybe ip had banned by waf, break (%d/%d), error: %s", pool.failedCount, pool.BreakThreshold, bl.Err)
 				pool.failedBaselines = append(pool.failedBaselines, bl)
 			} else if i := pool.base.Compare(bl); i < 1 {
 				if i == 0 {
@@ -128,10 +125,10 @@ func NewPool(ctx context.Context, config *pkg.Config) (*Pool, error) {
 			// 异步进行性能消耗较大的深度对比
 			pool.tempCh <- bl
 			pool.reqCount++
-			if pool.reqCount%pool.checkPeriod == 0 {
+			if pool.reqCount%pool.CheckPeriod == 0 {
 				pool.reqCount++
 				go pool.check()
-			} else if pool.failedCount%pool.errPeriod == 0 {
+			} else if pool.failedCount%pool.ErrPeriod == 0 {
 				pool.failedCount++
 				go pool.check()
 			}
@@ -168,8 +165,6 @@ type Pool struct {
 	tempCh          chan *pkg.Baseline // 待处理的baseline
 	reqCount        int
 	failedCount     int
-	checkPeriod     int
-	errPeriod       int
 	failedBaselines []*pkg.Baseline
 	base            *pkg.Baseline
 	baselines       map[int]*pkg.Baseline

internal/runner.go

@@ -17,28 +17,31 @@ var BlackStatus = []int{400, 404, 410}
 var FuzzyStatus = []int{403, 500, 501, 502, 503}
 
 type Runner struct {
-	URLList    chan string
+	URLList        chan string
-	Wordlist   []string
+	Wordlist       []string
-	Headers    http.Header
+	Headers        http.Header
-	Fns        []func(string) string
+	Fns            []func(string) string
-	Threads    int
+	Threads        int
-	PoolSize   int
+	PoolSize       int
-	Pools      *ants.PoolWithFunc
+	Pools          *ants.PoolWithFunc
-	poolwg     sync.WaitGroup
+	poolwg         sync.WaitGroup
-	Timeout    int
+	Timeout        int
-	Mod        string
+	Mod            string
-	Probes     []string
+	Probes         []string
-	OutputCh   chan *pkg.Baseline
+	OutputCh       chan *pkg.Baseline
-	FuzzyCh    chan *pkg.Baseline
+	FuzzyCh        chan *pkg.Baseline
-	Fuzzy      bool
+	Fuzzy          bool
-	OutputFile *files.File
+	OutputFile     *files.File
-	FuzzyFile  *files.File
+	FuzzyFile      *files.File
-	Force      bool
+	Force          bool
-	Progress   *uiprogress.Progress
+	Progress       *uiprogress.Progress
-	Offset     int
+	Offset         int
-	Limit      int
+	Limit          int
-	Deadline   int
+	Deadline       int
-	CheckOnly  bool
+	CheckPeriod    int
+	ErrPeriod      int
+	BreakThreshold int
+	CheckOnly      bool
 }
 
 func (r *Runner) Prepare(ctx context.Context) error {
@@ -55,15 +58,18 @@ func (r *Runner) Prepare(ctx context.Context) error {
 	r.Pools, err = ants.NewPoolWithFunc(r.PoolSize, func(i interface{}) {
 		u := i.(string)
 		config := &pkg.Config{
-			BaseURL:  u,
+			BaseURL:        u,
-			Wordlist: r.Wordlist,
+			Wordlist:       r.Wordlist,
-			Thread:   r.Threads,
+			Thread:         r.Threads,
-			Timeout:  r.Timeout,
+			Timeout:        r.Timeout,
-			Headers:  r.Headers,
+			Headers:        r.Headers,
-			Mod:      pkg.ModMap[r.Mod],
+			Mod:            pkg.ModMap[r.Mod],
-			Fns:      r.Fns,
+			Fns:            r.Fns,
-			OutputCh: r.OutputCh,
+			OutputCh:       r.OutputCh,
-			FuzzyCh:  r.FuzzyCh,
+			FuzzyCh:        r.FuzzyCh,
+			CheckPeriod:    r.CheckPeriod,
+			ErrPeriod:      r.ErrPeriod,
+			BreakThreshold: r.BreakThreshold,
 		}
 
 		if config.Mod == pkg.PathSpray {
@@ -84,6 +90,7 @@ func (r *Runner) Prepare(ctx context.Context) error {
 		if err != nil {
 			logs.Log.Error(err.Error())
 			if !r.Force {
+				// 如果没开启force, init失败将会关闭pool
 				pool.cancel()
 				r.poolwg.Done()
 				return

pkg/config.go

@@ -19,15 +19,18 @@ var ModMap = map[string]SprayMod{
 }
 
 type Config struct {
-	BaseURL    string
+	BaseURL        string
-	Wordlist   []string
+	Wordlist       []string
-	Thread     int
+	Thread         int
-	Timeout    int
+	Timeout        int
-	Method     string
+	CheckPeriod    int
-	Mod        SprayMod
+	ErrPeriod      int
-	Headers    http.Header
+	BreakThreshold int
-	ClientType int
+	Method         string
-	Fns        []func(string) string
+	Mod            SprayMod
-	OutputCh   chan *Baseline
+	Headers        http.Header
-	FuzzyCh    chan *Baseline
+	ClientType     int
+	Fns            []func(string) string
+	OutputCh       chan *Baseline
+	FuzzyCh        chan *Baseline
 }