Skip to content

Commit

Permalink
clean output code
Browse files Browse the repository at this point in the history 
enhance output format and color
support jsonline output in stdout
  • Loading branch information
@M09Ic
M09Ic committed Jul 4, 2024
1 parent 2371021 commit eb328a4
Show file tree
Hide file tree
Showing 4 changed files with 108 additions and 150 deletions.
2 changes: 1 addition & 1 deletion go.mod
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
4 changes: 4 additions & 0 deletions go.sum
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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=
Expand Down
38 changes: 8 additions & 30 deletions internal/option.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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"`
Expand Down Expand Up @@ -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
Expand Down
214 changes: 95 additions & 119 deletions internal/runner.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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

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
*Option

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,
Expand Down Expand Up @@ -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:
Expand All @@ -429,7 +422,7 @@ func (r *Runner) OutputHandler() {
r.DumpFile.SafeSync()
}
if bl.IsValid {
saveFunc(bl)
outputPrint(bl)
if bl.Recu {
r.AddRecursive(bl)
}
Expand All @@ -442,31 +435,14 @@ 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:
if !ok {
return
}
if r.Fuzzy {
fuzzySaveFunc(bl)
outputPrint(bl)
}
r.outwg.Done()
}
Expand Down

0 comments on commit eb328a4

Please sign in to comment.