From d95b0315ecda2a03ab140117122fc29f3c1c15a9 Mon Sep 17 00:00:00 2001
From: M09Ic <m09ician@gmail.com>
Date: Thu, 17 Nov 2022 16:27:44 +0800
Subject: [PATCH] =?UTF-8?q?=E5=AE=9E=E8=A3=85wafcheck,=20=E7=9B=AE?=
 =?UTF-8?q?=E5=89=8D=E5=8F=AA=E5=AF=B9=E7=8A=B6=E6=80=81=E7=A0=81=E5=81=9A?=
 =?UTF-8?q?=E7=AE=80=E5=8D=95=E7=9A=84=E5=88=A4=E6=96=AD.=20=E6=96=B0?=
 =?UTF-8?q?=E5=A2=9E--black-status,=20=E8=87=AA=E5=AE=9A=E4=B9=89=E9=BB=91?=
 =?UTF-8?q?=E5=90=8D=E5=8D=95=E7=8A=B6=E6=80=81=E7=A0=81?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 internal/option.go | 23 ++++++++++++++++++-----
 internal/pool.go   | 31 ++++++++++++++++++++-----------
 internal/runner.go | 13 ++++++++-----
 3 files changed, 46 insertions(+), 21 deletions(-)

diff --git a/internal/option.go b/internal/option.go
index 0300b12..7323295 100644
--- a/internal/option.go
+++ b/internal/option.go
@@ -55,11 +55,12 @@ type RequestOptions struct {
 }
 
 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"`
+	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"`
+	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
diff --git a/internal/pool.go b/internal/pool.go
index 78f0a09..d403f97 100644
--- a/internal/pool.go
+++ b/internal/pool.go
@@ -15,10 +15,19 @@ import (
 )
 
 var (
-	CheckStatusCode func(int) bool
-	CheckRedirect   func(string) bool
-	CheckWaf        func([]byte) bool
+	CheckBadStatus func(int) bool
+	CheckRedirect  func(string) 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
 }
 
diff --git a/internal/runner.go b/internal/runner.go
index 4f201a2..7ed61bb 100644
--- a/internal/runner.go
+++ b/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,13 +49,13 @@ 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 true
+		return false
 	}
 
 	r.Pools, err = ants.NewPoolWithFunc(r.PoolSize, func(i interface{}) {