添加进度条,修复多处bug, 初步可以使用

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
 )

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=

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)

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 {
+		pool.genReq = func(s string) (*http.Request, error) {
-			return pkg.BuildPathRequest(s, *config.BaseReq)
+			return pool.BuildPathRequest(s)
 		}
 	case pkg.HostSpray:
-		pool.genReq = func(s string) *http.Request {
+		pool.genReq = func(s string) (*http.Request, error) {
-			return pkg.BuildHostRequest(s, *config.BaseReq)
+			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))
+			p.wg.Add(1)
-		case <-time.NewTimer(time.Duration(p.DeadlineTime)).C:
+			_ = 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 (

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,
+				BaseURL:      u,
-				Wordlist: r.Wordlist,
+				Wordlist:     r.Wordlist,
-				Thread:   r.Threads,
+				Thread:       r.Threads,
-				Timeout:  2,
+				Timeout:      2,
-				Headers:  r.Headers,
+				Headers:      r.Headers,
-				Mod:      pkg.ModMap[r.Mod],
+				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())

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
+}

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{
+		MaxConnsPerHost: thread,
-			//Timeout:   time.Duration(delay) * time.Second,
+		IdleConnTimeout: time.Duration(timeout) * time.Second,
-			//KeepAlive: time.Duration(delay) * time.Second,
-			//DualStack: true,
-		}).DialContext,
-		MaxIdleConnsPerHost: thread,
-		MaxIdleConns:        thread,
-		IdleConnTimeout:     time.Duration(timeout) * time.Second,
-		DisableKeepAlives:   false,
 	}
 
 	c := &Client{

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

pkg/requests.go

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