From 1685331fa596586143a6747bf5c34ba851e20c75 Mon Sep 17 00:00:00 2001 From: M09Ic Date: Tue, 20 Sep 2022 04:01:38 +0800 Subject: [PATCH] =?UTF-8?q?=E6=B7=BB=E5=8A=A0=E8=BF=9B=E5=BA=A6=E6=9D=A1,?= =?UTF-8?q?=E4=BF=AE=E5=A4=8D=E5=A4=9A=E5=A4=84bug,=20=E5=88=9D=E6=AD=A5?= =?UTF-8?q?=E5=8F=AF=E4=BB=A5=E4=BD=BF=E7=94=A8?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- go.mod | 9 ++++-- go.sum | 10 ++++++ internal/baseline.go | 2 -- internal/pool.go | 72 ++++++++++++++++++++++++++------------------ internal/runner.go | 23 +++++++++----- pkg/bar.go | 57 +++++++++++++++++++++++++++++++++++ pkg/client.go | 12 ++------ pkg/config.go | 9 ------ pkg/requests.go | 15 --------- 9 files changed, 133 insertions(+), 76 deletions(-) create mode 100644 pkg/bar.go delete mode 100644 pkg/requests.go diff --git a/go.mod b/go.mod index be6cf3a..ed12bf0 100644 --- a/go.mod +++ b/go.mod @@ -3,10 +3,10 @@ module github.com/chainreactors/spray go 1.17 require ( - github.com/chainreactors/gogo/v2 v2.8.7 - github.com/chainreactors/logs v0.5.0 github.com/chainreactors/files v0.2.0 // indirect + github.com/chainreactors/gogo/v2 v2.8.7 github.com/chainreactors/ipcs v0.0.9 // indirect + github.com/chainreactors/logs v0.5.0 github.com/chainreactors/parsers v0.0.1 // indirect github.com/chainreactors/words v0.0.1 // indirect github.com/go-dedup/simhash v0.0.0-20170904020510-9ecaca7b509c @@ -20,12 +20,15 @@ require ( github.com/cnf/structhash v0.0.0-20201127153200-e1b16c1ebc08 // indirect github.com/go-dedup/megophone v0.0.0-20170830025436-f01be21026f5 // indirect github.com/go-dedup/text v0.0.0-20170907015346-8bb1b95e3cb7 // indirect + github.com/gosuri/uilive v0.0.4 // indirect + github.com/gosuri/uiprogress v0.0.1 // indirect github.com/jessevdk/go-flags v1.5.0 // indirect github.com/karrick/godirwalk v1.16.1 // indirect + github.com/mattn/go-isatty v0.0.16 // indirect github.com/pkg/errors v0.9.1 // indirect github.com/twmb/murmur3 v1.1.6 // indirect golang.org/x/crypto v0.0.0-20200604202706-70a84ac30bf9 // indirect golang.org/x/net v0.0.0-20190603091049-60506f45cf65 // indirect - golang.org/x/sys v0.0.0-20220209214540-3681064d5158 // indirect + golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab // indirect gopkg.in/yaml.v3 v3.0.1 // indirect ) diff --git a/go.sum b/go.sum index a443e3f..805ee6b 100644 --- a/go.sum +++ b/go.sum @@ -12,6 +12,8 @@ github.com/chainreactors/gogo/v2 v2.8.5 h1:7DfAQdyupToU4f+OYG+Hz4+0phyz3UGRrsAyT github.com/chainreactors/gogo/v2 v2.8.5/go.mod h1:SqlLDXXN/WgjVHpRIjB8OytRmHnj8l4Ic+h7pe/4wsY= github.com/chainreactors/gogo/v2 v2.8.6 h1:Xc9wi8K8/3u9s0V0vlch2rKHzQ+ATkLlHExGosUvVok= github.com/chainreactors/gogo/v2 v2.8.6/go.mod h1:SqlLDXXN/WgjVHpRIjB8OytRmHnj8l4Ic+h7pe/4wsY= +github.com/chainreactors/gogo/v2 v2.8.7 h1:pN83yw7mtLFcNV12ZwxKdTR2sDZFSiqva/JRQMMRvao= +github.com/chainreactors/gogo/v2 v2.8.7/go.mod h1:KFMyHLZSkpgDh/17N/k9gKYQ81MO+9MLdTMHQvoBSj0= github.com/chainreactors/ipcs v0.0.9 h1:4Onroq7gXLG5SLCCgNDx3JmtLxB4XgepGdHCtLp1Ows= github.com/chainreactors/ipcs v0.0.9/go.mod h1:E9M3Ohyq0TYQLlV4i2dbM9ThBZB1Nnd7Oexoie2xLII= github.com/chainreactors/logs v0.4.2/go.mod h1:Y0EtAnoF0kiASIJUnXN0pcOt420iRpHOAnOhEphzRHA= @@ -34,11 +36,17 @@ github.com/go-dedup/text v0.0.0-20170907015346-8bb1b95e3cb7 h1:11wFcswN+37U+Byjx github.com/go-dedup/text v0.0.0-20170907015346-8bb1b95e3cb7/go.mod h1:wSsK4VOECOSfSYTzkBFw+iGY7wj59e7X96ABtNj9aCQ= github.com/google/go-cmp v0.5.6/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= github.com/google/go-cmp v0.5.7/go.mod h1:n+brtR0CgQNWTVd5ZUFpTBC8YFBDLK/h/bpaJ8/DtOE= +github.com/gosuri/uilive v0.0.4 h1:hUEBpQDj8D8jXgtCdBu7sWsy5sbW/5GhuO8KBwJ2jyY= +github.com/gosuri/uilive v0.0.4/go.mod h1:V/epo5LjjlDE5RJUcqx8dbw+zc93y5Ya3yg8tfZ74VI= +github.com/gosuri/uiprogress v0.0.1 h1:0kpv/XY/qTmFWl/SkaJykZXrBBzwwadmW8fRb7RJSxw= +github.com/gosuri/uiprogress v0.0.1/go.mod h1:C1RTYn4Sc7iEyf6j8ft5dyoZ4212h8G1ol9QQluh5+0= github.com/jessevdk/go-flags v1.5.0 h1:1jKYvbxEjfUl0fmqTCOfonvskHHXMjBySTLW4y9LFvc= github.com/jessevdk/go-flags v1.5.0/go.mod h1:Fw0T6WPc1dYxT4mKEZRfG5kJhaTDP9pj1c2EWnYs/m4= github.com/josharian/native v1.0.0/go.mod h1:7X/raswPFr05uY3HiLlYeyQntB6OO7E/d2Cu7qoaN2w= github.com/karrick/godirwalk v1.16.1 h1:DynhcF+bztK8gooS0+NDJFrdNZjJ3gzVzC545UNA9iw= github.com/karrick/godirwalk v1.16.1/go.mod h1:j4mkqPuvaLI8mp1DroR3P6ad7cyYd4c1qeJ3RV7ULlk= +github.com/mattn/go-isatty v0.0.16 h1:bq3VjFmv/sOjHtdEhmkEV4x1AJtvUvOJ2PFAZ5+peKQ= +github.com/mattn/go-isatty v0.0.16/go.mod h1:kYGgaQfpe5nmfYZH+SKPsOc2e4SrIfOl2e/yFXSvRLM= github.com/mdlayher/arp v0.0.0-20220512170110-6706a2966875/go.mod h1:kfOoFJuHWp76v1RgZCb9/gVUc7XdY877S2uVYbNliGc= github.com/mdlayher/ethernet v0.0.0-20220221185849-529eae5b6118/go.mod h1:ZFUnHIVchZ9lJoWoEGUg8Q3M4U8aNNWA3CVSUTkW4og= github.com/mdlayher/packet v1.0.0/go.mod h1:eE7/ctqDhoiRhQ44ko5JZU2zxB88g+JH/6jmnjzPjOU= @@ -77,6 +85,8 @@ golang.org/x/sys v0.0.0-20210320140829-1e4c9ba3b0c4/go.mod h1:h1NjWce9XRLGQEsW7w golang.org/x/sys v0.0.0-20210927094055-39ccf1dd6fa6/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20220209214540-3681064d5158 h1:rm+CHSpPEEW2IsXUib1ThaHIjuBVZjxNgSKmBLFfD4c= golang.org/x/sys v0.0.0-20220209214540-3681064d5158/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab h1:2QkjZIsXupsJbJIdSjjUOgWK3aEtzyuh2mPt3l/CkeU= +golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= diff --git a/internal/baseline.go b/internal/baseline.go index 69e5d86..bd56f04 100644 --- a/internal/baseline.go +++ b/internal/baseline.go @@ -46,9 +46,7 @@ func NewBaseline(u *url.URL, resp *http.Response) *baseline { } else { logs.Log.Error("readfull failed" + err.Error()) } - _ = resp.Body.Close() } - if len(bl.Body) > 0 { bl.Md5 = parsers.Md5Hash(bl.Body) bl.Mmh3 = parsers.Mmh3Hash32(bl.Body) diff --git a/internal/pool.go b/internal/pool.go index 9e70881..a467a02 100644 --- a/internal/pool.go +++ b/internal/pool.go @@ -2,7 +2,7 @@ package internal import ( "context" - "fmt" + "github.com/chainreactors/logs" "github.com/chainreactors/spray/pkg" "github.com/chainreactors/words" "github.com/panjf2000/ants/v2" @@ -18,13 +18,7 @@ var ( ) func NewPool(ctx context.Context, config *pkg.Config, outputCh chan *baseline) (*Pool, error) { - err := config.Init() - if err != nil { - return nil, fmt.Errorf("pool init failed, %w", err) - } - pctx, cancel := context.WithCancel(ctx) - pool := &Pool{ Config: config, ctx: pctx, @@ -37,25 +31,31 @@ func NewPool(ctx context.Context, config *pkg.Config, outputCh chan *baseline) ( switch config.Mod { case pkg.PathSpray: - pool.genReq = func(s string) *http.Request { - return pkg.BuildPathRequest(s, *config.BaseReq) + pool.genReq = func(s string) (*http.Request, error) { + return pool.BuildPathRequest(s) } case pkg.HostSpray: - pool.genReq = func(s string) *http.Request { - return pkg.BuildHostRequest(s, *config.BaseReq) + pool.genReq = func(s string) (*http.Request, error) { + return pool.BuildHostRequest(s) } } p, _ := ants.NewPoolWithFunc(config.Thread, func(i interface{}) { var bl *baseline unit := i.(*Unit) - req := pool.genReq(unit.path) + req, err := pool.genReq(unit.path) + if err != nil { + logs.Log.Error(err.Error()) + return + } + resp, err := pool.client.Do(pctx, req) if err != nil { //logs.Log.Debugf("%s request error, %s", strurl, err.Error()) pool.errorCount++ bl = &baseline{Err: err} } else { + defer resp.Body.Close() // 必须要关闭body ,否则keep-alive无法生效 if err = pool.PreCompare(resp); err == nil { // 通过预对比跳过一些无用数据, 减少性能消耗 bl = NewBaseline(req.URL, resp) @@ -71,9 +71,14 @@ func NewPool(ctx context.Context, config *pkg.Config, outputCh chan *baseline) ( pool.baseline = bl case WordSource: // todo compare - pool.outputCh <- bl + //pool.outputCh <- bl + // todo 重构output + if bl.IsValid { + pool.bar.Print(bl.String()) + } } //todo connectivity check + pool.bar.Done() pool.wg.Done() }) @@ -88,34 +93,25 @@ type Pool struct { *pkg.Config client *pkg.Client pool *ants.PoolWithFunc + bar *pkg.Bar ctx context.Context //baseReq *http.Request baseline *baseline outputCh chan *baseline totalCount int errorCount int - genReq func(string) *http.Request + genReq func(s string) (*http.Request, error) //wordlist []string worder *words.Worder wg *sync.WaitGroup } -func (p *Pool) Add(u *Unit) error { - p.wg.Add(1) - _ = p.pool.Invoke(u) - p.wg.Wait() - - if p.baseline.Err != nil { - return p.baseline.Err - } - return nil -} - func (p *Pool) Init() error { //for i := 0; i < p.baseReqCount; i++ { - _ = p.Add(newUnit(pkg.RandPath(), InitSource)) + p.wg.Add(1) + _ = p.pool.Invoke(newUnit(pkg.RandPath(), InitSource)) //} - + p.wg.Wait() // todo 分析baseline // 检测基本访问能力 @@ -152,8 +148,9 @@ Loop: break Loop } p.totalCount++ - _ = p.Add(newUnit(u, WordSource)) - case <-time.NewTimer(time.Duration(p.DeadlineTime)).C: + p.wg.Add(1) + _ = p.pool.Invoke(newUnit(u, WordSource)) + case <-time.NewTimer(time.Duration(p.DeadlineTime) * time.Second).C: break Loop case <-ctx.Done(): break Loop @@ -185,6 +182,23 @@ func (p *Pool) RunWithWord(words []string) { } +func (p *Pool) BuildPathRequest(path string) (*http.Request, error) { + req, err := http.NewRequest("GET", p.BaseURL+path, nil) + if err != nil { + return nil, err + } + return req, nil +} + +func (p *Pool) BuildHostRequest(host string) (*http.Request, error) { + req, err := http.NewRequest("GET", p.BaseURL, nil) + req.Host = host + if err != nil { + return nil, err + } + return req, nil +} + type sourceType int const ( diff --git a/internal/runner.go b/internal/runner.go index 6b5d603..6df6c47 100644 --- a/internal/runner.go +++ b/internal/runner.go @@ -5,6 +5,7 @@ import ( "fmt" "github.com/chainreactors/logs" "github.com/chainreactors/spray/pkg" + "github.com/gosuri/uiprogress" "io/ioutil" "net/http" "os" @@ -28,16 +29,21 @@ type Runner struct { Threads int `short:"t" long:"thread" default:"20"` PoolSize int `short:"p" long:"pool"` Pools map[string]*Pool - Deadline int `long:"deadline"` // todo 总的超时时间,适配云函数的deadline + Deadline int `long:"deadline" default:"600"` // todo 总的超时时间,适配云函数的deadline Debug bool `long:"debug"` Mod string `short:"m" long:"mod" default:"path"` OutputCh chan *baseline + Progress *uiprogress.Progress } func (r *Runner) Prepare() error { + r.Progress = uiprogress.New() + r.Progress.Start() + if r.Debug { logs.Log.Level = logs.Debug } + var file *os.File var err error urlfrom := r.URLFile @@ -112,19 +118,20 @@ func (r *Runner) Run() { u := u go func() { config := &pkg.Config{ - BaseURL: u, - Wordlist: r.Wordlist, - Thread: r.Threads, - Timeout: 2, - Headers: r.Headers, - Mod: pkg.ModMap[r.Mod], + BaseURL: u, + Wordlist: r.Wordlist, + Thread: r.Threads, + Timeout: 2, + Headers: r.Headers, + Mod: pkg.ModMap[r.Mod], + DeadlineTime: r.Deadline, } pool, err := NewPool(ctx, config, r.OutputCh) if err != nil { logs.Log.Error(err.Error()) return } - + pool.bar = pkg.NewBar(u, len(r.Wordlist), r.Progress) err = pool.Init() if err != nil { logs.Log.Error(err.Error()) diff --git a/pkg/bar.go b/pkg/bar.go new file mode 100644 index 0000000..23196d0 --- /dev/null +++ b/pkg/bar.go @@ -0,0 +1,57 @@ +package pkg + +import ( + "fmt" + "github.com/gosuri/uiprogress" + "io" + "time" +) + +func NewBar(u string, total int, progress *uiprogress.Progress) *Bar { + bar := &Bar{ + Bar: progress.AddBar(total), + url: u, + writer: progress.Bypass(), + spend: 1, + } + + bar.AppendCompleted() + bar.PrependElapsed() + bar.PrependFunc(func(b *uiprogress.Bar) string { + return fmt.Sprintf("%v/s", bar.Current()/bar.spend) + }) + + bar.PrependFunc(func(b *uiprogress.Bar) string { + return u + }) + + go func() { + for !bar.close { + select { + case <-time.After(time.Duration(250) * time.Millisecond): + bar.spend++ + } + } + }() + return bar +} + +type Bar struct { + spend int + url string + close bool + writer io.Writer + *uiprogress.Bar +} + +func (bar *Bar) Done() { + bar.Incr() +} + +func (bar *Bar) Print(s string) { + fmt.Fprintln(bar.writer, s) +} + +func (bar *Bar) Close() { + bar.close = true +} diff --git a/pkg/client.go b/pkg/client.go index f6a0a10..8144edf 100644 --- a/pkg/client.go +++ b/pkg/client.go @@ -3,7 +3,6 @@ package pkg import ( "context" "crypto/tls" - "net" "net/http" "time" ) @@ -16,15 +15,8 @@ func NewClient(thread int, timeout int) *Client { Renegotiation: tls.RenegotiateOnceAsClient, InsecureSkipVerify: true, }, - DialContext: (&net.Dialer{ - //Timeout: time.Duration(delay) * time.Second, - //KeepAlive: time.Duration(delay) * time.Second, - //DualStack: true, - }).DialContext, - MaxIdleConnsPerHost: thread, - MaxIdleConns: thread, - IdleConnTimeout: time.Duration(timeout) * time.Second, - DisableKeepAlives: false, + MaxConnsPerHost: thread, + IdleConnTimeout: time.Duration(timeout) * time.Second, } c := &Client{ diff --git a/pkg/config.go b/pkg/config.go index 04537b3..e966946 100644 --- a/pkg/config.go +++ b/pkg/config.go @@ -29,12 +29,3 @@ type Config struct { Headers http.Header DeadlineTime int } - -func (c *Config) Init() (err error) { - c.BaseReq, err = http.NewRequest(c.Method, c.BaseURL, nil) - if err != nil { - return err - } - c.BaseReq.Header = c.Headers - return nil -} diff --git a/pkg/requests.go b/pkg/requests.go deleted file mode 100644 index 18454a7..0000000 --- a/pkg/requests.go +++ /dev/null @@ -1,15 +0,0 @@ -package pkg - -import ( - "net/http" -) - -func BuildPathRequest(path string, req http.Request) *http.Request { - req.URL.Path = path - return &req -} - -func BuildHostRequest(u string, req http.Request) *http.Request { - req.Host = u - return &req -}