Merge pull request #95 from chainreactors/dev

merge v1.2.1

core/ihttp/request.go

@@ -31,26 +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 header["User-Agent"] == "" {
+	if RandomUA {
-		if RandomUA {
+		r.SetHeader("User-Agent", pkg.RandomUA())
-			header["User-Agent"] = pkg.RandomUA()
-		} else {
-			header["User-Agent"] = pkg.DefaultUserAgent
-		}
-	}
-
-	if header["Accept"] == "" {
-		header["Accept"] = "*/*"
 	}
 
 	if r.StandardRequest != nil {
-		for k, v := range header {
+		r.StandardRequest.Header = header
-			r.StandardRequest.Header.Set(k, v)
-		}
 	} 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

@@ -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)

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

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,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
 }