diff --git a/core/ihttp/request.go b/core/ihttp/request.go index 93c2949..1302e17 100644 --- a/core/ihttp/request.go +++ b/core/ihttp/request.go @@ -31,26 +31,18 @@ type Request struct { ClientType int } -func (r *Request) SetHeaders(header map[string]string, RandomUA bool) { - if header["User-Agent"] == "" { - if RandomUA { - header["User-Agent"] = pkg.RandomUA() - } else { - header["User-Agent"] = pkg.DefaultUserAgent - } - } - - if header["Accept"] == "" { - header["Accept"] = "*/*" +func (r *Request) SetHeaders(header http.Header, RandomUA bool) { + if RandomUA { + r.SetHeader("User-Agent", pkg.RandomUA()) } if r.StandardRequest != nil { - for k, v := range header { - r.StandardRequest.Header.Set(k, v) - } + r.StandardRequest.Header = header } else if r.FastRequest != nil { for k, v := range header { - r.FastRequest.Header.Set(k, v) + for _, i := range v { + r.FastRequest.Header.Set(k, i) + } } } } diff --git a/core/option.go b/core/option.go index 7e790bb..14a38d1 100644 --- a/core/option.go +++ b/core/option.go @@ -99,7 +99,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.: --header 'Auth: example_auth'" config:"headers"` + Headers []string `short:"H" long:"header" description:"Strings, custom headers, e.g.: --header '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"` Cookie []string `long:"cookie" description:"Strings, custom cookie" config:"cookies"` diff --git a/core/pool/brutepool.go b/core/pool/brutepool.go index a7394a7..7460a0c 100644 --- a/core/pool/brutepool.go +++ b/core/pool/brutepool.go @@ -394,7 +394,6 @@ func (pool *BrutePool) NoScopeInvoke(v interface{}) { return } req.SetHeaders(pool.Headers, pool.RandomUserAgent) - req.SetHeader("User-Agent", pkg.RandomUA()) resp, reqerr := pool.client.Do(req) if pool.ClientType == ihttp.FAST { defer fasthttp.ReleaseResponse(resp.FastResponse) diff --git a/core/pool/checkpool.go b/core/pool/checkpool.go index b6040e2..136d460 100644 --- a/core/pool/checkpool.go +++ b/core/pool/checkpool.go @@ -36,7 +36,7 @@ func NewCheckPool(ctx context.Context, config *Config) (*CheckPool, error) { processCh: make(chan *baseline.Baseline, config.Thread), }, } - pool.Headers = map[string]string{"Connection": "close"} + pool.Headers.Set("Connection", "close") p, _ := ants.NewPoolWithFunc(config.Thread, pool.Invoke) pool.Pool = p diff --git a/core/pool/config.go b/core/pool/config.go index 6a0c1f2..5178dfe 100644 --- a/core/pool/config.go +++ b/core/pool/config.go @@ -7,6 +7,7 @@ import ( "github.com/chainreactors/words" "github.com/chainreactors/words/rule" "github.com/expr-lang/expr/vm" + "net/http" "sync" "time" ) @@ -27,7 +28,7 @@ type Config struct { BreakThreshold int32 Method string Mod SprayMod - Headers map[string]string + Headers http.Header ClientType int MatchExpr *vm.Program FilterExpr *vm.Program diff --git a/core/runner.go b/core/runner.go index 168c052..99b9809 100644 --- a/core/runner.go +++ b/core/runner.go @@ -15,6 +15,7 @@ import ( "github.com/panjf2000/ants/v2" "github.com/vbauerster/mpb/v8" "github.com/vbauerster/mpb/v8/decor" + "net/http" "strings" "sync" "time" @@ -67,7 +68,7 @@ func (r *Runner) PrepareConfig() *pool.Config { Thread: r.Threads, Timeout: time.Duration(r.Timeout) * time.Second, RateLimit: r.RateLimit, - Headers: r.Headers, + Headers: make(http.Header), Method: r.Method, Mod: pool.ModMap[r.Mod], OutputCh: r.outputCh, @@ -109,6 +110,18 @@ func (r *Runner) PrepareConfig() *pool.Config { } } + for k, v := range r.Headers { + config.Headers.Set(k, v) + } + + if config.Headers.Get("User-Agent") == "" { + config.Headers.Set("User-Agent", pkg.DefaultUserAgent) + } + + if config.Headers.Get("Accept") == "" { + config.Headers.Set("Accept", "*/*") + } + return config }