diff --git a/cmd/cmd.go b/cmd/cmd.go index 7d44fd2..44e115f 100644 --- a/cmd/cmd.go +++ b/cmd/cmd.go @@ -113,7 +113,7 @@ func Spray() { } if option.Format != "" { - internal.Format(option.Format, !option.NoColor) + internal.Format(option) return } diff --git a/internal/format.go b/internal/format.go index 59bde37..d15bde8 100644 --- a/internal/format.go +++ b/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 { - logs.Log.Info(result.String()) + + // 分组 + + for _, results := range group { + for _, result := range results { + if !opts.Fuzzy && result.IsFuzzy { + continue + } + if !opts.NoColor { + logs.Log.Info(result.ColorString()) + } else { + logs.Log.Info(result.String()) + } } } } diff --git a/internal/option.go b/internal/option.go index bfe6ada..94adf23 100644 --- a/internal/option.go +++ b/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" ` - 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"` - 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"` + 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"` + 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"` + //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 != "" { - r.FuzzyFile, err = files.NewFile(opt.FuzzyFile, false, false, true) - if err != nil { - return nil, err - } - } else if opt.AutoFile { - r.FuzzyFile, err = files.NewFile("fuzzy.json", false, false, true) - if err != nil { - return nil, err - } - } + //if opt.FuzzyFile != "" { + // r.FuzzyFile, err = files.NewFile(opt.FuzzyFile, false, false, true) + // if err != nil { + // return nil, err + // } + //} else if opt.AutoFile { + // r.FuzzyFile, err = files.NewFile("fuzzy.json", false, false, true) + // if err != nil { + // return nil, err + // } + //} if opt.DumpFile != "" { r.DumpFile, err = files.NewFile(opt.DumpFile, false, false, true) diff --git a/internal/runner.go b/internal/runner.go index 5cd2d9f..3aa4337 100644 --- a/internal/runner.go +++ b/internal/runner.go @@ -47,20 +47,20 @@ type Runner struct { MatchExpr *vm.Program RecursiveExpr *vm.Program OutputFile *files.File - FuzzyFile *files.File - DumpFile *files.File - StatFile *files.File - Progress *mpb.Progress - Fns []func(string) []string - Count int // tasks total number - Wordlist []string - AppendWords []string - RecuDepth int - ClientType int - Probes []string - Total int // wordlist total number - Color bool - Jsonify bool + //FuzzyFile *files.File + DumpFile *files.File + StatFile *files.File + Progress *mpb.Progress + Fns []func(string) []string + Count int // tasks total number + Wordlist []string + AppendWords []string + RecuDepth int + ClientType int + Probes []string + Total int // wordlist total number + Color bool + Jsonify bool } func (r *Runner) PrepareConfig() *pool.Config { @@ -358,62 +358,31 @@ func (r *Runner) saveStat(content string) { } } -func (r *Runner) OutputHandler() { - debugPrint := func(bl *pkg.Baseline) { - if r.Color { - logs.Log.Debug(bl.ColorString()) - } else { - logs.Log.Debug(bl.String()) - } +func (r *Runner) Output(bl *pkg.Baseline) { + var out string + if r.Option.Json { + out = bl.Jsonify() + } else if len(r.Probes) > 0 { + 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(line + "\n") - 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()) - } - } + r.OutputFile.SafeWrite(bl.Jsonify() + "\n") + r.OutputFile.SafeSync() } +} +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() }