From b1d4af93a09ec4d3daad6a2c2490bbd532bfc054 Mon Sep 17 00:00:00 2001 From: M09Ic Date: Thu, 8 Feb 2024 16:28:27 +0800 Subject: [PATCH] support append-file --- cmd/cmd.go | 2 ++ internal/config.go | 1 + internal/option.go | 28 ++++++++++++++++++----- internal/pool.go | 57 +++++++++++++++++++++++++++++++++------------- internal/runner.go | 2 ++ internal/types.go | 1 + internal/utils.go | 2 +- 7 files changed, 70 insertions(+), 23 deletions(-) diff --git a/cmd/cmd.go b/cmd/cmd.go index 4f552e7..b73b9d1 100644 --- a/cmd/cmd.go +++ b/cmd/cmd.go @@ -83,6 +83,8 @@ func Spray() { logs.AddLevel(internal.LogVerbose, "verbose", "[=] %s {{suffix}}") if option.Debug { logs.Log.SetLevel(logs.Debug) + } else if len(option.Verbose) > 0 { + logs.Log.SetLevel(internal.LogVerbose) } logs.Log.SetColorMap(map[logs.Level]func(string) string{ diff --git a/internal/config.go b/internal/config.go index cde9fe3..ba9f4b4 100644 --- a/internal/config.go +++ b/internal/config.go @@ -37,6 +37,7 @@ type Config struct { FilterExpr *vm.Program RecuExpr *vm.Program AppendRule *rule.Program + AppendWords []string OutputCh chan *Baseline FuzzyCh chan *Baseline Fuzzy bool diff --git a/internal/option.go b/internal/option.go index 0a1f260..6c9ebd4 100644 --- a/internal/option.go +++ b/internal/option.go @@ -1,6 +1,7 @@ package internal import ( + "bytes" "errors" "fmt" "github.com/antonmedv/expr" @@ -49,6 +50,7 @@ type InputOptions struct { Rules []string `short:"r" long:"rules" description:"Files, rule files, e.g.: -r rule1.txt -r rule2.txt"` AppendRule []string `long:"append-rule" description:"Files, when found valid path , use append rule generator new word with current path"` FilterRule string `long:"filter-rule" description:"String, filter rule, e.g.: --rule-filter '>8 <4'"` + AppendFile []string `long:"append-file" description:"Files, when found valid path , use append file new word with current path"` } type FunctionOptions struct { @@ -108,7 +110,7 @@ type ModeOptions struct { Scope []string `long:"scope" description:"String, custom scope, e.g.: --scope *.example.com"` Recursive string `long:"recursive" default:"current.IsDir()" description:"String,custom recursive rule, e.g.: --recursive current.IsDir()"` Depth int `long:"depth" default:"0" description:"Int, recursive depth"` - Index string `long:"index" default:"" description:"String, custom index path"` + Index string `long:"index" default:"/" description:"String, custom index path"` Random string `long:"random" default:"" description:"String, custom random path"` CheckPeriod int `long:"check-period" default:"200" description:"Int, check period when request"` ErrPeriod int `long:"error-period" default:"10" description:"Int, check period when error"` @@ -188,9 +190,6 @@ func (opt *Option) PrepareRunner() (*Runner, error) { r.Progress.Start() logs.Log.SetOutput(r.Progress.Bypass()) } - if len(opt.Verbose) == 1 { - logs.Log.SetLevel(LogVerbose) - } // configuration if opt.Force { @@ -328,7 +327,7 @@ func (opt *Option) PrepareRunner() (*Runner, error) { } if opt.Rules != nil { - rules, err := loadFileAndCombine(opt.Rules) + rules, err := loadRuleAndCombine(opt.Rules) if err != nil { return nil, err } @@ -357,13 +356,30 @@ func (opt *Option) PrepareRunner() (*Runner, error) { } if opt.AppendRule != nil { - content, err := loadFileAndCombine(opt.AppendRule) + content, err := loadRuleAndCombine(opt.AppendRule) if err != nil { return nil, err } r.AppendRules = rule.Compile(string(content), "") } + if opt.AppendFile != nil { + var bs bytes.Buffer + for _, f := range opt.AppendFile { + content, err := ioutil.ReadFile(f) + if err != nil { + return nil, err + } + bs.Write(bytes.TrimSpace(content)) + bs.WriteString("\n") + } + lines := strings.Split(bs.String(), "\n") + for i, line := range lines { + lines[i] = strings.TrimSpace(line) + } + r.AppendWords = lines + } + ports := utils.ParsePort(opt.PortRange) // prepare task diff --git a/internal/pool.go b/internal/pool.go index a724fd8..e19a616 100644 --- a/internal/pool.go +++ b/internal/pool.go @@ -5,7 +5,7 @@ import ( "fmt" "github.com/chainreactors/logs" "github.com/chainreactors/parsers" - ihttp2 "github.com/chainreactors/spray/internal/ihttp" + "github.com/chainreactors/spray/internal/ihttp" "github.com/chainreactors/spray/pkg" "github.com/chainreactors/utils/iutils" "github.com/chainreactors/words" @@ -48,7 +48,7 @@ func NewPool(ctx context.Context, config *Config) (*Pool, error) { url: u, ctx: pctx, cancel: cancel, - client: ihttp2.NewClient(&ihttp2.ClientConfig{ + client: ihttp.NewClient(&ihttp.ClientConfig{ Thread: config.Thread, Type: config.ClientType, Timeout: time.Duration(config.Timeout) * time.Second, @@ -92,7 +92,7 @@ type Pool struct { isDir bool url *url.URL Statistor *pkg.Statistor - client *ihttp2.Client + client *ihttp.Client reqPool *ants.PoolWithFunc scopePool *ants.PoolWithFunc bar *pkg.Bar @@ -133,22 +133,24 @@ func (pool *Pool) checkRedirect(redirectURL string) bool { } } -func (pool *Pool) genReq(mod SprayMod, s string) (*ihttp2.Request, error) { +func (pool *Pool) genReq(mod SprayMod, s string) (*ihttp.Request, error) { if mod == HostSpray { - return ihttp2.BuildHostRequest(pool.ClientType, pool.BaseURL, s) + return ihttp.BuildHostRequest(pool.ClientType, pool.BaseURL, s) } else if mod == PathSpray { - return ihttp2.BuildPathRequest(pool.ClientType, pool.base, s) + return ihttp.BuildPathRequest(pool.ClientType, pool.base, s) } return nil, fmt.Errorf("unknown mod") } func (pool *Pool) Init() error { pool.initwg.Add(2) - if pool.Index != "" { + if pool.Index != "/" { logs.Log.Logf(LogVerbose, "custom index url: %s", BaseURL(pool.url)+FormatURL(BaseURL(pool.url), pool.Index)) pool.reqPool.Invoke(newUnit(pool.Index, InitIndexSource)) + //pool.urls[Dir(pool.Index)] = struct{}{} } else { pool.reqPool.Invoke(newUnit(pool.url.Path, InitIndexSource)) + //pool.urls[Dir(pool.url.Path)] = struct{}{} } if pool.Random != "" { @@ -163,7 +165,7 @@ func (pool *Pool) Init() error { logs.Log.Error(pool.index.String()) return fmt.Errorf(pool.index.ErrString) } - if pool.index.Chunked && pool.ClientType == ihttp2.FAST { + if pool.index.Chunked && pool.ClientType == ihttp.FAST { logs.Log.Warn("chunk encoding! buf current client FASTHTTP not support chunk decode") } logs.Log.Logf(LogVerbose, "[baseline.index] "+pool.index.Format([]string{"status", "length", "spend", "title", "frame", "redirect"})) @@ -286,7 +288,7 @@ func (pool *Pool) Invoke(v interface{}) { atomic.AddInt32(&pool.Statistor.ReqTotal, 1) unit := v.(*Unit) - var req *ihttp2.Request + var req *ihttp.Request var err error if unit.source == WordSource { req, err = pool.genReq(pool.Mod, unit.path) @@ -304,7 +306,7 @@ func (pool *Pool) Invoke(v interface{}) { start := time.Now() resp, reqerr := pool.client.Do(pool.ctx, req) - if pool.ClientType == ihttp2.FAST { + if pool.ClientType == ihttp.FAST { defer fasthttp.ReleaseResponse(resp.FastResponse) defer fasthttp.ReleaseRequest(req.FastRequest) } @@ -317,7 +319,6 @@ func (pool *Pool) Invoke(v interface{}) { bl = &Baseline{ SprayResult: &parsers.SprayResult{ UrlString: pool.base + unit.path, - IsValid: false, ErrString: reqerr.Error(), Reason: ErrRequestFailed.Error(), }, @@ -346,7 +347,7 @@ func (pool *Pool) Invoke(v interface{}) { pool.doRedirect(bl, unit.depth) } - if ihttp2.DefaultMaxBodySize != 0 && bl.BodyLength > ihttp2.DefaultMaxBodySize { + if ihttp.DefaultMaxBodySize != 0 && bl.BodyLength > ihttp.DefaultMaxBodySize { bl.ExceedLength = true } bl.Source = unit.source @@ -412,7 +413,7 @@ func (pool *Pool) Invoke(v interface{}) { func (pool *Pool) NoScopeInvoke(v interface{}) { defer pool.waiter.Done() unit := v.(*Unit) - req, err := ihttp2.BuildPathRequest(pool.ClientType, unit.path, "") + req, err := ihttp.BuildPathRequest(pool.ClientType, unit.path, "") if err != nil { logs.Log.Error(err.Error()) return @@ -420,7 +421,7 @@ func (pool *Pool) NoScopeInvoke(v interface{}) { req.SetHeaders(pool.Headers) req.SetHeader("User-Agent", RandomUA()) resp, reqerr := pool.client.Do(pool.ctx, req) - if pool.ClientType == ihttp2.FAST { + if pool.ClientType == ihttp.FAST { defer fasthttp.ReleaseResponse(resp.FastResponse) defer fasthttp.ReleaseRequest(req.FastRequest) } @@ -507,9 +508,12 @@ func (pool *Pool) Handler() { } if bl.IsValid || bl.IsFuzzy { - pool.waiter.Add(2) + pool.waiter.Add(3) pool.doCrawl(bl) pool.doRule(bl) + if _, ok := pool.urls[Dir(bl.Url.Path)]; !ok { + pool.doAppendWords(bl) + } } // 如果要进行递归判断, 要满足 bl有效, mod为path-spray, 当前深度小于最大递归深度 if bl.IsValid { @@ -530,7 +534,7 @@ func (pool *Pool) Handler() { pool.analyzeDone = true } -func (pool *Pool) PreCompare(resp *ihttp2.Response) error { +func (pool *Pool) PreCompare(resp *ihttp.Response) error { status := resp.StatusCode() if iutils.IntsContains(WhiteStatus, status) { // 如果为白名单状态码则直接返回 @@ -720,6 +724,27 @@ func (pool *Pool) doRule(bl *Baseline) { }() } +func (pool *Pool) doAppendWords(bl *Baseline) { + if pool.AppendWords == nil { + pool.waiter.Done() + return + } + if bl.Source == AppendSource { + pool.waiter.Done() + return + } + + go func() { + defer pool.waiter.Done() + for _, u := range pool.AppendWords { + pool.addAddition(&Unit{ + path: Dir(bl.Url.Path) + u, + source: AppendSource, + }) + } + }() +} + func (pool *Pool) doRetry(bl *Baseline) { if bl.Retry >= pool.Retry { return diff --git a/internal/runner.go b/internal/runner.go index 0138f02..0e60663 100644 --- a/internal/runner.go +++ b/internal/runner.go @@ -41,6 +41,7 @@ type Runner struct { Wordlist []string Rules *rule.Program AppendRules *rule.Program + AppendWords []string Headers map[string]string Fns []func(string) []string FilterExpr *vm.Program @@ -104,6 +105,7 @@ func (r *Runner) PrepareConfig() *Config { FilterExpr: r.FilterExpr, RecuExpr: r.RecursiveExpr, AppendRule: r.AppendRules, + AppendWords: r.AppendWords, IgnoreWaf: r.IgnoreWaf, Crawl: r.Crawl, Scope: r.Scope, diff --git a/internal/types.go b/internal/types.go index f7ebcd9..568b404 100644 --- a/internal/types.go +++ b/internal/types.go @@ -20,6 +20,7 @@ const ( CommonFileSource UpgradeSource RetrySource + AppendSource ) func newUnit(path string, source int) *Unit { diff --git a/internal/utils.go b/internal/utils.go index 94cb971..ad114af 100644 --- a/internal/utils.go +++ b/internal/utils.go @@ -97,7 +97,7 @@ func loadFileToSlice(filename string) ([]string, error) { return ss, nil } -func loadFileAndCombine(filename []string) (string, error) { +func loadRuleAndCombine(filename []string) (string, error) { var bs bytes.Buffer for _, f := range filename { if data, ok := pkg.Rules[f]; ok {