实装--cookie --user-agent, --header参数

internal/option.go

@@ -67,12 +67,11 @@ type OutputOptions struct {
 }
 
 type RequestOptions struct {
-	Headers []string `long:"header" description:"String, custom headers, e.g.: --headers 'Auth: example_auth'"`
+	Headers         []string `long:"header" description:"String, custom headers, e.g.: --headers 'Auth: example_auth'"`
-	//UserAgent       string   `long:"user-agent" description:"String, custom user-agent, e.g.: --user-agent Custom"`
+	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"`
+	RandomUserAgent bool     `long:"random-agent" description:"Bool, use random with default user-agent"`
-	//Method          string   `long:"method" default:"GET" description:"String, custom method"`
+	Cookie          []string `long:"cookie" description:"String, custom cookie"`
-	//Cookie          string   `long:"cookie" description:"String, custom cookie"`
+	MaxBodyLength   int      `long:"max-length" default:"100" description:"Int, max response body length (kb), default 100k, e.g. -max-length 1000"`
-	MaxBodyLength int `long:"max-length" default:"100" description:"Int, max response body length, default 100k, e.g. -max-length 1000"`
 }
 
 type ModeOptions struct {
@@ -97,7 +96,7 @@ type MiscOptions struct {
 	Deadline int    `long:"deadline" default:"999999" description:"Int, deadline (seconds)"` // todo 总的超时时间,适配云函数的deadline
 	Timeout  int    `long:"timeout" default:"2" description:"Int, timeout with request (seconds)"`
 	PoolSize int    `short:"p" long:"pool" default:"5" description:"Int, Pool size"`
-	Threads  int    `short:"t" long:"thread" default:"20" description:"Int, number of threads per pool (seconds)"`
+	Threads  int    `short:"t" long:"thread" default:"20" description:"Int, number of threads per pool"`
 	Debug    bool   `long:"debug" description:"Bool, output debug info"`
 	NoColor  bool   `long:"no-color" description:"Bool, no color"`
 	Quiet    bool   `short:"q" long:"quiet" description:"Bool, Quiet"`
@@ -119,6 +118,7 @@ func (opt *Option) PrepareRunner() (*Runner, error) {
 		Mod:            opt.Mod,
 		Timeout:        opt.Timeout,
 		Deadline:       opt.Deadline,
+		Headers:        make(map[string]string),
 		Offset:         opt.Offset,
 		Total:          opt.Limit,
 		taskCh:         make(chan *Task),
@@ -414,10 +414,17 @@ func (opt *Option) PrepareRunner() (*Runner, error) {
 		if i == -1 {
 			logs.Log.Warn("invalid header")
 		} else {
-			r.Headers.Add(h[:i], h[i+2:])
+			r.Headers[h[:i]] = h[i+2:]
 		}
 	}
 
+	if opt.UserAgent != "" {
+		r.Headers["User-Agent"] = opt.UserAgent
+	}
+	if opt.Cookie != nil {
+		r.Headers["Cookie"] = strings.Join(opt.Cookie, "; ")
+	}
+
 	if opt.OutputProbe != "" {
 		r.Probes = strings.Split(opt.OutputProbe, ",")
 	}

internal/pool.go

@@ -52,6 +52,8 @@ func NewPool(ctx context.Context, config *pkg.Config) (*Pool, error) {
 			logs.Log.Error(err.Error())
 			return
 		}
+		req.SetHeaders(pool.Headers)
+
 		start := time.Now()
 		resp, reqerr := pool.client.Do(pctx, req)
 		if pool.ClientType == ihttp.FAST {

internal/runner.go

@@ -12,7 +12,6 @@ import (
 	"github.com/chainreactors/words/rule"
 	"github.com/gosuri/uiprogress"
 	"github.com/panjf2000/ants/v2"
-	"net/http"
 	"sync"
 	"time"
 )
@@ -40,7 +39,7 @@ type Runner struct {
 	URLList        []string
 	Wordlist       []string
 	Rules          []rule.Expression
-	Headers        http.Header
+	Headers        map[string]string
 	Fns            []func(string) string
 	FilterExpr     *vm.Program
 	MatchExpr      *vm.Program

pkg/config.go

@@ -2,7 +2,6 @@ package pkg
 
 import (
 	"github.com/antonmedv/expr/vm"
-	"net/http"
 )
 
 type SprayMod int
@@ -29,7 +28,7 @@ type Config struct {
 	BreakThreshold int
 	Method         string
 	Mod            SprayMod
-	Headers        http.Header
+	Headers        map[string]string
 	ClientType     int
 	MatchExpr      *vm.Program
 	FilterExpr     *vm.Program

pkg/ihttp/request.go

@@ -35,7 +35,7 @@ type Request struct {
 	ClientType      int
 }
 
-func (r *Request) SetHeader(header map[string]string) {
+func (r *Request) SetHeaders(header map[string]string) {
 	if r.StandardRequest != nil {
 		for k, v := range header {
 			r.StandardRequest.Header.Set(k, v)
@@ -47,6 +47,14 @@ func (r *Request) SetHeader(header map[string]string) {
 	}
 }
 
+func (r *Request) SetHeader(key, value string) {
+	if r.StandardRequest != nil {
+		r.StandardRequest.Header.Set(key, value)
+	} else if r.FastRequest != nil {
+		r.FastRequest.Header.Set(key, value)
+	}
+}
+
 func (r *Request) URI() string {
 	if r.FastRequest != nil {
 		return r.FastRequest.URI().String()