From e483bb44398f010d23724b6e0ebaf7f40c29979c Mon Sep 17 00:00:00 2001 From: M09Ic Date: Mon, 14 Oct 2024 02:20:16 +0800 Subject: [PATCH] baseline add from and parent prop --- go.mod | 2 +- go.sum | 2 ++ internal/pool/brutepool.go | 38 ++++++++++++++++++++++++++++++++++---- internal/pool/checkpool.go | 4 ++++ internal/pool/pool.go | 20 ++------------------ internal/pool/type.go | 10 ++++++++++ 6 files changed, 53 insertions(+), 23 deletions(-) diff --git a/go.mod b/go.mod index 638c2f5..fafa128 100644 --- 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 index 0647b56..f10bbfb 100644 --- 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 index a5d270b..3242329 100644 --- 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 index 62a9b36..1b84cdc 100644 --- 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 index 34f3a36..d183b79 100644 --- 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 index ed4381f..9967169 100644 --- 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{},