From 70b9e25cf40b547ee5d94ee1d9ccff5295727697 Mon Sep 17 00:00:00 2001
From: M09Ic <m09ician@gmail.com>
Date: Wed, 9 Nov 2022 16:05:17 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BC=98=E5=8C=96compare=E9=80=BB=E8=BE=91?=
 =?UTF-8?q?=E4=B8=8E=E6=97=A5=E5=BF=97?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 go.mod               |  2 ++
 go.sum               |  2 ++
 internal/baseline.go | 37 ++++++++++++++++++++++++++++---------
 internal/pool.go     |  9 +++++----
 4 files changed, 37 insertions(+), 13 deletions(-)

diff --git a/go.mod b/go.mod
index 5b5eda3..42b4ca1 100644
--- a/go.mod
+++ b/go.mod
@@ -30,5 +30,7 @@ require (
 	github.com/twmb/murmur3 v1.1.6 // indirect
 	github.com/valyala/bytebufferpool v1.0.0 // indirect
 	golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab // indirect
+	gopkg.in/yaml.v2 v2.4.0 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
+	sigs.k8s.io/yaml v1.3.0 // indirect
 )
diff --git a/go.sum b/go.sum
index c64dcc1..339fe8f 100644
--- a/go.sum
+++ b/go.sum
@@ -75,9 +75,11 @@ golang.org/x/text v0.3.6/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
 golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ=
 golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/yaml.v2 v2.4.0 h1:D8xgwECY7CYvx+Y2n4sBz93Jn9JRvxdiyyo8CTfuKaY=
 gopkg.in/yaml.v2 v2.4.0/go.mod h1:RDklbk79AGWmwhnvt/jBztapEOGDOx6ZbXqjP6csGnQ=
 gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
 gopkg.in/yaml.v3 v3.0.0-20210107192922-496545a6307b/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
 gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
+sigs.k8s.io/yaml v1.3.0 h1:a2VclLzOGrwOHDiV8EfBGhvjHvP46CtW5j6POvhYGGo=
 sigs.k8s.io/yaml v1.3.0/go.mod h1:GeOyir5tyXNByN85N/dRIT9es5UQNerPYEKK56eTBm8=
diff --git a/internal/baseline.go b/internal/baseline.go
index 420fc1a..dd69d89 100644
--- a/internal/baseline.go
+++ b/internal/baseline.go
@@ -5,6 +5,7 @@ import (
 	"github.com/chainreactors/parsers"
 	"github.com/chainreactors/spray/pkg"
 	"github.com/chainreactors/spray/pkg/ihttp"
+	"net/url"
 	"strconv"
 	"strings"
 )
@@ -15,7 +16,10 @@ func NewBaseline(u, host string, resp *ihttp.Response) *baseline {
 		Status:  resp.StatusCode(),
 		IsValid: true,
 	}
-
+	uu, err := url.Parse(u)
+	if err == nil {
+		bl.Path = uu.Path
+	}
 	if resp.ClientType == ihttp.STANDARD {
 		bl.Host = host
 	}
@@ -36,6 +40,11 @@ func NewInvalidBaseline(u, host string, resp *ihttp.Response) *baseline {
 		IsValid: false,
 	}
 
+	uu, err := url.Parse(u)
+	if err == nil {
+		bl.Path = uu.Path
+	}
+
 	if resp.ClientType == ihttp.STANDARD {
 		bl.Host = host
 	}
@@ -47,6 +56,7 @@ func NewInvalidBaseline(u, host string, resp *ihttp.Response) *baseline {
 
 type baseline struct {
 	Url          string         `json:"url"`
+	Path         string         `json:"path"`
 	Host         string         `json:"host"`
 	Body         []byte         `json:"-"`
 	BodyLength   int            `json:"body_length"`
@@ -75,21 +85,30 @@ func (bl *baseline) Collect() {
 	bl.Frameworks = pkg.FingerDetect(string(bl.Raw))
 }
 
-// Equal if this equal other return true
-func (bl *baseline) Equal(other *baseline) bool {
+// Compare if this equal other return true
+func (bl *baseline) Compare(other *baseline) int {
 	if other.RedirectURL != "" && bl.RedirectURL == other.RedirectURL {
-		// 如果重定向url不为空, 且与bl不相同, 则说明不是同一个页面
-		return true
+		// 如果重定向url不为空, 且与base不相同, 则说明不是同一个页面
+		return 1
 	}
 
 	if bl.BodyLength == other.BodyLength {
 		// 如果body length相等且md5相等, 则说明是同一个页面
-		if bl.BodyMd5 == parsers.Md5Hash(other.Raw) {
-			return true
+		if bl.BodyMd5 == parsers.Md5Hash(other.Body) {
+			// 如果length相等, md5也相等, 则判断为全同
+			return 1
+		} else {
+			// 如果长度相等, 但是md5不相等, 可能是存在csrftoken之类的随机值
+			return 0
+		}
+	} else {
+		if strings.Contains(string(other.Body), other.Path) {
+			return 0
+		} else {
+			return -1
 		}
 	}
-
-	return false
+	return -1
 }
 
 func (bl *baseline) FuzzyEqual(other *baseline) bool {
diff --git a/internal/pool.go b/internal/pool.go
index c307365..533e00d 100644
--- a/internal/pool.go
+++ b/internal/pool.go
@@ -87,7 +87,7 @@ func NewPool(ctx context.Context, config *pkg.Config, outputCh chan *baseline) (
 			bl = &baseline{Url: pool.BaseURL + unit.path, Err: reqerr}
 		} else {
 			pool.failedCount = 0 // 如果后续访问正常, 重置错误次数
-			if err = pool.PreCompare(resp); err == nil || unit.source == CheckSource {
+			if err = pool.PreCompare(resp); err == nil || unit.source == CheckSource || unit.source == InitSource {
 				// 通过预对比跳过一些无用数据, 减少性能消耗
 				bl = NewBaseline(req.URI(), req.Host(), resp)
 			} else {
@@ -99,12 +99,13 @@ func NewPool(ctx context.Context, config *pkg.Config, outputCh chan *baseline) (
 		case InitSource:
 			pool.base = bl
 			pool.initwg.Done()
+			logs.Log.Important("[init] " + bl.String())
 			return
 		case CheckSource:
-			logs.Log.Debugf("check: " + bl.String())
+			logs.Log.Debugf("[check] " + bl.String())
 			if bl.Err != nil {
 				logs.Log.Warnf("maybe ip has banned by waf, break (%d/%d), error: %s", pool.failedCount, breakThreshold, bl.Err.Error())
-			} else if !pool.base.Equal(bl) {
+			} else if pool.base.Compare(bl) < 1 {
 				logs.Log.Warn("maybe trigger risk control")
 			}
 
@@ -220,7 +221,7 @@ func (p *Pool) PreCompare(resp *ihttp.Response) error {
 
 func (p *Pool) comparing() {
 	for bl := range p.tempCh {
-		if p.base.Equal(bl) {
+		if p.base.Compare(bl) == 1 {
 			// 如果是同一个包则设置为无效包
 			bl.IsValid = false
 			p.outputCh <- bl