admin/spray
1
0
Fork 0
mirror of https://github.com/chainreactors/spray.git synced 2025-09-15 11:40:13 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor output and format

Browse Source
This commit is contained in:
M09Ic 2024-08-26 00:04:44 +08:00
parent de168e0be9
commit 1c28898631
4 changed files with 81 additions and 95 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
cmd/cmd.go
View File

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

25
internal/format.go
View File

@ -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 _, results := range group {
for _, result := range results {
if color {
if !opts.Fuzzy && result.IsFuzzy {
continue
}
if !opts.NoColor {
logs.Log.Info(result.ColorString())
} else {
logs.Log.Info(result.String())
}
}
}
}

24
internal/option.go
View File

@ -84,7 +84,7 @@ type OutputOptions struct {
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"`
//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)

79
internal/runner.go
View File

@ -47,7 +47,7 @@ type Runner struct {
MatchExpr *vm.Program
RecursiveExpr *vm.Program
OutputFile *files.File
FuzzyFile *files.File
//FuzzyFile *files.File
DumpFile *files.File
StatFile *files.File
Progress *mpb.Progress
@ -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())
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 {
logs.Log.Debug(bl.String())
}
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.SafeWrite(bl.Jsonify() + "\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())
}
}
}
}
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()
}