admin/spray
1
0
Fork 0
mirror of https://github.com/chainreactors/spray.git synced 2025-06-21 18:30:49 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix config load priority, now config < cmd

Browse Source 
enhance bar print
fix config.yaml default
This commit is contained in:
M09Ic 2024-03-07 04:15:28 +08:00
parent b1e42e763d
commit 4e78e55b6e
5 changed files with 80 additions and 53 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
cmd/cmd.go
View File

@ -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,16 +84,12 @@ 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
}
}
}
if option.Version {
fmt.Println(ver)

14
config.yaml
View File

@ -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: ""

20
internal/option.go
View File

@ -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)

40
internal/runner.go
View File

@ -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() {
if r.bar != nil {
r.bar.Increment()
}
r.finished++
r.poolwg.Done()
}

35
pkg/bar.go
View File

@ -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)
}