适配进度条, 优化性能

internal/baseline.go

@@ -21,15 +21,17 @@ func NewBaseline(u *url.URL, resp *http.Response) *baseline {
 		IsValid:    true,
 	}
 
-	var header string
+	var header strings.Builder
 	for k, v := range resp.Header {
-		// stringbuilder
 		for _, i := range v {
-			header += fmt.Sprintf("%s: %s\r\n", k, i)
+			header.WriteString(k)
+			header.WriteString(": ")
+			header.WriteString(i)
+			header.WriteString("\r\n")
 		}
 	}
-	bl.Header = header
+	bl.Header = header.String()
-	bl.HeaderLength = len(header)
+	bl.HeaderLength = header.Len()
 
 	redirectURL, err := resp.Location()
 	if err == nil {
@@ -120,7 +122,17 @@ func (bl *baseline) FuzzyCompare() bool {
 }
 
 func (bl *baseline) String() string {
-	return fmt.Sprintf("%s - %d - %d [%s]", bl.UrlString, bl.Status, bl.BodyLength, bl.Frameworks.ToString())
+	var line strings.Builder
+	line.WriteString("[+] ")
+	line.WriteString(bl.UrlString)
+	line.WriteString(fmt.Sprintf(" - %d - %d ", bl.Status, bl.BodyLength))
+	if bl.RedirectURL != "" {
+		line.WriteString("-> ")
+		line.WriteString(bl.RedirectURL)
+	}
+	line.WriteString(bl.Frameworks.ToString())
+	//line.WriteString(bl.Extracteds)
+	return line.String()
 }
 
 func (bl *baseline) Jsonify() string {

internal/pool.go

@@ -48,7 +48,6 @@ func NewPool(ctx context.Context, config *pkg.Config, outputCh chan *baseline) (
 			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())
@@ -56,7 +55,7 @@ func NewPool(ctx context.Context, config *pkg.Config, outputCh chan *baseline) (
 			bl = &baseline{Err: err}
 		} else {
 			defer resp.Body.Close() // 必须要关闭body ,否则keep-alive无法生效
-			if err = pool.PreCompare(resp); err == nil {
+			if err = pool.PreCompare(resp); err == nil || unit.source == CheckSource {
 				// 通过预对比跳过一些无用数据, 减少性能消耗
 				bl = NewBaseline(req.URL, resp)
 			} else if err == ErrWaf {
@@ -67,15 +66,11 @@ func NewPool(ctx context.Context, config *pkg.Config, outputCh chan *baseline) (
 		}
 
 		switch unit.source {
-		case InitSource:
+		case CheckSource:
 			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()
@@ -109,7 +104,7 @@ type Pool struct {
 func (p *Pool) Init() error {
 	//for i := 0; i < p.baseReqCount; i++ {
 	p.wg.Add(1)
-	_ = p.pool.Invoke(newUnit(pkg.RandPath(), InitSource))
+	_ = p.pool.Invoke(newUnit(pkg.RandPath(), CheckSource))
 	//}
 	p.wg.Wait()
 	// todo 分析baseline
@@ -158,7 +153,7 @@ Loop:
 			break Loop
 		}
 	}
-
+	p.bar.Close()
 	p.wg.Wait()
 }
 
@@ -202,7 +197,7 @@ func (p *Pool) BuildHostRequest(host string) (*http.Request, error) {
 type sourceType int
 
 const (
-	InitSource sourceType = iota + 1
+	CheckSource sourceType = iota + 1
 	WordSource
 	WafSource
 )

internal/runner.go

@@ -13,7 +13,7 @@ import (
 	"sync"
 )
 
-var BlackStatus = []int{404, 410}
+var BlackStatus = []int{400, 404, 410}
 var FuzzyStatus = []int{403, 500, 501, 502, 503}
 
 type Runner struct {
@@ -31,6 +31,7 @@ type Runner struct {
 	Pools      map[string]*Pool
 	Deadline   int    `long:"deadline" default:"600"` // todo 总的超时时间,适配云函数的deadline
 	Debug      bool   `long:"debug"`
+	Quiet      bool   `short:"q" long:"quiet"`
 	Mod        string `short:"m" long:"mod" default:"path"`
 	OutputCh   chan *baseline
 	Progress   *uiprogress.Progress
@@ -38,11 +39,14 @@ type Runner struct {
 
 func (r *Runner) Prepare() error {
 	r.Progress = uiprogress.New()
-	r.Progress.Start()
 
 	if r.Debug {
 		logs.Log.Level = logs.Debug
 	}
+	if !r.Quiet {
+		r.Progress.Start()
+		logs.Log.Writer = r.Progress.Bypass()
+	}
 
 	var file *os.File
 	var err error
@@ -157,7 +161,7 @@ func (r *Runner) Outputting() {
 		select {
 		case bl := <-r.OutputCh:
 			if bl.IsValid {
-				logs.Log.Console(bl.String() + "\n")
+				logs.Log.Console(bl.String())
 			} else {
 				logs.Log.Debug(bl.String())
 			}

pkg/bar.go

@@ -3,16 +3,14 @@ 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),
+		Bar:   progress.AddBar(total),
-		url:    u,
+		url:   u,
-		writer: progress.Bypass(),
+		spend: 1,
-		spend:  1,
 	}
 
 	bar.AppendCompleted()
@@ -37,10 +35,9 @@ func NewBar(u string, total int, progress *uiprogress.Progress) *Bar {
 }
 
 type Bar struct {
-	spend  int
+	spend int
-	url    string
+	url   string
-	close  bool
+	close bool
-	writer io.Writer
 	*uiprogress.Bar
 }
 
@@ -48,10 +45,6 @@ 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

@@ -18,7 +18,14 @@ func NewClient(thread int, timeout int) *Client {
 		MaxConnsPerHost: thread,
 		IdleConnTimeout: time.Duration(timeout) * time.Second,
 	}
-
+	//c := &fasthttp.Client{
+	//	TLSConfig: &tls.Config{
+	//		Renegotiation:      tls.RenegotiateOnceAsClient,
+	//		InsecureSkipVerify: true,
+	//	},
+	//	MaxConnsPerHost:     thread,
+	//	MaxIdleConnDuration: time.Duration(timeout) * time.Second,
+	//}
 	c := &Client{
 		client: &http.Client{
 			Transport:     tr,