mirror of
https://github.com/chainreactors/spray.git
synced 2025-11-04 09:58:03 +00:00
添加进度条,修复多处bug, 初步可以使用
This commit is contained in:
M09Ic
parent
26d084eb6e
commit
1685331fa5
9
go.mod
9
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
|
||||
)
|
||||
|
||||
10
go.sum
10
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=
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -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 (
|
||||
|
||||
@ -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
|
||||
@ -118,13 +124,14 @@ func (r *Runner) Run() {
|
||||
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())
|
||||
|
||||
57
pkg/bar.go
Normal file
57
pkg/bar.go
Normal file
@ -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
|
||||
}
|
||||
@ -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,
|
||||
MaxConnsPerHost: thread,
|
||||
IdleConnTimeout: time.Duration(timeout) * time.Second,
|
||||
DisableKeepAlives: false,
|
||||
}
|
||||
|
||||
c := &Client{
|
||||
|
||||
@ -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
|
||||
}
|
||||
|
||||
@ -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
|
||||
}
|
||||
Loading…
x
Reference in New Issue
Block a user