From 5142012ceb94aa50b958a3f365695d7de4c4ac42 Mon Sep 17 00:00:00 2001 From: M09Ic Date: Tue, 3 Jan 2023 18:43:12 +0800 Subject: [PATCH] =?UTF-8?q?=E5=AE=9E=E8=A3=85--cookie=20--user-agent,=20--?= =?UTF-8?q?header=E5=8F=82=E6=95=B0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- internal/option.go | 23 +++++++++++++++-------- internal/pool.go | 2 ++ internal/runner.go | 3 +-- pkg/config.go | 3 +-- pkg/ihttp/request.go | 10 +++++++++- 5 files changed, 28 insertions(+), 13 deletions(-) diff --git a/internal/option.go b/internal/option.go index e96dc32..c52cf68 100644 --- a/internal/option.go +++ b/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'"` - //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"` - //Method string `long:"method" default:"GET" description:"String, custom method"` - //Cookie string `long:"cookie" description:"String, custom cookie"` - MaxBodyLength int `long:"max-length" default:"100" description:"Int, max response body length, default 100k, e.g. -max-length 1000"` + 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"` + RandomUserAgent bool `long:"random-agent" description:"Bool, use random with default user-agent"` + 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"` } 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, ",") } diff --git a/internal/pool.go b/internal/pool.go index dd0a721..c9c65c2 100644 --- a/internal/pool.go +++ b/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 { diff --git a/internal/runner.go b/internal/runner.go index be6bfd3..a46e538 100644 --- a/internal/runner.go +++ b/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 diff --git a/pkg/config.go b/pkg/config.go index e017bd9..334025b 100644 --- a/pkg/config.go +++ b/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 diff --git a/pkg/ihttp/request.go b/pkg/ihttp/request.go index e67bbbc..6fe89d2 100644 --- a/pkg/ihttp/request.go +++ b/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()