refactor output and format

cmd/cmd.go

@@ -113,7 +113,7 @@ func Spray() {
 	}
 
 	if option.Format != "" {
-		internal.Format(option.Format, !option.NoColor)
+		internal.Format(option)
 		return
 	}
 

internal/format.go

@@ -6,22 +6,23 @@ import (
 	"github.com/chainreactors/logs"
 	"github.com/chainreactors/spray/pkg"
 	"io"
+	"net/url"
 	"os"
 )
 
-func Format(filename string, color bool) {
+func Format(opts Option) {
 	var content []byte
 	var err error
-	if filename == "stdin" {
+	if opts.Format == "stdin" {
 		content, err = io.ReadAll(os.Stdin)
 	} else {
-		content, err = os.ReadFile(filename)
+		content, err = os.ReadFile(opts.Format)
 	}
 
 	if err != nil {
 		return
 	}
-	var results []*pkg.Baseline
+	group := make(map[string][]*pkg.Baseline)
 	for _, line := range bytes.Split(bytes.TrimSpace(content), []byte("\n")) {
 		var result pkg.Baseline
 		err := json.Unmarshal(line, &result)
@@ -29,13 +30,25 @@ func Format(filename string, color bool) {
 			logs.Log.Error(err.Error())
 			return
 		}
-		results = append(results, &result)
+		result.Url, err = url.Parse(result.UrlString)
+		if err != nil {
+			continue
+		}
+		group[result.Url.Host] = append(group[result.Url.Host], &result)
 	}
-	for _, result := range results {
+
-		if color {
+	// 分组
-			logs.Log.Info(result.ColorString())
+
-		} else {
+	for _, results := range group {
-			logs.Log.Info(result.String())
+		for _, result := range results {
+			if !opts.Fuzzy && result.IsFuzzy {
+				continue
+			}
+			if !opts.NoColor {
+				logs.Log.Info(result.ColorString())
+			} else {
+				logs.Log.Info(result.String())
+			}
 		}
 	}
 }

internal/option.go

@@ -80,11 +80,11 @@ type FunctionOptions struct {
 }
 
 type OutputOptions struct {
-	Match       string `long:"match" description:"String, custom match function, e.g.: --match 'current.Status != 200''" config:"match" `
+	Match      string `long:"match" description:"String, custom match function, e.g.: --match 'current.Status != 200''" config:"match" `
-	Filter      string `long:"filter" description:"String, custom filter function, e.g.: --filter 'current.Body contains \"hello\"'" config:"filter"`
+	Filter     string `long:"filter" description:"String, custom filter function, e.g.: --filter 'current.Body contains \"hello\"'" config:"filter"`
-	Fuzzy       bool   `long:"fuzzy" description:"String, open fuzzy output" config:"fuzzy"`
+	Fuzzy      bool   `long:"fuzzy" description:"String, open fuzzy output" config:"fuzzy"`
-	OutputFile  string `short:"f" long:"file" description:"String, output filename" json:"output_file,omitempty" config:"output-file"`
+	OutputFile string `short:"f" long:"file" description:"String, output filename" json:"output_file,omitempty" config:"output-file"`
-	FuzzyFile   string `long:"fuzzy-file" description:"String, fuzzy output filename" json:"fuzzy_file,omitempty" config:"fuzzy-file"`
+	//FuzzyFile   string `long:"fuzzy-file" description:"String, fuzzy output filename" json:"fuzzy_file,omitempty" config:"fuzzy-file"`
 	DumpFile    string `long:"dump-file" description:"String, dump all request, and write to filename" config:"dump-file"`
 	Dump        bool   `long:"dump" description:"Bool, dump all request" config:"dump"`
 	AutoFile    bool   `long:"auto-file" description:"Bool, auto generator output and fuzzy filename" config:"auto-file"`
@@ -450,17 +450,17 @@ func (opt *Option) NewRunner() (*Runner, error) {
 		}
 	}
 
-	if opt.FuzzyFile != "" {
+	//if opt.FuzzyFile != "" {
-		r.FuzzyFile, err = files.NewFile(opt.FuzzyFile, false, false, true)
+	//	r.FuzzyFile, err = files.NewFile(opt.FuzzyFile, false, false, true)
-		if err != nil {
+	//	if err != nil {
-			return nil, err
+	//		return nil, err
-		}
+	//	}
-	} else if opt.AutoFile {
+	//} else if opt.AutoFile {
-		r.FuzzyFile, err = files.NewFile("fuzzy.json", false, false, true)
+	//	r.FuzzyFile, err = files.NewFile("fuzzy.json", false, false, true)
-		if err != nil {
+	//	if err != nil {
-			return nil, err
+	//		return nil, err
-		}
+	//	}
-	}
+	//}
 
 	if opt.DumpFile != "" {
 		r.DumpFile, err = files.NewFile(opt.DumpFile, false, false, true)

internal/runner.go

@@ -47,20 +47,20 @@ type Runner struct {
 	MatchExpr     *vm.Program
 	RecursiveExpr *vm.Program
 	OutputFile    *files.File
-	FuzzyFile     *files.File
+	//FuzzyFile     *files.File
-	DumpFile      *files.File
+	DumpFile    *files.File
-	StatFile      *files.File
+	StatFile    *files.File
-	Progress      *mpb.Progress
+	Progress    *mpb.Progress
-	Fns           []func(string) []string
+	Fns         []func(string) []string
-	Count         int // tasks total number
+	Count       int // tasks total number
-	Wordlist      []string
+	Wordlist    []string
-	AppendWords   []string
+	AppendWords []string
-	RecuDepth     int
+	RecuDepth   int
-	ClientType    int
+	ClientType  int
-	Probes        []string
+	Probes      []string
-	Total         int // wordlist total number
+	Total       int // wordlist total number
-	Color         bool
+	Color       bool
-	Jsonify       bool
+	Jsonify     bool
 }
 
 func (r *Runner) PrepareConfig() *pool.Config {
@@ -358,62 +358,31 @@ func (r *Runner) saveStat(content string) {
 	}
 }
 
-func (r *Runner) OutputHandler() {
+func (r *Runner) Output(bl *pkg.Baseline) {
-	debugPrint := func(bl *pkg.Baseline) {
+	var out string
-		if r.Color {
+	if r.Option.Json {
-			logs.Log.Debug(bl.ColorString())
+		out = bl.Jsonify()
-		} else {
+	} else if len(r.Probes) > 0 {
-			logs.Log.Debug(bl.String())
+		out = bl.Format(r.Probes)
-		}
+	} else if r.Color {
+		out = bl.ColorString()
+	} else {
+		out = bl.String()
+	}
+
+	if bl.IsFuzzy {
+		logs.Log.Console("[fuzzy] " + out + "\n")
+	} else {
+		logs.Log.Console(out + "\n")
 	}
 
-	var saveFunc func(string)
 	if r.OutputFile != nil {
-		saveFunc = func(line string) {
+		r.OutputFile.SafeWrite(bl.Jsonify() + "\n")
-			r.OutputFile.SafeWrite(line + "\n")
+		r.OutputFile.SafeSync()
-			r.OutputFile.SafeSync()
-		}
-	} else {
-		saveFunc = func(line string) {
-			logs.Log.Console(line + "\n")
-		}
-	}
-
-	var fuzzySaveFunc func(string)
-	if r.FuzzyFile != nil {
-		fuzzySaveFunc = func(line string) {
-			r.FuzzyFile.SafeWrite(line + "\n")
-			r.FuzzyFile.SafeSync()
-		}
-	} else {
-		fuzzySaveFunc = func(line string) {
-			logs.Log.Console("[fuzzy] " + line + "\n")
-		}
-	}
-	outputPrint := func(bl *pkg.Baseline) {
-		var outFunc func(string)
-		if bl.IsFuzzy {
-			outFunc = fuzzySaveFunc
-		} else {
-			outFunc = saveFunc
-		}
-		if r.Option.Json {
-			outFunc(bl.Jsonify())
-		} else if r.Color {
-			if len(r.Probes) > 0 {
-				outFunc(logs.GreenBold(bl.Format(r.Probes)))
-			} else {
-				outFunc(logs.GreenBold(bl.ColorString()))
-			}
-		} else {
-			if len(r.Probes) > 0 {
-				outFunc(bl.Format(r.Probes))
-			} else {
-				outFunc(bl.String())
-			}
-		}
 	}
+}
 
+func (r *Runner) OutputHandler() {
 	go func() {
 		for {
 			select {
@@ -426,12 +395,16 @@ func (r *Runner) OutputHandler() {
 					r.DumpFile.SafeSync()
 				}
 				if bl.IsValid {
-					outputPrint(bl)
+					r.Output(bl)
 					if bl.Recu {
 						r.AddRecursive(bl)
 					}
 				} else {
-					debugPrint(bl)
+					if r.Color {
+						logs.Log.Debug(bl.ColorString())
+					} else {
+						logs.Log.Debug(bl.String())
+					}
 				}
 				r.outwg.Done()
 			}
@@ -446,7 +419,7 @@ func (r *Runner) OutputHandler() {
 					return
 				}
 				if r.Fuzzy {
-					outputPrint(bl)
+					r.Output(bl)
 				}
 				r.outwg.Done()
 			}