Skip to content

Commit

Permalink
新增--rate-limit, 用来限制单个pool请求速率
Browse files Browse the repository at this point in the history
  • Loading branch information
@M09Ic
M09Ic committed Jan 12, 2023
1 parent 78ee22b commit 758a274
Show file tree
Hide file tree
Showing 6 changed files with 67 additions and 48 deletions.
1 change: 1 addition & 0 deletions go.mod
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,5 +32,6 @@ require (
github.com/twmb/murmur3 v1.1.6 // indirect
github.com/valyala/bytebufferpool v1.0.0 // indirect
golang.org/x/sys v0.2.0 // indirect
golang.org/x/time v0.3.0 // indirect
gopkg.in/yaml.v2 v2.4.0 // indirect
)
2 changes: 2 additions & 0 deletions go.sum
View file
Original file line number Diff line number Diff line change
Expand Up @@ -103,6 +103,8 @@ golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
golang.org/x/text v0.3.2/go.mod h1:bEr9sfX3Q8Zfm5fL9x+3itogRgK3+ptLWKqgva+5dAk=
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/time v0.3.0 h1:rg5rLMjNzMS1RkNLzCG38eapWhnYLFYXDXj2gOlr8j4=
golang.org/x/time v0.3.0/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
Expand Down
40 changes: 23 additions & 17 deletions internal/option.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,22 +21,23 @@ type Option struct {
InputOptions `group:"Input Options"`
FunctionOptions `group:"Function Options"`
OutputOptions `group:"Output Options"`
PluginOptions `group:"Plugin Options"`
RequestOptions `group:"Request Options"`
ModeOptions `group:"Modify Options"`
MiscOptions `group:"Miscellaneous Options"`
}

type InputOptions struct {
ResumeFrom string `long:"resume"`
URL []string `short:"u" long:"url" description:"String, Multi, input baseurl, e.g.: http://google.com"`
URL []string `short:"u" long:"url" description:"Strings, input baseurl, e.g.: http://google.com"`
URLFile string `short:"l" long:"list" description:"File, input filename"`
Raw string `long:"raw" description:"File, input raw request filename"`
Offset int `long:"offset" description:"Int, wordlist offset"`
Limit int `long:"limit" description:"Int, wordlist limit, start with offset. e.g.: --offset 1000 --limit 100"`
Dictionaries []string `short:"d" long:"dict" description:"Files, Multi,dict files, e.g.: -d 1.txt -d 2.txt"`
Word string `short:"w" long:"word" description:"String, word generate dsl, e.g.: -w test{?ld#4}"`
Rules []string `short:"r" long:"rules" description:"Files, Multi, rule files, e.g.: -r rule1.txt -r rule2.txt"`
AppendRule []string `long:"append-rule" description:"File, when found valid path , use append rule generator new word with current path"`
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'"`
}

Expand All @@ -46,15 +47,15 @@ type FunctionOptions struct {
RemoveExtensions string `long:"remove-extension" description:"String, remove extensions (separated by commas), e.g.: --remove-extension jsp,jspx"`
Uppercase bool `short:"U" long:"uppercase" desvcription:"Bool, upper wordlist, e.g.: --uppercase"`
Lowercase bool `short:"L" long:"lowercase" description:"Bool, lower wordlist, e.g.: --lowercase"`
Prefixes []string `long:"prefix" description:"Strings, Multi, add prefix, e.g.: --prefix aaa --prefix bbb"`
Suffixes []string `long:"suffix" description:"Strings, Multi, add suffix, e.g.: --suffix aaa --suffix bbb"`
Replaces map[string]string `long:"replace" description:"Strings, Multi, replace string, e.g.: --replace aaa:bbb --replace ccc:ddd"`
Prefixes []string `long:"prefix" description:"Strings, add prefix, e.g.: --prefix aaa --prefix bbb"`
Suffixes []string `long:"suffix" description:"Strings, add suffix, e.g.: --suffix aaa --suffix bbb"`
Replaces map[string]string `long:"replace" description:"Strings, replace string, e.g.: --replace aaa:bbb --replace ccc:ddd"`
}

type OutputOptions struct {
Match string `long:"match" description:"String, custom match function, e.g.: --match current.Status != 200" json:"match,omitempty"`
Filter string `long:"filter" description:"String, custom filter function, e.g.: --filter current.Body contains 'hello'" json:"filter,omitempty"`
Extracts []string `long:"extract" description:"String, Multi, extract response, e.g.: --extract js --extract ip --extract version:(.*?)" json:"extracts,omitempty"`
Extracts []string `long:"extract" description:"Strings, extract response, e.g.: --extract js --extract ip --extract version:(.*?)" json:"extracts,omitempty"`
OutputFile string `short:"f" long:"file" description:"String, output filename" json:"output_file,omitempty"`
Format string `short:"F" long:"format" description:"String, output format, e.g.: --format 1.json"`
FuzzyFile string `long:"fuzzy-file" description:"String, fuzzy output filename" json:"fuzzy_file,omitempty"`
Expand All @@ -66,27 +67,31 @@ type OutputOptions struct {
}

type RequestOptions struct {
Headers []string `long:"header" description:"String, Multi, custom headers, e.g.: --headers 'Auth: example_auth'"`
Headers []string `long:"header" description:"Strings, custom headers, e.g.: --headers 'Auth: example_auth'"`
UserAgent string `long:"user-agent" description:"String, custom user-agent, e.g.: --user-agent Custom"`
RandomUserAgent bool `long:"random-agent" description:"Bool, use random with default user-agent"`
Cookie []string `long:"cookie" description:"String, Multi, custom cookie"`
Cookie []string `long:"cookie" description:"Strings, custom cookie"`
ReadAll bool `long:"read-all" description:"Bool, read all response body"`
MaxBodyLength int `long:"max-length" default:"100" description:"Int, max response body length (kb), default 100k, e.g. -max-length 1000"`
}

type PluginOptions struct {
Advance bool `short:"a" long:"advance" description:"Bool, enable crawl and active"`
Active bool `long:"active" description:"Bool, enable active finger detect"`
Bak bool `long:"bak" description:"Bool, enable bak found"`
FileBak bool `long:"file-bak" description:"Bool, enable valid result bak found, equal --append-rule rule/filebak.txt"`
Common bool `long:"common" description:"Bool, enable common file found"`
Crawl bool `long:"crawl" description:"Bool, enable crawl"`
CrawlDepth int `long:"crawl-depth" default:"3" description:"Int, crawl depth"`
CrawlScope string `long:"crawl-scope" description:"Int, crawl scope (todo)"`
}

type ModeOptions struct {
Advance bool `short:"a" long:"advance" description:"Bool, enable crawl and active"`
Active bool `long:"active" description:"Bool, enable active finger detect"`
Crawl bool `long:"crawl" description:"Bool, enable crawl"`
Bak bool `long:"bak" description:"Bool, enable bak found"`
FileBak bool `long:"file-bak" description:"Bool, enable valid result bak found, equal --append-rule rule/filebak.txt"`
Common bool `long:"common" description:"Bool, enable common file found"`
RateLimit int `long:"rate-limit" default:"0" description:"Int, request rate limit (rate/s), e.g.: --rate-limit 100"`
Force bool `long:"force" description:"Bool, skip error break"`
CheckOnly bool `long:"check-only" description:"Bool, check only"`
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"`
CrawlDepth int `long:"crawl-depth" default:"3" description:"Int, crawl depth"`
CrawlScope string `long:"crawl-scope" description:"Int, crawl scope (todo)"`
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"`
BreakThreshold int `long:"error-threshold" default:"20" description:"Int, break when the error exceeds the threshold "`
Expand Down Expand Up @@ -121,6 +126,7 @@ func (opt *Option) PrepareRunner() (*Runner, error) {
PoolSize: opt.PoolSize,
Mod: opt.Mod,
Timeout: opt.Timeout,
RateLimit: opt.RateLimit,
Deadline: opt.Deadline,
Headers: make(map[string]string),
Offset: opt.Offset,
Expand Down
69 changes: 38 additions & 31 deletions internal/pool.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ import (
"github.com/chainreactors/words/rule"
"github.com/panjf2000/ants/v2"
"github.com/valyala/fasthttp"
"golang.org/x/time/rate"
"net/url"
"path"
"strconv"
Expand Down Expand Up @@ -50,8 +51,9 @@ func NewPool(ctx context.Context, config *pkg.Config) (*Pool, error) {
tempCh: make(chan *pkg.Baseline, config.Thread),
checkCh: make(chan int),
additionCh: make(chan *Unit, 100),
wg: sync.WaitGroup{},
waiter: sync.WaitGroup{},
initwg: sync.WaitGroup{},
limiter: rate.NewLimiter(rate.Limit(config.RateLimit), 1),
reqCount: 1,
failedCount: 1,
}
Expand Down Expand Up @@ -127,7 +129,7 @@ func NewPool(ctx context.Context, config *pkg.Config) (*Pool, error) {

// 如果要进行递归判断, 要满足 bl有效, mod为path-spray, 当前深度小于最大递归深度
if bl.IsValid {
pool.wg.Add(2)
pool.waiter.Add(2)
pool.doCrawl(bl)
pool.doRule(bl)
if bl.RecuDepth < MaxRecursion {
Expand All @@ -137,7 +139,7 @@ func NewPool(ctx context.Context, config *pkg.Config) (*Pool, error) {
}
}
pool.OutputCh <- bl
pool.wg.Done()
pool.waiter.Done()
}

pool.analyzeDone = true
Expand Down Expand Up @@ -170,8 +172,9 @@ type Pool struct {
urls map[string]struct{}
analyzeDone bool
worder *words.Worder
limiter *rate.Limiter
locker sync.Mutex
wg sync.WaitGroup
waiter sync.WaitGroup
initwg sync.WaitGroup // 初始化用, 之后改成锁
}

Expand Down Expand Up @@ -237,17 +240,17 @@ func (pool *Pool) genReq(s string) (*ihttp.Request, error) {
func (pool *Pool) Run(ctx context.Context, offset, limit int) {
pool.worder.RunWithRules()
if pool.Active {
pool.wg.Add(1)
pool.waiter.Add(1)
go pool.doActive()
}

if pool.Bak {
pool.wg.Add(1)
pool.waiter.Add(1)
go pool.doBak()
}

if pool.Common {
pool.wg.Add(1)
pool.waiter.Add(1)
go pool.doCommonFile()
}

Expand All @@ -256,7 +259,7 @@ func (pool *Pool) Run(ctx context.Context, offset, limit int) {
wait := func() {
if !worderDone {
worderDone = true
pool.wg.Wait()
pool.waiter.Wait()
close(closeCh)
}
}
Expand All @@ -280,7 +283,7 @@ Loop:
continue
}

pool.wg.Add(1)
pool.waiter.Add(1)
pool.urls[u] = struct{}{}
pool.reqPool.Invoke(newUnit(pool.safePath(u), WordSource)) // 原样的目录拼接, 输入了几个"/"就是几个, 适配java的目录解析
case source := <-pool.checkCh:
Expand All @@ -296,7 +299,7 @@ Loop:
}
if _, ok := pool.urls[unit.path]; ok {
logs.Log.Debugf("[%s] duplicate path: %s, skipped", pkg.GetSourceName(unit.source), pool.base+unit.path)
pool.wg.Done()
pool.waiter.Done()
} else {
pool.urls[unit.path] = struct{}{}
pool.reqPool.Invoke(unit)
Expand All @@ -310,12 +313,16 @@ Loop:
}
}

pool.wg.Wait()
pool.waiter.Wait()
pool.Statistor.EndTime = time.Now().Unix()
pool.Close()
}

func (pool *Pool) Invoke(v interface{}) {
if pool.RateLimit != 0 {
pool.limiter.Wait(pool.ctx)
}

atomic.AddInt32(&pool.Statistor.ReqTotal, 1)
unit := v.(*Unit)
req, err := pool.genReq(unit.path)
Expand Down Expand Up @@ -355,7 +362,7 @@ func (pool *Pool) Invoke(v interface{}) {

// 手动处理重定向
if bl.IsValid && unit.source != CheckSource && bl.RedirectURL != "" {
pool.wg.Add(1)
pool.waiter.Add(1)
pool.doRedirect(bl, unit.depth)
}

Expand All @@ -378,7 +385,7 @@ func (pool *Pool) Invoke(v interface{}) {
pool.locker.Lock()
pool.index = bl
pool.locker.Unlock()
pool.wg.Add(1)
pool.waiter.Add(1)
pool.doCrawl(bl)
if bl.Status == 200 || (bl.Status/100) == 3 {
pool.OutputCh <- bl
Expand Down Expand Up @@ -524,13 +531,13 @@ func (pool *Pool) Upgrade(bl *pkg.Baseline) error {
}

func (pool *Pool) doRedirect(bl *pkg.Baseline, depth int) {
defer pool.wg.Done()
defer pool.waiter.Done()
if depth >= MaxRedirect {
return
}
reURL := FormatURL(bl.Url.Path, bl.RedirectURL)

pool.wg.Add(1)
pool.waiter.Add(1)
go pool.addAddition(&Unit{
path: reURL,
source: RedirectSource,
Expand All @@ -541,24 +548,24 @@ func (pool *Pool) doRedirect(bl *pkg.Baseline, depth int) {

func (pool *Pool) doCrawl(bl *pkg.Baseline) {
if !pool.Crawl || bl.ReqDepth >= MaxCrawl {
pool.wg.Done()
pool.waiter.Done()
return
}
bl.CollectURL()
if bl.URLs == nil {
pool.wg.Done()
pool.waiter.Done()
return
}

go func() {
defer pool.wg.Done()
defer pool.waiter.Done()
for _, u := range bl.URLs {
if u = FormatURL(bl.Url.Path, u); u == "" {
continue
}

// 通过map去重, 只有新的url才会进入到该逻辑
pool.wg.Add(1)
pool.waiter.Add(1)
pool.addAddition(&Unit{
path: u,
source: CrawlSource,
Expand All @@ -571,18 +578,18 @@ func (pool *Pool) doCrawl(bl *pkg.Baseline) {

func (pool *Pool) doRule(bl *pkg.Baseline) {
if pool.AppendRule == nil {
pool.wg.Done()
pool.waiter.Done()
return
}
if bl.Source == int(RuleSource) {
pool.wg.Done()
pool.waiter.Done()
return
}

go func() {
defer pool.wg.Done()
defer pool.waiter.Done()
for u := range rule.RunAsStream(pool.AppendRule.Expressions, path.Base(bl.Path)) {
pool.wg.Add(1)
pool.waiter.Add(1)
pool.addAddition(&Unit{
path: Dir(bl.Url.Path) + u,
source: RuleSource,
Expand All @@ -592,9 +599,9 @@ func (pool *Pool) doRule(bl *pkg.Baseline) {
}

func (pool *Pool) doActive() {
defer pool.wg.Done()
defer pool.waiter.Done()
for _, u := range pkg.ActivePath {
pool.wg.Add(1)
pool.waiter.Add(1)
pool.addAddition(&Unit{
path: pool.dir + u[1:],
source: ActiveSource,
Expand All @@ -603,14 +610,14 @@ func (pool *Pool) doActive() {
}

func (pool *Pool) doBak() {
defer pool.wg.Done()
defer pool.waiter.Done()
worder, err := words.NewWorderWithDsl("{?0}.{@bak_ext}", [][]string{pkg.BakGenerator(pool.url.Host)}, nil)
if err != nil {
return
}
worder.Run()
for w := range worder.C {
pool.wg.Add(1)
pool.waiter.Add(1)
pool.addAddition(&Unit{
path: pool.dir + w,
source: BakSource,
Expand All @@ -623,7 +630,7 @@ func (pool *Pool) doBak() {
}
worder.Run()
for w := range worder.C {
pool.wg.Add(1)
pool.waiter.Add(1)
pool.addAddition(&Unit{
path: pool.dir + w,
source: BakSource,
Expand All @@ -632,9 +639,9 @@ func (pool *Pool) doBak() {
}

func (pool *Pool) doCommonFile() {
defer pool.wg.Done()
defer pool.waiter.Done()
for _, u := range mask.SpecialWords["common_file"] {
pool.wg.Add(1)
pool.waiter.Add(1)
pool.addAddition(&Unit{
path: pool.dir + u,
source: CommonFileSource,
Expand Down Expand Up @@ -665,7 +672,7 @@ func (pool *Pool) addAddition(u *Unit) {
func (pool *Pool) addFuzzyBaseline(bl *pkg.Baseline) {
if _, ok := pool.baselines[bl.Status]; !ok && pkg.IntsContains(FuzzyStatus, bl.Status) {
bl.Collect()
pool.wg.Add(1)
pool.waiter.Add(1)
pool.doCrawl(bl)
pool.baselines[bl.Status] = bl
logs.Log.Infof("[baseline.%dinit] %s", bl.Status, bl.Format([]string{"status", "length", "spend", "title", "frame", "redirect"}))
Expand Down
2 changes: 2 additions & 0 deletions internal/runner.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,7 @@ type Runner struct {
Progress *uiprogress.Progress
Offset int
Limit int
RateLimit int
Total int
Deadline int
CheckPeriod int
Expand All @@ -83,6 +84,7 @@ func (r *Runner) PrepareConfig() *pkg.Config {
config := &pkg.Config{
Thread: r.Threads,
Timeout: r.Timeout,
RateLimit: r.RateLimit,
Headers: r.Headers,
Mod: pkg.ModMap[r.Mod],
OutputCh: r.OutputCh,
Expand Down
1 change: 1 addition & 0 deletions pkg/config.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ type Config struct {
Thread int
Wordlist []string
Timeout int
RateLimit int
CheckPeriod int
ErrPeriod int
BreakThreshold int
Expand Down

0 comments on commit 758a274

Please sign in to comment.