From 9e82bb1ab397e3a373e83442c9ddbd1732ef2625 Mon Sep 17 00:00:00 2001 From: M09Ic Date: Thu, 7 Mar 2024 00:24:30 +0800 Subject: [PATCH] use mpb replace uiprogress --- go.mod | 7 +++-- go.sum | 15 +++++++--- internal/option.go | 8 +++--- internal/pool/brutepool.go | 2 +- internal/runner.go | 35 ++++++++++++++--------- pkg/bar.go | 57 +++++++++++++++++++++----------------- 6 files changed, 74 insertions(+), 50 deletions(-) diff --git a/go.mod b/go.mod index 46cf2ab..ce986df 100644 --- a/go.mod +++ b/go.mod @@ -14,10 +14,10 @@ require ( github.com/chainreactors/words v0.4.1-0.20240220104223-153f52e53f37 github.com/goccy/go-yaml v1.11.2 github.com/gookit/config/v2 v2.2.5 - github.com/gosuri/uiprogress v0.0.1 github.com/jessevdk/go-flags v1.5.0 github.com/panjf2000/ants/v2 v2.7.0 github.com/valyala/fasthttp v1.43.0 + github.com/vbauerster/mpb/v8 v8.7.2 golang.org/x/net v0.21.0 golang.org/x/time v0.3.0 sigs.k8s.io/yaml v1.4.0 @@ -25,6 +25,8 @@ require ( require ( dario.cat/mergo v1.0.0 // indirect + github.com/VividCortex/ewma v1.2.0 // indirect + github.com/acarl005/stripansi v0.0.0-20180116102854-5a71ef0e047d // indirect github.com/andybalholm/brotli v1.1.0 // indirect github.com/fatih/color v1.15.0 // indirect github.com/go-dedup/megophone v0.0.0-20170830025436-f01be21026f5 // indirect @@ -33,11 +35,12 @@ require ( github.com/go-playground/validator/v10 v10.14.1 // indirect github.com/gookit/color v1.5.4 // indirect github.com/gookit/goutil v0.6.15 // indirect - github.com/gosuri/uilive v0.0.4 // indirect github.com/klauspost/compress v1.17.4 // indirect github.com/mattn/go-colorable v0.1.13 // indirect github.com/mattn/go-isatty v0.0.19 // indirect + github.com/mattn/go-runewidth v0.0.15 // indirect github.com/mitchellh/mapstructure v1.5.0 // indirect + github.com/rivo/uniseg v0.4.4 // indirect github.com/twmb/murmur3 v1.1.8 // indirect github.com/valyala/bytebufferpool v1.0.0 // indirect github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e // indirect diff --git a/go.sum b/go.sum index aa8c2c2..9453d17 100644 --- a/go.sum +++ b/go.sum @@ -1,5 +1,9 @@ dario.cat/mergo v1.0.0 h1:AGCNq9Evsj31mOgNPcLyXc+4PNABt905YmuqPYYpBWk= dario.cat/mergo v1.0.0/go.mod h1:uNxQE+84aUszobStD9th8a29P2fMDhsBdgRYvZOxGmk= +github.com/VividCortex/ewma v1.2.0 h1:f58SaIzcDXrSy3kWaHNvuJgJ3Nmz59Zji6XoJR/q1ow= +github.com/VividCortex/ewma v1.2.0/go.mod h1:nz4BbCtbLyFDeC9SUHbtcT5644juEuWfUAUnGx7j5l4= +github.com/acarl005/stripansi v0.0.0-20180116102854-5a71ef0e047d h1:licZJFw2RwpHMqeKTCYkitsPqHNxTmd4SNR5r94FGM8= +github.com/acarl005/stripansi v0.0.0-20180116102854-5a71ef0e047d/go.mod h1:asat636LX7Bqt5lYEZ27JNDcqxfjdBQuJ/MM4CN/Lzo= github.com/andybalholm/brotli v1.0.4/go.mod h1:fO7iG3H7G2nSZ7m0zPUDn85XEX2GTukHGRSepvi9Eig= github.com/andybalholm/brotli v1.1.0 h1:eLKJA0d02Lf0mVpIDgYnqXcUn0GqVmEFny3VuID1U3M= github.com/andybalholm/brotli v1.1.0/go.mod h1:sms7XGricyQI9K10gOSf56VKKWS4oLer58Q+mhRPtnY= @@ -53,10 +57,6 @@ github.com/gookit/goutil v0.6.15 h1:mMQ0ElojNZoyPD0eVROk5QXJPh2uKR4g06slgPDF5Jo= github.com/gookit/goutil v0.6.15/go.mod h1:qdKdYEHQdEtyH+4fNdQNZfJHhI0jUZzHxQVAV3DaMDY= github.com/gookit/ini/v2 v2.2.3 h1:nSbN+x9OfQPcMObTFP+XuHt8ev6ndv/fWWqxFhPMu2E= github.com/gookit/ini/v2 v2.2.3/go.mod h1:Vu6p7P7xcfmb8KYu3L0ek8bqu/Im63N81q208SCCZY4= -github.com/gosuri/uilive v0.0.4 h1:hUEBpQDj8D8jXgtCdBu7sWsy5sbW/5GhuO8KBwJ2jyY= -github.com/gosuri/uilive v0.0.4/go.mod h1:V/epo5LjjlDE5RJUcqx8dbw+zc93y5Ya3yg8tfZ74VI= -github.com/gosuri/uiprogress v0.0.1 h1:0kpv/XY/qTmFWl/SkaJykZXrBBzwwadmW8fRb7RJSxw= -github.com/gosuri/uiprogress v0.0.1/go.mod h1:C1RTYn4Sc7iEyf6j8ft5dyoZ4212h8G1ol9QQluh5+0= github.com/jessevdk/go-flags v1.5.0 h1:1jKYvbxEjfUl0fmqTCOfonvskHHXMjBySTLW4y9LFvc= github.com/jessevdk/go-flags v1.5.0/go.mod h1:Fw0T6WPc1dYxT4mKEZRfG5kJhaTDP9pj1c2EWnYs/m4= github.com/klauspost/compress v1.15.9/go.mod h1:PhcZ0MbTNciWF3rruxRgKxI5NkcHHrHUDtV4Yw2GlzU= @@ -69,12 +69,17 @@ github.com/mattn/go-colorable v0.1.13/go.mod h1:7S9/ev0klgBDR4GtXTXX8a3vIGJpMovk github.com/mattn/go-isatty v0.0.16/go.mod h1:kYGgaQfpe5nmfYZH+SKPsOc2e4SrIfOl2e/yFXSvRLM= github.com/mattn/go-isatty v0.0.19 h1:JITubQf0MOLdlGRuRq+jtsDlekdYPia9ZFsB8h/APPA= github.com/mattn/go-isatty v0.0.19/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y= +github.com/mattn/go-runewidth v0.0.15 h1:UNAjwbU9l54TA3KzvqLGxwWjHmMgBUVhBiTjelZgg3U= +github.com/mattn/go-runewidth v0.0.15/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w= github.com/mitchellh/mapstructure v1.5.0 h1:jeMsZIYE/09sWLaz43PL7Gy6RuMjD2eJVyuac5Z2hdY= github.com/mitchellh/mapstructure v1.5.0/go.mod h1:bFUtVrKA4DC2yAKiSyO/QUcy7e+RRV2QTWOzhPopBRo= github.com/panjf2000/ants/v2 v2.7.0 h1:Y3Bgpfo9HDkBoHNVFbMfY5mAvi5TAA17y3HbzQ74p5Y= github.com/panjf2000/ants/v2 v2.7.0/go.mod h1:KIBmYG9QQX5U2qzFP/yQJaq/nSb6rahS9iEHkrCMgM8= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/rivo/uniseg v0.2.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc= +github.com/rivo/uniseg v0.4.4 h1:8TfxU8dW6PdqD27gjM8MVNuicgxIjxpm4K7x4jp8sis= +github.com/rivo/uniseg v0.4.4/go.mod h1:FN3SvrM+Zdj16jyLfmOkMNblXMcoc8DfTHruCPUcx88= github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw= github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo= @@ -92,6 +97,8 @@ github.com/valyala/bytebufferpool v1.0.0/go.mod h1:6bBcMArwyJ5K/AmCkWv1jt77kVWyC github.com/valyala/fasthttp v1.43.0 h1:Gy4sb32C98fbzVWZlTM1oTMdLWGyvxR03VhM6cBIU4g= github.com/valyala/fasthttp v1.43.0/go.mod h1:f6VbjjoI3z1NDOZOv17o6RvtRSWxC77seBFc2uWtgiY= github.com/valyala/tcplisten v1.0.0/go.mod h1:T0xQ8SeCZGxckz9qRXTfG43PvQ/mcWh7FwZEA7Ioqkc= +github.com/vbauerster/mpb/v8 v8.7.2 h1:SMJtxhNho1MV3OuFgS1DAzhANN1Ejc5Ct+0iSaIkB14= +github.com/vbauerster/mpb/v8 v8.7.2/go.mod h1:ZFnrjzspgDHoxYLGvxIruiNk73GNTPG4YHgVNpR10VY= github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e h1:JVG44RsyaB9T2KIHavMF/ppJZNG9ZpyihvCd0w101no= github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e/go.mod h1:RbqR21r5mrJuqunuUZ/Dhy/avygyECGrLceyNeo4LiM= github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY= diff --git a/internal/option.go b/internal/option.go index f2a200f..c1daa4c 100644 --- a/internal/option.go +++ b/internal/option.go @@ -15,7 +15,7 @@ import ( "github.com/chainreactors/utils/iutils" "github.com/chainreactors/words/mask" "github.com/chainreactors/words/rule" - "github.com/gosuri/uiprogress" + "github.com/vbauerster/mpb/v8" "io/ioutil" "net/url" "os" @@ -23,6 +23,7 @@ import ( "strconv" "strings" "sync" + "time" ) var ( @@ -152,7 +153,7 @@ func (opt *Option) PrepareRunner() (*Runner, error) { return nil, err } r := &Runner{ - Progress: uiprogress.New(), + Progress: mpb.New(mpb.WithRefreshRate(100 * time.Millisecond)), Threads: opt.Threads, PoolSize: opt.PoolSize, Mod: opt.Mod, @@ -195,8 +196,7 @@ func (opt *Option) PrepareRunner() (*Runner, error) { r.Color = false } if !(opt.Quiet || opt.NoBar) { - r.Progress.Start() - logs.Log.SetOutput(r.Progress.Bypass()) + logs.Log.SetOutput(r.Progress) } // configuration diff --git a/internal/pool/brutepool.go b/internal/pool/brutepool.go index 2634d88..3b5ccee 100644 --- a/internal/pool/brutepool.go +++ b/internal/pool/brutepool.go @@ -237,7 +237,7 @@ Loop: pool.Statistor.End++ if w == "" { pool.Statistor.Skipped++ - continue + pool.Bar.Done() } pool.wordOffset++ diff --git a/internal/runner.go b/internal/runner.go index 56c4eca..23b122d 100644 --- a/internal/runner.go +++ b/internal/runner.go @@ -2,7 +2,6 @@ package internal import ( "context" - "fmt" "github.com/antonmedv/expr/vm" "github.com/chainreactors/files" "github.com/chainreactors/logs" @@ -11,8 +10,9 @@ import ( "github.com/chainreactors/spray/pkg" "github.com/chainreactors/words" "github.com/chainreactors/words/rule" - "github.com/gosuri/uiprogress" "github.com/panjf2000/ants/v2" + "github.com/vbauerster/mpb/v8" + "github.com/vbauerster/mpb/v8/decor" "sync" "time" ) @@ -33,7 +33,7 @@ type Runner struct { outwg *sync.WaitGroup outputCh chan *pkg.Baseline fuzzyCh chan *pkg.Baseline - bar *uiprogress.Bar + bar *mpb.Bar finished int Tasks chan *Task @@ -61,7 +61,7 @@ type Runner struct { FuzzyFile *files.File DumpFile *files.File StatFile *files.File - Progress *uiprogress.Progress + Progress *mpb.Progress Offset int Limit int RateLimit int @@ -157,7 +157,7 @@ func (r *Runner) Prepare(ctx context.Context) error { }() pool.Worder = words.NewWorderWithChan(ch) pool.Worder.Fns = r.Fns - pool.Bar = pkg.NewBar("check", r.Count-r.Offset, r.Progress) + pool.Bar = pkg.NewBar("check", r.Count-r.Offset, pool.Statistor, r.Progress) pool.Run(ctx, r.Offset, r.Count) r.poolwg.Done() }) @@ -171,12 +171,21 @@ func (r *Runner) Prepare(ctx context.Context) error { }() if r.Count > 0 { - r.bar = r.Progress.AddBar(r.Count) - r.bar.PrependCompleted() - r.bar.PrependFunc(func(b *uiprogress.Bar) string { - return fmt.Sprintf("total progressive: %d/%d ", r.finished, r.Count) - }) - r.bar.AppendElapsed() + 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.Pools, err = ants.NewPoolWithFunc(r.PoolSize, func(i interface{}) { @@ -219,7 +228,7 @@ func (r *Runner) Prepare(ctx context.Context) error { } else { limit = pool.Statistor.Total } - pool.Bar = pkg.NewBar(config.BaseURL, limit-pool.Statistor.Offset, r.Progress) + pool.Bar = pkg.NewBar(config.BaseURL, limit-pool.Statistor.Offset, pool.Statistor, r.Progress) logs.Log.Importantf("[pool] task: %s, total %d words, %d threads, proxy: %s", pool.BaseURL, limit-pool.Statistor.Offset, pool.Thread, pool.ProxyAddr) err = pool.Init() if err != nil { @@ -331,7 +340,7 @@ Loop: } func (r *Runner) Done() { - r.bar.Incr() + r.bar.Increment() r.finished++ r.poolwg.Done() } diff --git a/pkg/bar.go b/pkg/bar.go index ac335e2..98612e3 100644 --- a/pkg/bar.go +++ b/pkg/bar.go @@ -1,45 +1,50 @@ package pkg import ( - "fmt" "github.com/chainreactors/go-metrics" - "github.com/gosuri/uiprogress" + "github.com/vbauerster/mpb/v8" + "github.com/vbauerster/mpb/v8/decor" ) -func NewBar(u string, total int, progress *uiprogress.Progress) *Bar { - bar := &Bar{ - Bar: progress.AddBar(total), +func NewBar(u string, total int, stat *Statistor, p *mpb.Progress) *Bar { + m := metrics.NewMeter() + metrics.Register(u, m) + + // 在mpb v8中，Name装饰器的使用方式略有不同 + 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"), + ), + mpb.AppendDecorators( + // 显示经过的时间 + decor.Elapsed(decor.ET_STYLE_GO, decor.WC{W: 4}), + ), + ) + + return &Bar{ url: u, - m: metrics.NewMeter(), + bar: bar, + m: m, } - - metrics.Register(bar.url, bar.m) - bar.PrependCompleted() - bar.PrependFunc(func(b *uiprogress.Bar) string { - return fmt.Sprintf("%f/s %d/%d", bar.m.Rate1(), bar.m.Count(), bar.Bar.Total) - }) - bar.PrependFunc(func(b *uiprogress.Bar) string { - return u - }) - bar.AppendElapsed() - - return bar } type Bar struct { - url string - total int - close bool - *uiprogress.Bar - m metrics.Meter + url string + bar *mpb.Bar + m metrics.Meter } func (bar *Bar) Done() { bar.m.Mark(1) - bar.Incr() + bar.bar.Increment() } func (bar *Bar) Close() { - metrics.Unregister(bar.url) - bar.close = true + //metrics.Unregister(bar.url) + // 标记进度条为完成状态 + //bar.bar.Abort(false) }