baseline add from and parent prop

chainreactors · Oct 13, 2024 · e483bb4 · e483bb4
1 parent 344e560
commit e483bb4
Show file tree

Hide file tree

Showing 6 changed files with 53 additions and 23 deletions.
diff --git a/go.mod b/go.mod
@@ -6,7 +6,7 @@ require (
 	github.com/chainreactors/files v0.0.0-20240716182835-7884ee1e77f0
 	github.com/chainreactors/fingers v0.0.0-20240716172449-2fc3147b9c2a
 	github.com/chainreactors/logs v0.0.0-20240207121836-c946f072f81f
-	github.com/chainreactors/parsers v0.0.0-20240910081704-fd57f462fc65
+	github.com/chainreactors/parsers v0.0.0-20241013180542-88e2dc355c57
 	github.com/chainreactors/utils v0.0.0-20240805193040-ff3b97aa3c3f
 	github.com/expr-lang/expr v1.16.9
 	github.com/gookit/config/v2 v2.2.5

diff --git a/go.sum b/go.sum
@@ -101,6 +101,8 @@ github.com/chainreactors/parsers v0.0.0-20240829055950-923f89a92b84 h1:F6umsdHLx
 github.com/chainreactors/parsers v0.0.0-20240829055950-923f89a92b84/go.mod h1:7rXdYz6jrdjF0WUH1ICcAXKIKKjKmJo2PU8u43V7jkA=
 github.com/chainreactors/parsers v0.0.0-20240910081704-fd57f462fc65 h1:subSvyczsErYMRnCD07s4Ub6zOSaw2xZ1/O9t3tHkuw=
 github.com/chainreactors/parsers v0.0.0-20240910081704-fd57f462fc65/go.mod h1:7rXdYz6jrdjF0WUH1ICcAXKIKKjKmJo2PU8u43V7jkA=
+github.com/chainreactors/parsers v0.0.0-20241013180542-88e2dc355c57 h1:KuijtekTNtSpQbKf2jqKp99gxnGQXffPeEF+EOHnXBE=
+github.com/chainreactors/parsers v0.0.0-20241013180542-88e2dc355c57/go.mod h1:7rXdYz6jrdjF0WUH1ICcAXKIKKjKmJo2PU8u43V7jkA=
 github.com/chainreactors/utils v0.0.0-20240528085651-ba1b255482c1/go.mod h1:JA4eiQZm+7AsfjXBcIzIdVKBEhDCb16eNtWFCGTxlvs=
 github.com/chainreactors/utils v0.0.0-20240704062557-662d623b74f4/go.mod h1:JA4eiQZm+7AsfjXBcIzIdVKBEhDCb16eNtWFCGTxlvs=
 github.com/chainreactors/utils v0.0.0-20240715080349-d2d0484c95ed/go.mod h1:LajXuvESQwP+qCMAvlcoSXppQCjuLlBrnQpu9XQ1HtU=

diff --git a/internal/pool/brutepool.go b/internal/pool/brutepool.go
@@ -322,9 +322,7 @@ func (pool *BrutePool) Invoke(v interface{}) {
 	if !ihttp.CheckBodySize(int64(bl.BodyLength)) {
 		bl.ExceedLength = true
 	}
-	bl.Source = unit.source
-	bl.ReqDepth = unit.depth
-	bl.Number = unit.number
+	unit.Update(bl)
 	bl.Spended = time.Since(start).Milliseconds()
 	switch unit.source {
 	case parsers.InitRandomSource:
@@ -684,6 +682,26 @@ func (pool *BrutePool) doCheck() {
 	}
 }
 
+func (pool *BrutePool) doRedirect(bl *pkg.Baseline, depth int) {
+	if depth >= pool.MaxRedirect {
+		return
+	}
+	reURL := pkg.FormatURL(bl.Url.Path, bl.RedirectURL)
+	pool.wg.Add(1)
+	go func() {
+		defer pool.wg.Done()
+		pool.addAddition(&Unit{
+			path:     reURL,
+			parent:   bl.Number,
+			host:     bl.Host,
+			source:   parsers.RedirectSource,
+			from:     bl.Source,
+			frontUrl: bl.UrlString,
+			depth:    depth + 1,
+		})
+	}()
+}
+
 func (pool *BrutePool) doCrawl(bl *pkg.Baseline) {
 	if !pool.Crawl || bl.ReqDepth >= pool.MaxCrawlDepth {
 		return
@@ -705,8 +723,10 @@ func (pool *BrutePool) doCrawl(bl *pkg.Baseline) {
 			}
 			pool.addAddition(&Unit{
 				path:   u,
+				parent: bl.Number,
 				host:   bl.Host,
 				source: parsers.CrawlSource,
+				from:   bl.Source,
 				depth:  bl.ReqDepth + 1,
 			})
 		}
@@ -731,7 +751,13 @@ func (pool *BrutePool) doScopeCrawl(bl *pkg.Baseline) {
 				if _, ok := pool.scopeurls[u]; !ok {
 					pool.urls.Store(u, nil)
 					pool.wg.Add(1)
-					pool.scopePool.Invoke(&Unit{path: u, source: parsers.CrawlSource, depth: bl.ReqDepth + 1})
+					pool.scopePool.Invoke(&Unit{
+						path:   u,
+						parent: bl.Number,
+						source: parsers.CrawlSource,
+						from:   bl.Source,
+						depth:  bl.ReqDepth + 1,
+					})
 				}
 				pool.scopeLocker.Unlock()
 			}
@@ -776,8 +802,10 @@ func (pool *BrutePool) doAppendRule(bl *pkg.Baseline) {
 		for u := range rule.RunAsStream(pool.AppendRule.Expressions, path.Base(bl.Path)) {
 			pool.addAddition(&Unit{
 				path:   pkg.Dir(bl.Url.Path) + u,
+				parent: bl.Number,
 				host:   bl.Host,
 				source: parsers.AppendRuleSource,
+				from:   bl.Source,
 				depth:  bl.ReqDepth + 1,
 			})
 		}
@@ -797,8 +825,10 @@ func (pool *BrutePool) doAppendWords(bl *pkg.Baseline) {
 		for u := range NewBruteWords(pool.Config, pool.AppendWords).Output {
 			pool.addAddition(&Unit{
 				path:   pkg.SafePath(bl.Path, u),
+				parent: bl.Number,
 				host:   bl.Host,
 				source: parsers.AppendSource,
+				from:   bl.Source,
 				depth:  bl.RecuDepth + 1,
 			})
 		}

diff --git a/internal/pool/checkpool.go b/internal/pool/checkpool.go
@@ -199,9 +199,11 @@ func (pool *CheckPool) doRedirect(bl *pkg.Baseline, depth int) {
 	go func() {
 		pool.additionCh <- &Unit{
 			path:     reURL,
+			parent:   bl.Number,
 			source:   parsers.RedirectSource,
 			frontUrl: bl.UrlString,
 			depth:    depth + 1,
+			from:     bl.Source,
 		}
 	}()
 }
@@ -221,8 +223,10 @@ func (pool *CheckPool) doUpgrade(bl *pkg.Baseline) {
 	go func() {
 		pool.additionCh <- &Unit{
 			path:   reurl,
+			parent: bl.Number,
 			source: parsers.UpgradeSource,
 			depth:  bl.ReqDepth + 1,
+			from:   bl.Source,
 		}
 	}()
 }
diff --git a/internal/pool/pool.go b/internal/pool/pool.go
@@ -28,24 +28,6 @@ type BasePool struct {
 	isFallback  atomic.Bool
 }
 
-func (pool *BasePool) doRedirect(bl *pkg.Baseline, depth int) {
-	if depth >= pool.MaxRedirect {
-		return
-	}
-	reURL := pkg.FormatURL(bl.Url.Path, bl.RedirectURL)
-	pool.wg.Add(1)
-	go func() {
-		defer pool.wg.Done()
-		pool.addAddition(&Unit{
-			path:     reURL,
-			host:     bl.Host,
-			source:   parsers.RedirectSource,
-			frontUrl: bl.UrlString,
-			depth:    depth + 1,
-		})
-	}()
-}
-
 func (pool *BasePool) doRetry(bl *pkg.Baseline) {
 	if bl.Retry >= pool.RetryLimit {
 		return
@@ -55,8 +37,10 @@ func (pool *BasePool) doRetry(bl *pkg.Baseline) {
 		defer pool.wg.Done()
 		pool.addAddition(&Unit{
 			path:   bl.Path,
+			parent: bl.Number,
 			host:   bl.Host,
 			source: parsers.RetrySource,
+			from:   bl.Source,
 			retry:  bl.Retry + 1,
 		})
 	}()

diff --git a/internal/pool/type.go b/internal/pool/type.go
@@ -11,14 +11,24 @@ func newUnit(path string, source parsers.SpraySource) *Unit {
 
 type Unit struct {
 	number   int
+	parent   int
 	host     string
 	path     string
+	from     parsers.SpraySource
 	source   parsers.SpraySource
 	retry    int
 	frontUrl string
 	depth    int
 }
 
+func (u *Unit) Update(bl *pkg.Baseline) {
+	bl.Number = u.number
+	bl.Parent = u.parent
+	bl.Host = u.host
+	bl.Path = u.path
+	bl.Source = u.source
+}
+
 func NewBaselines() *Baselines {
 	return &Baselines{
 		baselines: map[int]*pkg.Baseline{},