From 4e78e55b6ed2fdc794eaed2d8c704a1e5995ccd5 Mon Sep 17 00:00:00 2001 From: M09Ic Date: Thu, 7 Mar 2024 04:15:28 +0800 Subject: [PATCH] fix config load priority, now config < cmd enhance bar print fix config.yaml default --- cmd/cmd.go | 22 ++++++++++++++-------- config.yaml | 14 +++++++------- internal/option.go | 20 ++++++++++---------- internal/runner.go | 42 ++++++++++++++++++++++++++---------------- pkg/bar.go | 35 +++++++++++++++++++++++------------ 5 files changed, 80 insertions(+), 53 deletions(-) diff --git a/cmd/cmd.go b/cmd/cmd.go index 0f833a6..17aea49 100644 --- a/cmd/cmd.go +++ b/cmd/cmd.go @@ -18,9 +18,19 @@ import ( ) var ver = "v0.9.5" +var DefaultConfig = "config.yaml" func Spray() { var option internal.Option + + if files.IsExist(DefaultConfig) { + err := internal.LoadConfig(DefaultConfig, &option) + if err != nil { + logs.Log.Error(err.Error()) + return + } + } + parser := flags.NewParser(&option, flags.Default) parser.Usage = ` @@ -74,14 +84,10 @@ func Spray() { return } if option.Config != "" { - if !files.IsExist(option.Config) { - logs.Log.Warnf("config file %s not found", option.Config) - } else { - err := internal.LoadConfig(option.Config, &option) - if err != nil { - logs.Log.Error(err.Error()) - return - } + err := internal.LoadConfig(option.Config, &option) + if err != nil { + logs.Log.Error(err.Error()) + return } } diff --git a/config.yaml b/config.yaml index d6e89ce..72d45df 100644 --- a/config.yaml +++ b/config.yaml @@ -2,7 +2,7 @@ input: # Files, Multi,dict files, e.g.: -d 1.txt -d 2.txt dictionaries: [] # Bool, no dictionary - no-dict: true + no-dict: false # String, word generate dsl, e.g.: -w test{?ld#4} word: "" # Files, rule files, e.g.: -r rule1.txt -r rule2.txt @@ -55,6 +55,12 @@ output: format: "" # String, output format output_probe: "" + # Bool, Quiet + quiet: false + # Bool, no color + no-color: false + # Bool, No progress bar + no-bar: false plugins: # Bool, enable all plugin all: false @@ -144,11 +150,5 @@ misc: debug: false # Bool, log verbose level ,default 0, level1: -v level2 -vv verbose: [] - # Bool, Quiet - quiet: false - # Bool, no color - no-color: false - # Bool, No progress bar - no-bar: false # String, proxy address, e.g.: --proxy socks5://127.0.0.1:1080 proxy: "" diff --git a/internal/option.go b/internal/option.go index 5ce52c0..b739a38 100644 --- a/internal/option.go +++ b/internal/option.go @@ -43,7 +43,7 @@ type Option struct { type InputOptions struct { ResumeFrom string `long:"resume" description:"File, resume filename" ` - Config string `short:"c" long:"config" default:"config.yaml" description:"File, config filename"` + Config string `short:"c" long:"config" description:"File, config filename"` URL []string `short:"u" long:"url" description:"Strings, input baseurl, e.g.: http://google.com"` URLFile string `short:"l" long:"list" description:"File, input filename"` PortRange string `short:"p" long:"port" description:"String, input port range, e.g.: 80,8080-8090,db"` @@ -85,6 +85,9 @@ type OutputOptions struct { 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"` 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"` + NoBar bool `long:"no-bar" description:"Bool, No progress bar" config:"no-bar"` } type RequestOptions struct { @@ -141,9 +144,6 @@ type MiscOptions struct { Debug bool `long:"debug" description:"Bool, output debug info" config:"debug"` Version bool `long:"version" description:"Bool, show version"` Verbose []bool `short:"v" description:"Bool, log verbose level ,default 0, level1: -v level2 -vv " config:"verbose"` - Quiet bool `short:"q" long:"quiet" description:"Bool, Quiet" config:"quiet"` - NoColor bool `long:"no-color" description:"Bool, no color" config:"no-color"` - NoBar bool `long:"no-bar" description:"Bool, No progress bar" config:"no-bar"` Proxy string `long:"proxy" description:"String, proxy address, e.g.: --proxy socks5://127.0.0.1:1080" config:"proxy"` InitConfig bool `long:"init" description:"Bool, init config file"` } @@ -154,7 +154,6 @@ func (opt *Option) PrepareRunner() (*Runner, error) { return nil, err } r := &Runner{ - Progress: mpb.New(mpb.WithRefreshRate(100 * time.Millisecond)), Threads: opt.Threads, PoolSize: opt.PoolSize, Mod: opt.Mod, @@ -197,6 +196,7 @@ func (opt *Option) PrepareRunner() (*Runner, error) { r.Color = false } if !(opt.Quiet || opt.NoBar) { + r.Progress = mpb.New(mpb.WithRefreshRate(100 * time.Millisecond)) logs.Log.SetOutput(r.Progress) } @@ -334,11 +334,11 @@ func (opt *Option) PrepareRunner() (*Runner, error) { opt.Word += "}" } - if opt.Suffixes != nil { + if len(opt.Suffixes) != 0 { mask.SpecialWords["suffix"] = opt.Suffixes opt.Word += "{@suffix}" } - if opt.Prefixes != nil { + if len(opt.Prefixes) != 0 { mask.SpecialWords["prefix"] = opt.Prefixes opt.Word = "{@prefix}" + opt.Word } @@ -362,7 +362,7 @@ func (opt *Option) PrepareRunner() (*Runner, error) { logs.Log.Logf(pkg.LogVerbose, "Parsed %d words by %s", len(r.Wordlist), opt.Word) } - if opt.Rules != nil { + if len(opt.Rules) != 0 { rules, err := loadRuleAndCombine(opt.Rules) if err != nil { return nil, err @@ -391,7 +391,7 @@ func (opt *Option) PrepareRunner() (*Runner, error) { Total: r.Total, } - if opt.AppendRule != nil { + if len(opt.AppendRule) != 0 { content, err := loadRuleAndCombine(opt.AppendRule) if err != nil { return nil, err @@ -399,7 +399,7 @@ func (opt *Option) PrepareRunner() (*Runner, error) { r.AppendRules = rule.Compile(string(content), "") } - if opt.AppendFile != nil { + if len(opt.AppendFile) != 0 { var bs bytes.Buffer for _, f := range opt.AppendFile { content, err := ioutil.ReadFile(f) diff --git a/internal/runner.go b/internal/runner.go index 23b122d..8700f65 100644 --- a/internal/runner.go +++ b/internal/runner.go @@ -171,21 +171,7 @@ func (r *Runner) Prepare(ctx context.Context) error { }() if r.Count > 0 { - prompt := "total progressive:" - r.bar = r.Progress.AddBar(int64(r.Count), - mpb.BarFillerClearOnComplete(), // 可选：当进度条完成时清除 - mpb.PrependDecorators( - // 显示自定义的信息，比如下载速度和进度 - decor.Name(prompt, decor.WC{W: len(prompt) + 1, C: decor.DindentRight}), // 这里调整了装饰器的参数 - decor.OnComplete( // 当进度完成时显示的文本 - decor.Counters(0, "% d/% d"), " done!", - ), - ), - mpb.AppendDecorators( - // 显示经过的时间 - decor.Elapsed(decor.ET_STYLE_GO, decor.WC{W: 4}), - ), - ) + r.addBar(r.Count) } r.Pools, err = ants.NewPoolWithFunc(r.PoolSize, func(i interface{}) { @@ -339,8 +325,32 @@ Loop: time.Sleep(100 * time.Millisecond) // 延迟100ms, 等所有数据处理完毕 } +func (r *Runner) addBar(total int) { + if r.Progress == nil { + return + } + + prompt := "total progressive:" + r.bar = r.Progress.AddBar(int64(total), + mpb.BarFillerClearOnComplete(), // 可选：当进度条完成时清除 + mpb.PrependDecorators( + // 显示自定义的信息，比如下载速度和进度 + decor.Name(prompt, decor.WC{W: len(prompt) + 1, C: decor.DindentRight}), // 这里调整了装饰器的参数 + decor.OnComplete( // 当进度完成时显示的文本 + decor.Counters(0, "% d/% d"), " done!", + ), + ), + mpb.AppendDecorators( + // 显示经过的时间 + decor.Elapsed(decor.ET_STYLE_GO, decor.WC{W: 4}), + ), + ) +} + func (r *Runner) Done() { - r.bar.Increment() + if r.bar != nil { + r.bar.Increment() + } r.finished++ r.poolwg.Done() } diff --git a/pkg/bar.go b/pkg/bar.go index 98612e3..fe20e11 100644 --- a/pkg/bar.go +++ b/pkg/bar.go @@ -1,26 +1,31 @@ package pkg import ( - "github.com/chainreactors/go-metrics" + "fmt" "github.com/vbauerster/mpb/v8" "github.com/vbauerster/mpb/v8/decor" + "time" ) func NewBar(u string, total int, stat *Statistor, p *mpb.Progress) *Bar { - m := metrics.NewMeter() - metrics.Register(u, m) - - // 在mpb v8中，Name装饰器的使用方式略有不同 + if p == nil { + return &Bar{ + url: u, + } + } bar := p.AddBar(int64(total), mpb.BarFillerClearOnComplete(), mpb.BarRemoveOnComplete(), mpb.PrependDecorators( - // 显示自定义的信息，比如下载速度和进度 decor.Name(u, decor.WC{W: len(u) + 1, C: decor.DindentRight}), // 这里调整了装饰器的参数 - decor.Counters(0, "% d/% d"), + decor.NewAverageSpeed(0, "% .0f/s ", time.Now()), + decor.Counters(0, "%d/%d"), + decor.Any(func(s decor.Statistics) string { + return fmt.Sprintf(" found: %d", stat.FoundNumber) + }), ), mpb.AppendDecorators( - // 显示经过的时间 + decor.Percentage(), decor.Elapsed(decor.ET_STYLE_GO, decor.WC{W: 4}), ), ) @@ -28,23 +33,29 @@ func NewBar(u string, total int, stat *Statistor, p *mpb.Progress) *Bar { return &Bar{ url: u, bar: bar, - m: m, + //m: m, } } type Bar struct { url string bar *mpb.Bar - m metrics.Meter + //m metrics.Meter } func (bar *Bar) Done() { - bar.m.Mark(1) + //bar.m.Mark(1) + if bar.bar == nil { + return + } bar.bar.Increment() } func (bar *Bar) Close() { //metrics.Unregister(bar.url) // 标记进度条为完成状态 - //bar.bar.Abort(false) + if bar.bar == nil { + return + } + bar.bar.Abort(false) }