实装wafcheck, 目前只对状态码做简单的判断.

新增--black-status, 自定义黑名单状态码

internal/option.go

@@ -60,6 +60,7 @@ type ModeOptions struct {
 	CheckPeriod    int    `long:"check-period" default:"100"`
 	ErrPeriod      int    `long:"error-period" default:"10"`
 	BreakThreshold int    `long:"error-threshold" default:"20"`
+	BlackStatus    string `long:"black-status" default:"default"`
 }
 
 type MiscOptions struct {
@@ -124,6 +125,18 @@ func (opt *Option) PrepareRunner() (*Runner, error) {
 		r.ErrPeriod = max
 	}
 
+	if opt.BlackStatus != "default" {
+		for _, s := range strings.Split(opt.BlackStatus, ",") {
+			si, err := strconv.Atoi(s)
+			if err != nil {
+				return nil, err
+			}
+			BlackStatus = append(BlackStatus, si)
+		}
+	} else {
+		BlackStatus = []int{400, 404, 410}
+	}
+
 	// prepare url
 	var urls []string
 	var file *os.File

internal/pool.go

@@ -15,10 +15,19 @@ import (
 )
 
 var (
-	CheckStatusCode func(int) bool
+	CheckBadStatus func(int) bool
 	CheckRedirect  func(string) bool
-	CheckWaf        func([]byte) bool
 )
+
+func CheckWaf(status int) bool {
+	for _, s := range WAFStatus {
+		if status == s {
+			return true
+		}
+	}
+	return false
+}
+
 var max = 2147483647
 
 func NewPool(ctx context.Context, config *pkg.Config) (*Pool, error) {
@@ -202,7 +211,7 @@ func (p *Pool) Init() error {
 	p.index.Collect()
 
 	logs.Log.Important("[baseline.random] " + p.base.String())
-	logs.Log.Important("[baseline.index] " + p.base.String())
+	logs.Log.Important("[baseline.index] " + p.index.String())
 
 	if p.base.RedirectURL != "" {
 		CheckRedirect = func(redirectURL string) bool {
@@ -256,23 +265,23 @@ Loop:
 }
 
 func (p *Pool) PreCompare(resp *ihttp.Response) error {
-	if p.base != nil && p.base.Status != 200 && p.base.Status == resp.StatusCode() {
+	status := resp.StatusCode()
+	if p.base != nil && p.base.Status != 200 && p.base.Status == status {
 		return ErrSameStatus
 	}
 
-	if !CheckStatusCode(resp.StatusCode()) {
+	if CheckBadStatus(status) {
 		return ErrBadStatus
 	}
 
+	if CheckWaf(status) {
+		return ErrWaf
+	}
+
 	if CheckRedirect != nil && !CheckRedirect(string(resp.GetHeader("Location"))) {
 		return ErrRedirect
 	}
 
-	if CheckWaf != nil && !CheckWaf(nil) {
-		// todo check waf
-		return ErrWaf
-	}
-
 	return nil
 }
 

internal/runner.go

@@ -13,8 +13,11 @@ import (
 	"time"
 )
 
-var BlackStatus = []int{400, 404, 410}
-var FuzzyStatus = []int{403, 500, 501, 502, 503}
+var (
+	BlackStatus = []int{}
+	FuzzyStatus = []int{403, 500, 501, 502, 503}
+	WAFStatus   = []int{493, 418}
+)
 
 type Runner struct {
 	URLList        chan string
@@ -46,14 +49,14 @@ type Runner struct {
 
 func (r *Runner) Prepare(ctx context.Context) error {
 	var err error
-	CheckStatusCode = func(status int) bool {
+	CheckBadStatus = func(status int) bool {
 		for _, black := range BlackStatus {
 			if black == status {
-				return false
-			}
-		}
 				return true
 			}
+		}
+		return false
+	}
 
 	r.Pools, err = ants.NewPoolWithFunc(r.PoolSize, func(i interface{}) {
 		u := i.(string)