fix: brute init not set header, #94

Author: M09Ic

core/ihttp/request.go

@@ -31,18 +31,18 @@ type Request struct {
 	ClientType      int
 }
 
-func (r *Request) SetHeaders(header map[string]string, RandomUA bool) {
+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)
+			}
 		}
 	}
 }

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"`

core/pool/brutepool.go

@@ -107,12 +107,12 @@ type BrutePool struct {
 }
 
 func (pool *BrutePool) Init() error {
-	if pool.Headers["User-Agent"] == "" {
-		pool.Headers["User-Agent"] = pkg.DefaultUserAgent
+	if pool.Headers.Get("User-Agent") == "" {
+		pool.Headers.Set("User-Agent", pkg.DefaultUserAgent)
 	}
 
-	if pool.Headers["Accept"] == "" {
-		pool.Headers["Accept"] = "*/*"
+	if pool.Headers.Get("Accept") == "" {
+		pool.Headers.Set("Accept", "*/*")
 	}
 
 	pool.initwg.Add(2)

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

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,10 @@ func (r *Runner) PrepareConfig() *pool.Config {
 		}
 	}
 
+	for k, v := range r.Headers {
+		config.Headers.Set(k, v)
+	}
+
 	return config
 }