From eb328a4b5238bb8d8b1eebc4ebcbf2d153b4abd3 Mon Sep 17 00:00:00 2001 From: M09Ic Date: Thu, 4 Jul 2024 15:18:46 +0800 Subject: [PATCH] clean output code enhance output format and color support jsonline output in stdout --- go.mod | 2 +- go.sum | 4 + internal/option.go | 38 ++------ internal/runner.go | 212 ++++++++++++++++++++------------------------- 4 files changed, 107 insertions(+), 149 deletions(-) diff --git a/go.mod b/go.mod index 47ebcaa..a698c0a 100644 --- a/go.mod +++ b/go.mod @@ -8,7 +8,7 @@ require ( github.com/chainreactors/files v0.0.0-20231123083421-cea5b4ad18a8 github.com/chainreactors/fingers v0.0.0-20240704063230-de8fec05ff8b github.com/chainreactors/logs v0.0.0-20240207121836-c946f072f81f - github.com/chainreactors/parsers v0.0.0-20240704062910-decf861def9e + github.com/chainreactors/parsers v0.0.0-20240704071623-9d0ee90230a6 github.com/chainreactors/utils v0.0.0-20240704062557-662d623b74f4 github.com/chainreactors/words v0.4.1-0.20240510105042-5ba5c2edc508 github.com/expr-lang/expr v1.16.9 diff --git a/go.sum b/go.sum index 63447bc..0c56205 100644 --- a/go.sum +++ b/go.sum @@ -119,6 +119,10 @@ github.com/chainreactors/parsers v0.0.0-20240702104902-1ce563b7ef76 h1:i4sHuonM5 github.com/chainreactors/parsers v0.0.0-20240702104902-1ce563b7ef76/go.mod h1:G/XLE5RAaUdqADkbhQ59mPrUAbsJLiQ2DN6CwtwNpBQ= github.com/chainreactors/parsers v0.0.0-20240704062910-decf861def9e h1:42ILX5kS76M1D9IQvXgfelpgUJDi/K+4/egE0tLzuSE= github.com/chainreactors/parsers v0.0.0-20240704062910-decf861def9e/go.mod h1:7rXdYz6jrdjF0WUH1ICcAXKIKKjKmJo2PU8u43V7jkA= +github.com/chainreactors/parsers v0.0.0-20240704071443-8a8558f34cf9 h1:XxPUVhP29vnbLuhxFt8VT3eyBR8d/GfHR7YK44zyDVo= +github.com/chainreactors/parsers v0.0.0-20240704071443-8a8558f34cf9/go.mod h1:7rXdYz6jrdjF0WUH1ICcAXKIKKjKmJo2PU8u43V7jkA= +github.com/chainreactors/parsers v0.0.0-20240704071623-9d0ee90230a6 h1:jUxPo0RJ/f+/4x3ydeXqCeMq5VbvYBjtmpBePWFfNc8= +github.com/chainreactors/parsers v0.0.0-20240704071623-9d0ee90230a6/go.mod h1:7rXdYz6jrdjF0WUH1ICcAXKIKKjKmJo2PU8u43V7jkA= github.com/chainreactors/utils v0.0.0-20240302165634-2b8494c9cfc3/go.mod h1:JA4eiQZm+7AsfjXBcIzIdVKBEhDCb16eNtWFCGTxlvs= github.com/chainreactors/utils v0.0.0-20240528085651-ba1b255482c1 h1:+awuysRKLmdLQbVK+HPSOGvO3dFGdNSbM2jyLh+VYOA= github.com/chainreactors/utils v0.0.0-20240528085651-ba1b255482c1/go.mod h1:JA4eiQZm+7AsfjXBcIzIdVKBEhDCb16eNtWFCGTxlvs= diff --git a/internal/option.go b/internal/option.go index c82e6a7..d038e50 100644 --- a/internal/option.go +++ b/internal/option.go @@ -87,6 +87,7 @@ type OutputOptions struct { 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"` Format string `short:"F" long:"format" description:"String, output format, e.g.: --format 1.json" config:"format"` + Json bool `short:"j" long:"json" description:"Bool, output json" config:"json"` OutputProbe string `short:"o" long:"probe" description:"String, output format" config:"output_probe"` Quiet bool `short:"q" long:"quiet" description:"Bool, Quiet" config:"quiet"` NoColor bool `long:"no-color" description:"Bool, no color" config:"no-color"` @@ -158,36 +159,13 @@ func (opt *Option) PrepareRunner() (*Runner, error) { return nil, err } r := &Runner{ - Threads: opt.Threads, - PoolSize: opt.PoolSize, - Mod: opt.Mod, - Timeout: opt.Timeout, - RateLimit: opt.RateLimit, - Deadline: opt.Deadline, - Headers: make(map[string]string), - Method: opt.Method, - Offset: opt.Offset, - Total: opt.Limit, - taskCh: make(chan *Task), - outputCh: make(chan *pkg.Baseline, 256), - outwg: &sync.WaitGroup{}, - fuzzyCh: make(chan *pkg.Baseline, 256), - Fuzzy: opt.Fuzzy, - Force: opt.Force, - CheckOnly: opt.CheckOnly, - CheckPeriod: opt.CheckPeriod, - ErrPeriod: opt.ErrPeriod, - BreakThreshold: opt.BreakThreshold, - Crawl: opt.Crawl, - Scope: opt.Scope, - Finger: opt.Finger, - Bak: opt.Bak, - Common: opt.Common, - RetryCount: opt.RetryCount, - RandomUserAgent: opt.RandomUserAgent, - Random: opt.Random, - Index: opt.Index, - Proxy: opt.Proxy, + Option: opt, + taskCh: make(chan *Task), + outputCh: make(chan *pkg.Baseline, 256), + outwg: &sync.WaitGroup{}, + fuzzyCh: make(chan *pkg.Baseline, 256), + Headers: make(map[string]string), + Total: opt.Limit, } // log and bar diff --git a/internal/runner.go b/internal/runner.go index 1d0e59d..ff9fe1e 100644 --- a/internal/runner.go +++ b/internal/runner.go @@ -28,86 +28,62 @@ var ( ) type Runner struct { - taskCh chan *Task - poolwg sync.WaitGroup - outwg *sync.WaitGroup - outputCh chan *pkg.Baseline - fuzzyCh chan *pkg.Baseline - bar *mpb.Bar - finished int + *Option - Tasks chan *Task - Count int // tasks total number - Wordlist []string - Rules *rule.Program - AppendRules *rule.Program - AppendWords []string - Headers map[string]string - Method string - Fns []func(string) []string - FilterExpr *vm.Program - MatchExpr *vm.Program - RecursiveExpr *vm.Program - RecuDepth int - Threads int - PoolSize int - ClientType int - Pools *ants.PoolWithFunc - PoolName map[string]bool - Timeout int - Mod string - Probes []string - Fuzzy bool - OutputFile *files.File - FuzzyFile *files.File - DumpFile *files.File - StatFile *files.File - Progress *mpb.Progress - Offset int - Limit int - RateLimit int - Total int // wordlist total number - Deadline int - CheckPeriod int - ErrPeriod int - BreakThreshold int - Color bool - CheckOnly bool - Force bool - IgnoreWaf bool - Crawl bool - Scope []string - Finger bool - Bak bool - Common bool - RetryCount int - RandomUserAgent bool - Random string - Index string - Proxy string + taskCh chan *Task + poolwg sync.WaitGroup + outwg *sync.WaitGroup + outputCh chan *pkg.Baseline + fuzzyCh chan *pkg.Baseline + bar *mpb.Bar + finished int + Pools *ants.PoolWithFunc + PoolName map[string]bool + Tasks chan *Task + Rules *rule.Program + AppendRules *rule.Program + Headers map[string]string + FilterExpr *vm.Program + 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 } func (r *Runner) PrepareConfig() *pool.Config { config := &pool.Config{ - Thread: r.Threads, - Timeout: r.Timeout, - RateLimit: r.RateLimit, - Headers: r.Headers, - Method: r.Method, - Mod: pool.ModMap[r.Mod], - OutputCh: r.outputCh, - FuzzyCh: r.fuzzyCh, - OutLocker: r.outwg, - Fuzzy: r.Fuzzy, - CheckPeriod: r.CheckPeriod, - ErrPeriod: int32(r.ErrPeriod), - BreakThreshold: int32(r.BreakThreshold), - MatchExpr: r.MatchExpr, - FilterExpr: r.FilterExpr, - RecuExpr: r.RecursiveExpr, - AppendRule: r.AppendRules, - AppendWords: r.AppendWords, - IgnoreWaf: r.IgnoreWaf, + Thread: r.Threads, + Timeout: r.Timeout, + RateLimit: r.RateLimit, + Headers: r.Headers, + Method: r.Method, + Mod: pool.ModMap[r.Mod], + OutputCh: r.outputCh, + FuzzyCh: r.fuzzyCh, + OutLocker: r.outwg, + Fuzzy: r.Fuzzy, + CheckPeriod: r.CheckPeriod, + ErrPeriod: int32(r.ErrPeriod), + BreakThreshold: int32(r.BreakThreshold), + MatchExpr: r.MatchExpr, + FilterExpr: r.FilterExpr, + RecuExpr: r.RecursiveExpr, + AppendRule: r.AppendRules, + AppendWords: r.AppendWords, + //IgnoreWaf: r.IgnoreWaf, Crawl: r.Crawl, Scope: r.Scope, Active: r.Finger, @@ -386,38 +362,55 @@ func (r *Runner) OutputHandler() { logs.Log.Debug(bl.String()) } } - go func() { - var saveFunc func(*pkg.Baseline) - if r.OutputFile != nil { - saveFunc = func(bl *pkg.Baseline) { - r.OutputFile.SafeWrite(bl.Jsonify() + "\n") - r.OutputFile.SafeSync() + 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 { - if r.Color { - saveFunc = func(bl *pkg.Baseline) { - logs.Log.Console(logs.GreenBold("[+] " + bl.Format(r.Probes) + "\n")) - } - } else { - saveFunc = func(bl *pkg.Baseline) { - logs.Log.Console("[+] " + bl.Format(r.Probes) + "\n") - } - } + outFunc(bl.Format(r.Probes)) } else { - if r.Color { - saveFunc = func(bl *pkg.Baseline) { - logs.Log.Console(logs.GreenBold("[+] " + bl.ColorString() + "\n")) - } - } else { - saveFunc = func(bl *pkg.Baseline) { - logs.Log.Console("[+] " + bl.String() + "\n") - } - } + outFunc(bl.String()) } } + } + go func() { for { select { case bl, ok := <-r.outputCh: @@ -429,7 +422,7 @@ func (r *Runner) OutputHandler() { r.DumpFile.SafeSync() } if bl.IsValid { - saveFunc(bl) + outputPrint(bl) if bl.Recu { r.AddRecursive(bl) } @@ -442,23 +435,6 @@ func (r *Runner) OutputHandler() { }() go func() { - var fuzzySaveFunc func(*pkg.Baseline) - if r.FuzzyFile != nil { - fuzzySaveFunc = func(bl *pkg.Baseline) { - r.FuzzyFile.SafeWrite(bl.Jsonify() + "\n") - } - } else { - if r.Color { - fuzzySaveFunc = func(bl *pkg.Baseline) { - logs.Log.Console(logs.GreenBold("[fuzzy] " + bl.ColorString() + "\n")) - } - } else { - fuzzySaveFunc = func(bl *pkg.Baseline) { - logs.Log.Console("[fuzzy] " + bl.String() + "\n") - } - } - } - for { select { case bl, ok := <-r.fuzzyCh: @@ -466,7 +442,7 @@ func (r *Runner) OutputHandler() { return } if r.Fuzzy { - fuzzySaveFunc(bl) + outputPrint(bl) } r.outwg.Done() }