实装对probes的支持

新增自动检测退出任务时的细节输出
优化错误请求的输出
优化options的结构

internal/baseline.go

@@ -165,6 +165,27 @@ func (bl *baseline) Additional(key string) string {
 	}
 }
 
+func (bl *baseline) Format(probes []string) string {
+	var line strings.Builder
+	line.WriteString(bl.Url)
+	if bl.Host != "" {
+		line.WriteString(" (" + bl.Host + ")")
+	}
+
+	if bl.Err != nil {
+		line.WriteString("err: ")
+		line.WriteString(bl.Err.Error())
+		return line.String()
+	}
+
+	for _, p := range probes {
+		line.WriteString(" ")
+		line.WriteString(bl.Additional(p))
+	}
+
+	return line.String()
+}
+
 func (bl *baseline) String() string {
 	var line strings.Builder
 	//line.WriteString("[+] ")
@@ -172,6 +193,13 @@ func (bl *baseline) String() string {
 	if bl.Host != "" {
 		line.WriteString(" (" + bl.Host + ")")
 	}
+
+	if bl.Err != nil {
+		line.WriteString("err: ")
+		line.WriteString(bl.Err.Error())
+		return line.String()
+	}
+
 	line.WriteString(" - ")
 	line.WriteString(strconv.Itoa(bl.Status))
 	line.WriteString(" - ")
@@ -182,14 +210,8 @@ func (bl *baseline) String() string {
 		line.WriteString(" ")
 	}
 	line.WriteString(bl.Additional("title"))
-	line.WriteString(bl.Additional("mmh3"))
 	line.WriteString(bl.Frameworks.ToString())
-	//line.WriteString(bl.Extracteds)
-	//line.WriteString("\n")
-	if bl.Err != nil {
-		line.WriteString("err: ")
-		line.WriteString(bl.Err.Error())
-	}
+
 	return line.String()
 }
 

internal/option.go

@@ -13,31 +13,52 @@ import (
 )
 
 type Option struct {
-	URL               string            `short:"u" long:"url"`
-	URLFile           string            `short:"l" long:"list"`
-	Dictionaries      []string          `short:"d" long:"dict"`
-	Word              string            `short:"w" long:"word"`
-	Extensions        string            `short:"e" long:"extension"`
-	ExcludeExtensions string            `long:"exclude-extension"`
-	RemoveExtensions  string            `long:"remove-extension"`
-	Uppercase         bool              `short:"U" long:"uppercase"`
-	Lowercase         bool              `short:"L" long:"lowercase"`
-	Prefixes          []string          `long:"prefix"`
-	Suffixes          []string          `long:"suffix"`
-	Replaces          map[string]string `long:"replace"`
-	Deadline          int               `long:"deadline" default:"600"` // todo 总的超时时间,适配云函数的deadline
-	Timeout           int               `long:"timeout" default:"2"`
-	Headers           []string          `long:"header"`
-	OutputFile        string            `short:"f"`
-	OutputProbe       string            `long:"probe"`
-	Offset            int               `long:"offset"`
-	Limit             int               `long:"limit"`
-	Threads           int               `short:"t" long:"thread" default:"20"`
-	PoolSize          int               `short:"p" long:"pool" default:"5"`
-	Debug             bool              `long:"debug"`
-	Quiet             bool              `short:"q" long:"quiet"`
-	Mod               string            `short:"m" long:"mod" default:"path"`
-	Client            string            `short:"c" long:"client" default:"auto"`
+	InputOptions
+	OutputOptions
+	RequestOptions
+	MiscOptions
+}
+
+type InputOptions struct {
+	URL               string            `short:"u" long:"url" description:"String, input baseurl (separated by commas), e.g.: http://google.com, http://baidu.com"`
+	URLFile           string            `short:"l" long:"list" description:"File, input filename"`
+	Offset            int               `long:"offset" description:"Int, wordlist offset"`
+	Limit             int               `long:"limit" description:"Int, wordlist limit, start with offset. e.g.: --offset 1000 --limit 100"`
+	Dictionaries      []string          `short:"d" long:"dict" description:"Files, dict files, e.g.: -d 1.txt -d 2.txt"`
+	Word              string            `short:"w" long:"word" description:"String, word generate dsl, e.g.: -w test{?ld#4}"`
+	Extensions        string            `short:"e" long:"extension" description:"String, add extensions (separated by commas), e.g.: -e jsp,jspx"`
+	ExcludeExtensions string            `long:"exclude-extension" description:"String, exclude extensions (separated by commas), e.g.: --exclude-extension jsp,jspx"`
+	RemoveExtensions  string            `long:"remove-extension" description:"String, remove extensions (separated by commas), e.g.: --remove-extension jsp,jspx"`
+	Uppercase         bool              `short:"U" long:"uppercase" description:"Bool, upper wordlist, e.g.: --uppercase"`
+	Lowercase         bool              `short:"L" long:"lowercase" description:"Bool, lower wordlist, e.g.: --lowercase"`
+	Prefixes          []string          `long:"prefix" description:"Strings, add prefix, e.g.: --prefix aaa --prefix bbb"`
+	Suffixes          []string          `long:"suffix" description:"Strings, add suffix, e.g.: --suffix aaa --suffix bbb"`
+	Replaces          map[string]string `long:"replace" description:"Strings, replace string, e.g.: --replace aaa:bbb --replace ccc:ddd"`
+}
+
+type OutputOptions struct {
+	Matches     map[string]string `short:"m" long:"match" description:"String, "`
+	Filters     map[string]string `long:"filter" description:"String, "`
+	Extracts    []string          `long:"extract" description:"String, "`
+	OutputFile  string            `short:"f" description:"String, output filename"`
+	OutputProbe string            `long:"probe" description:"String, output format"`
+}
+
+type RequestOptions struct {
+	Headers []string `long:"header"`
+	Method  string   `long:"method"`
+	Cookie  string   `long:"cookie"`
+}
+
+type MiscOptions struct {
+	Deadline int    `long:"deadline" default:"600" 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)"`
+	Debug    bool   `long:"debug" description:"Bool, output debug info"`
+	Quiet    bool   `short:"q" long:"quiet" description:"Bool, Quiet"`
+	Mod      string `short:"m" long:"mod" default:"path" choice:"path" choice:"host" description:"String, path/host spray"`
+	Client   string `short:"c" long:"client" default:"auto" choice:"fast" choice:"standard" choice:"auto" description:"String, Client type"`
 }
 
 func (opt *Option) PrepareRunner() (*Runner, error) {
@@ -52,6 +73,7 @@ func (opt *Option) PrepareRunner() (*Runner, error) {
 		PoolSize: opt.PoolSize,
 		Mod:      opt.Mod,
 		Timeout:  opt.Timeout,
+		Probes:   strings.Split(opt.OutputProbe, ","),
 	}
 
 	err = pkg.LoadTemplates()

internal/pool.go

@@ -48,6 +48,7 @@ func NewPool(ctx context.Context, config *pkg.Config, outputCh chan *baseline) (
 
 			if pool.failedCount > breakThreshold {
 				// 当报错次数超过上限是, 结束任务
+				pool.Recover()
 				pool.cancel()
 			}
 		}
@@ -62,6 +63,7 @@ func NewPool(ctx context.Context, config *pkg.Config, outputCh chan *baseline) (
 
 			if pool.failedCount > breakThreshold {
 				// 当报错次数超过上限是, 结束任务
+				pool.Recover()
 				pool.cancel()
 			}
 		}
@@ -86,7 +88,6 @@ func NewPool(ctx context.Context, config *pkg.Config, outputCh chan *baseline) (
 			pool.failedCount++
 			bl = &baseline{Url: pool.BaseURL + unit.path, Err: reqerr}
 		} else {
-			pool.failedCount = 0 // 如果后续访问正常, 重置错误次数
 			if err = pool.PreCompare(resp); err == nil || unit.source == CheckSource || unit.source == InitSource {
 				// 通过预对比跳过一些无用数据, 减少性能消耗
 				bl = NewBaseline(req.URI(), req.Host(), resp)
@@ -99,14 +100,18 @@ func NewPool(ctx context.Context, config *pkg.Config, outputCh chan *baseline) (
 		case InitSource:
 			pool.base = bl
 			pool.initwg.Done()
-			logs.Log.Important("[init] " + bl.String())
+			logs.Log.Important("[baseline] " + bl.String())
 			return
 		case CheckSource:
-			logs.Log.Debugf("[check] " + bl.String())
 			if bl.Err != nil {
-				logs.Log.Warnf("maybe ip has banned by waf, break (%d/%d), error: %s", pool.failedCount, breakThreshold, bl.Err.Error())
+				logs.Log.Warnf("[check.error] maybe ip had banned by waf, break (%d/%d), error: %s", pool.failedCount, breakThreshold, bl.Err.Error())
+				pool.failedBaselines = append(pool.failedBaselines, bl)
 			} else if pool.base.Compare(bl) < 1 {
-				logs.Log.Warn("maybe trigger risk control")
+				logs.Log.Warn("[check.failed] maybe trigger risk control, " + bl.String())
+				pool.failedBaselines = append(pool.failedBaselines, bl)
+			} else {
+				pool.ResetFailed() // 如果后续访问正常, 重置错误次数
+				logs.Log.Debug("[check.pass] " + bl.String())
 			}
 
 		case WordSource:
@@ -119,9 +124,9 @@ func NewPool(ctx context.Context, config *pkg.Config, outputCh chan *baseline) (
 			} else if pool.failedCount%pool.errPeriod == 0 {
 				go pool.check()
 			}
+			pool.bar.Done()
 		}
 
-		pool.bar.Done()
 		pool.wg.Done()
 	})
 
@@ -138,19 +143,20 @@ type Pool struct {
 	ctx    context.Context
 	cancel context.CancelFunc
 	//baseReq      *http.Request
-	base        *baseline
-	outputCh    chan *baseline // 输出的chan, 全局统一
-	tempCh      chan *baseline // 待处理的baseline
-	reqCount    int
-	failedCount int
-	checkPeriod int
-	errPeriod   int
-	analyzeDone bool
-	genReq      func(s string) (*ihttp.Request, error)
-	check       func()
-	worder      *words.Worder
-	wg          sync.WaitGroup
-	initwg      sync.WaitGroup // 初始化用, 之后改成锁
+	base            *baseline
+	outputCh        chan *baseline // 输出的chan, 全局统一
+	tempCh          chan *baseline // 待处理的baseline
+	reqCount        int
+	failedCount     int
+	checkPeriod     int
+	errPeriod       int
+	failedBaselines []*baseline
+	analyzeDone     bool
+	genReq          func(s string) (*ihttp.Request, error)
+	check           func()
+	worder          *words.Worder
+	wg              sync.WaitGroup
+	initwg          sync.WaitGroup // 初始化用, 之后改成锁
 }
 
 func (p *Pool) Init() error {
@@ -247,10 +253,22 @@ func (p *Pool) comparing() {
 	p.analyzeDone = true
 }
 
+func (p *Pool) ResetFailed() {
+	p.failedCount = 0
+	p.failedBaselines = nil
+}
+
+func (p *Pool) Recover() {
+	logs.Log.Errorf("failed request exceeds the threshold , task will exit. Breakpoint %d", p.reqCount)
+	logs.Log.Error("collecting failed check")
+	for _, bl := range p.failedBaselines {
+		logs.Log.Error(bl.String())
+	}
+}
+
 func (p *Pool) Close() {
 	p.wg.Wait()
 	p.bar.Close()
-
 	close(p.tempCh)
 	for !p.analyzeDone {
 		time.Sleep(time.Duration(100) * time.Millisecond)

internal/runner.go

@@ -25,6 +25,7 @@ type Runner struct {
 	poolwg   sync.WaitGroup
 	Timeout  int
 	Mod      string
+	Probes   []string
 	OutputCh chan *baseline
 	Progress *uiprogress.Progress
 }
@@ -104,7 +105,11 @@ func (r *Runner) Outputting() {
 		select {
 		case bl := <-r.OutputCh:
 			if bl.IsValid {
-				logs.Log.Console("[+] " + bl.String() + "\n")
+				if len(r.Probes) > 0 {
+					logs.Log.Console("[+]" + bl.Format(r.Probes) + "\n")
+				} else {
+					logs.Log.Console("[+] " + bl.String() + "\n")
+				}
 			} else {
 				logs.Log.Debug(bl.String())
 			}