From f755fc3816b443c702e08a5b9d72d1d5a425993d Mon Sep 17 00:00:00 2001 From: M09ic Date: Fri, 17 May 2024 17:47:11 +0800 Subject: [PATCH] support -x/--method custom http method support --raw parser input from raw http --- internal/ihttp/request.go | 5 +++-- internal/option.go | 38 ++++++++++++++++++++++++++++++-------- internal/pool/brutepool.go | 15 +++++++++------ internal/pool/checkpool.go | 10 +++++----- internal/pool/pool.go | 26 +++++++++++++------------- internal/runner.go | 2 ++ 6 files changed, 62 insertions(+), 34 deletions(-) diff --git a/internal/ihttp/request.go b/internal/ihttp/request.go index 6fe89d2..34b72ab 100644 --- a/internal/ihttp/request.go +++ b/internal/ihttp/request.go @@ -5,13 +5,14 @@ import ( "net/http" ) -func BuildPathRequest(clientType int, base, path string) (*Request, error) { +func BuildPathRequest(clientType int, base, path, method string) (*Request, error) { if clientType == FAST { req := fasthttp.AcquireRequest() + req.Header.SetMethod(method) req.SetRequestURI(base + path) return &Request{FastRequest: req, ClientType: FAST}, nil } else { - req, err := http.NewRequest("GET", base+path, nil) + req, err := http.NewRequest(method, base+path, nil) return &Request{StandardRequest: req, ClientType: STANDARD}, err } } diff --git a/internal/option.go b/internal/option.go index a5e30aa..fc35556 100644 --- a/internal/option.go +++ b/internal/option.go @@ -1,6 +1,7 @@ package internal import ( + "bufio" "bytes" "errors" "fmt" @@ -17,6 +18,7 @@ import ( "github.com/chainreactors/words/rule" "github.com/vbauerster/mpb/v8" "io/ioutil" + "net/http" "net/url" "os" "path/filepath" @@ -43,13 +45,13 @@ type Option struct { } type InputOptions struct { - ResumeFrom string `long:"resume" description:"File, resume filename" ` - Config string `short:"c" long:"config" description:"File, config filename"` - 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"` - PortRange string `short:"p" long:"port" description:"String, input port range, e.g.: 80,8080-8090,db"` - CIDRs string `long:"cidr" description:"String, input cidr, e.g.: 1.1.1.1/24 "` - //Raw string `long:"raw" description:"File, input raw request filename"` + ResumeFrom string `long:"resume" description:"File, resume filename" ` + Config string `short:"c" long:"config" description:"File, config filename"` + 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"` + PortRange string `short:"p" long:"port" description:"String, input port range, e.g.: 80,8080-8090,db"` + CIDRs string `long:"cidr" description:"String, input cidr, e.g.: 1.1.1.1/24 "` + RawFile string `long:"raw" description:"File, input raw request filename"` Dictionaries []string `short:"d" long:"dict" description:"Files, Multi,dict files, e.g.: -d 1.txt -d 2.txt" config:"dictionaries"` NoDict bool `long:"no-dict" description:"Bool, no dictionary" config:"no-dict"` Word string `short:"w" long:"word" description:"String, word generate dsl, e.g.: -w test{?ld#4}" config:"word"` @@ -92,6 +94,7 @@ type OutputOptions struct { } type RequestOptions struct { + Method string `short:"x" long:"method" default:"GET" description:"String, request method, e.g.: --method POST" config:"method"` Headers []string `long:"header" description:"Strings, custom headers, e.g.: --headers 'Auth: example_auth'" config:"headers"` UserAgent string `long:"user-agent" description:"String, custom user-agent, e.g.: --user-agent Custom" config:"useragent"` RandomUserAgent bool `long:"random-agent" description:"Bool, use random with default user-agent" config:"random-useragent"` @@ -162,6 +165,7 @@ func (opt *Option) PrepareRunner() (*Runner, error) { RateLimit: opt.RateLimit, Deadline: opt.Deadline, Headers: make(map[string]string), + Method: opt.Method, Offset: opt.Offset, Total: opt.Limit, taskCh: make(chan *Task), @@ -459,6 +463,24 @@ func (opt *Option) PrepareRunner() (*Runner, error) { taskfrom = "cmd" r.Count = len(opt.URL) + } else if opt.RawFile != "" { + raw, err := os.Open(opt.RawFile) + if err != nil { + return nil, err + } + + req, err := http.ReadRequest(bufio.NewReader(raw)) + if err != nil { + return nil, err + } + go func() { + opt.GenerateTasks(tasks, fmt.Sprintf("http://%s%s", req.Host, req.URL.String()), ports) + close(tasks) + }() + r.Method = req.Method + for k, _ := range req.Header { + r.Headers[k] = req.Header.Get(k) + } } else if opt.CIDRs != "" { if len(ports) == 0 { ports = []string{"80", "443"} @@ -739,7 +761,7 @@ func (opt *Option) Validate() error { return errors.New("--resume and --depth cannot be used at the same time") } - if opt.ResumeFrom == "" && opt.URL == nil && opt.URLFile == "" && opt.CIDRs == "" { + if opt.ResumeFrom == "" && opt.URL == nil && opt.URLFile == "" && opt.CIDRs == "" && opt.RawFile == "" { return fmt.Errorf("without any target, please use -u/-l/-c/--resume to set targets") } return nil diff --git a/internal/pool/brutepool.go b/internal/pool/brutepool.go index 16d3e11..e3777d6 100644 --- a/internal/pool/brutepool.go +++ b/internal/pool/brutepool.go @@ -2,6 +2,7 @@ package pool import ( "context" + "errors" "fmt" "github.com/chainreactors/logs" "github.com/chainreactors/parsers" @@ -37,7 +38,7 @@ func NewBrutePool(ctx context.Context, config *Config) (*BrutePool, error) { pctx, cancel := context.WithCancel(ctx) pool := &BrutePool{ Baselines: NewBaselines(), - This: &This{ + BasePool: &BasePool{ Config: config, ctx: pctx, Cancel: cancel, @@ -83,7 +84,7 @@ func NewBrutePool(ctx context.Context, config *Config) (*BrutePool, error) { type BrutePool struct { *Baselines - *This + *BasePool base string // url的根目录, 在爬虫或者redirect时, 会需要用到根目录进行拼接 isDir bool url *url.URL @@ -125,7 +126,7 @@ func (pool *BrutePool) genReq(mod SprayMod, s string) (*ihttp.Request, error) { if mod == HostSpray { return ihttp.BuildHostRequest(pool.ClientType, pool.BaseURL, s) } else if mod == PathSpray { - return ihttp.BuildPathRequest(pool.ClientType, pool.base, s) + return ihttp.BuildPathRequest(pool.ClientType, pool.base, s, pool.Method) } return nil, fmt.Errorf("unknown mod") } @@ -311,7 +312,9 @@ func (pool *BrutePool) Invoke(v interface{}) { } req.SetHeaders(pool.Headers) - req.SetHeader("User-Agent", pkg.RandomUA()) + if pool.RandomUserAgent { + req.SetHeader("User-Agent", pkg.RandomUA()) + } start := time.Now() resp, reqerr := pool.client.Do(pool.ctx, req) @@ -322,7 +325,7 @@ func (pool *BrutePool) Invoke(v interface{}) { // compare与各种错误处理 var bl *pkg.Baseline - if reqerr != nil && reqerr != fasthttp.ErrBodyTooLarge { + if reqerr != nil && !errors.Is(reqerr, fasthttp.ErrBodyTooLarge) { atomic.AddInt32(&pool.failedCount, 1) atomic.AddInt32(&pool.Statistor.FailedNumber, 1) bl = &pkg.Baseline{ @@ -422,7 +425,7 @@ func (pool *BrutePool) Invoke(v interface{}) { func (pool *BrutePool) NoScopeInvoke(v interface{}) { defer pool.wg.Done() unit := v.(*Unit) - req, err := ihttp.BuildPathRequest(pool.ClientType, unit.path, "") + req, err := ihttp.BuildPathRequest(pool.ClientType, unit.path, "", pool.Method) if err != nil { logs.Log.Error(err.Error()) return diff --git a/internal/pool/checkpool.go b/internal/pool/checkpool.go index 89678f0..5cb0998 100644 --- a/internal/pool/checkpool.go +++ b/internal/pool/checkpool.go @@ -19,7 +19,7 @@ import ( func NewCheckPool(ctx context.Context, config *Config) (*CheckPool, error) { pctx, cancel := context.WithCancel(ctx) pool := &CheckPool{ - &This{ + &BasePool{ Config: config, Statistor: pkg.NewStatistor(""), ctx: pctx, @@ -38,12 +38,12 @@ func NewCheckPool(ctx context.Context, config *Config) (*CheckPool, error) { pool.Headers = map[string]string{"Connection": "close"} p, _ := ants.NewPoolWithFunc(config.Thread, pool.Invoke) - pool.This.Pool = p + pool.BasePool.Pool = p return pool, nil } type CheckPool struct { - *This + *BasePool } func (pool *CheckPool) Run(ctx context.Context, offset, limit int) { @@ -81,12 +81,12 @@ Loop: } pool.wg.Add(1) - _ = pool.This.Pool.Invoke(newUnit(u, parsers.CheckSource)) + _ = pool.BasePool.Pool.Invoke(newUnit(u, parsers.CheckSource)) case u, ok := <-pool.additionCh: if !ok { continue } - _ = pool.This.Pool.Invoke(u) + _ = pool.BasePool.Pool.Invoke(u) case <-pool.closeCh: break Loop case <-ctx.Done(): diff --git a/internal/pool/pool.go b/internal/pool/pool.go index f4bb8ab..6cfefe8 100644 --- a/internal/pool/pool.go +++ b/internal/pool/pool.go @@ -14,7 +14,7 @@ import ( "sync" ) -type This struct { +type BasePool struct { *Config Statistor *pkg.Statistor Pool *ants.PoolWithFunc @@ -31,7 +31,7 @@ type This struct { wg sync.WaitGroup } -func (pool *This) doRedirect(bl *pkg.Baseline, depth int) { +func (pool *BasePool) doRedirect(bl *pkg.Baseline, depth int) { if depth >= MaxRedirect { return } @@ -48,7 +48,7 @@ func (pool *This) doRedirect(bl *pkg.Baseline, depth int) { }() } -func (pool *This) doRule(bl *pkg.Baseline) { +func (pool *BasePool) doRule(bl *pkg.Baseline) { if pool.AppendRule == nil { pool.wg.Done() return @@ -69,7 +69,7 @@ func (pool *This) doRule(bl *pkg.Baseline) { }() } -func (pool *This) doAppendWords(bl *pkg.Baseline) { +func (pool *BasePool) doAppendWords(bl *pkg.Baseline) { if pool.AppendWords == nil { pool.wg.Done() return @@ -90,7 +90,7 @@ func (pool *This) doAppendWords(bl *pkg.Baseline) { }() } -func (pool *This) doRetry(bl *pkg.Baseline) { +func (pool *BasePool) doRetry(bl *pkg.Baseline) { if bl.Retry >= pool.Retry { return } @@ -105,7 +105,7 @@ func (pool *This) doRetry(bl *pkg.Baseline) { }() } -func (pool *This) doActive() { +func (pool *BasePool) doActive() { defer pool.wg.Done() for _, u := range pkg.ActivePath { pool.addAddition(&Unit{ @@ -115,7 +115,7 @@ func (pool *This) doActive() { } } -func (pool *This) doCommonFile() { +func (pool *BasePool) doCommonFile() { defer pool.wg.Done() for _, u := range mask.SpecialWords["common_file"] { pool.addAddition(&Unit{ @@ -125,7 +125,7 @@ func (pool *This) doCommonFile() { } } -func (pool *This) addAddition(u *Unit) { +func (pool *BasePool) addAddition(u *Unit) { // 强行屏蔽报错, 防止goroutine泄露 pool.wg.Add(1) defer func() { @@ -135,25 +135,25 @@ func (pool *This) addAddition(u *Unit) { pool.additionCh <- u } -func (pool *This) Close() { +func (pool *BasePool) Close() { pool.Bar.Close() } -func (pool *This) genReq(s string) (*ihttp.Request, error) { +func (pool *BasePool) genReq(s string) (*ihttp.Request, error) { if pool.Mod == HostSpray { return ihttp.BuildHostRequest(pool.ClientType, pool.BaseURL, s) } else if pool.Mod == PathSpray { - return ihttp.BuildPathRequest(pool.ClientType, pool.BaseURL, s) + return ihttp.BuildPathRequest(pool.ClientType, pool.BaseURL, s, pool.Method) } return nil, fmt.Errorf("unknown mod") } -func (pool *This) putToOutput(bl *pkg.Baseline) { +func (pool *BasePool) putToOutput(bl *pkg.Baseline) { pool.OutLocker.Add(1) pool.OutputCh <- bl } -func (pool *This) putToFuzzy(bl *pkg.Baseline) { +func (pool *BasePool) putToFuzzy(bl *pkg.Baseline) { pool.OutLocker.Add(1) bl.IsFuzzy = true pool.FuzzyCh <- bl diff --git a/internal/runner.go b/internal/runner.go index 8700f65..e28a99c 100644 --- a/internal/runner.go +++ b/internal/runner.go @@ -43,6 +43,7 @@ type Runner struct { AppendRules *rule.Program AppendWords []string Headers map[string]string + Method string Fns []func(string) []string FilterExpr *vm.Program MatchExpr *vm.Program @@ -92,6 +93,7 @@ func (r *Runner) PrepareConfig() *pool.Config { Timeout: r.Timeout, RateLimit: r.RateLimit, Headers: r.Headers, + Method: r.Method, Mod: pool.ModMap[r.Mod], OutputCh: r.outputCh, FuzzyCh: r.fuzzyCh,