admin/spray
1
0
Fork 0
mirror of https://github.com/chainreactors/spray.git synced 2025-09-15 11:40:13 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
M09Ic 2023-01-03 18:43:12 +08:00
parent e6aabe44c5
commit 5142012ceb
5 changed files with 28 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

23
internal/option.go
View File

@ -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, ",")
}

2
internal/pool.go
View File

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

3
internal/runner.go
View File

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

3
pkg/config.go
View File

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

10
pkg/ihttp/request.go
View File

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