实现文件写入

This commit is contained in:
M09Ic 2022-11-29 15:08:10 +08:00
parent 922ed78611
commit 779c98487c
2 changed files with 52 additions and 10 deletions

View File

@ -3,6 +3,7 @@ package internal
import ( import (
"fmt" "fmt"
"github.com/antonmedv/expr" "github.com/antonmedv/expr"
"github.com/chainreactors/files"
"github.com/chainreactors/logs" "github.com/chainreactors/logs"
"github.com/chainreactors/spray/pkg" "github.com/chainreactors/spray/pkg"
"github.com/chainreactors/words/mask" "github.com/chainreactors/words/mask"
@ -295,6 +296,21 @@ func (opt *Option) PrepareRunner() (*Runner, error) {
if opt.OutputProbe != "" { if opt.OutputProbe != "" {
r.Probes = strings.Split(opt.OutputProbe, ",") r.Probes = strings.Split(opt.OutputProbe, ",")
} }
if opt.OutputFile != "" {
r.OutputFile, err = files.NewFile(opt.OutputFile, false, false, true)
if err != nil {
return nil, err
}
r.FuzzyFile, err = files.NewFile(opt.OutputFile+"fuzzy", false, false, true)
if err != nil {
return nil, err
}
r.StatFile, err = files.NewFile(opt.OutputFile+"fuzzy", false, false, true)
if err != nil {
return nil, err
}
}
return r, nil return r, nil
} }

View File

@ -46,6 +46,7 @@ type Runner struct {
Fuzzy bool Fuzzy bool
OutputFile *files.File OutputFile *files.File
FuzzyFile *files.File FuzzyFile *files.File
StatFile *files.File
Force bool Force bool
Progress *uiprogress.Progress Progress *uiprogress.Progress
Offset int Offset int
@ -142,6 +143,10 @@ func (r *Runner) Prepare(ctx context.Context) error {
pool.Run(ctx, r.Offset, r.Limit) pool.Run(ctx, r.Offset, r.Limit)
logs.Log.Important(pool.Statistor.String()) logs.Log.Important(pool.Statistor.String())
logs.Log.Important(pool.Statistor.Detail()) logs.Log.Important(pool.Statistor.Detail())
if r.StatFile != nil {
r.StatFile.SafeWrite(pool.Statistor.Json() + "\n")
r.StatFile.SafeSync()
}
r.Done() r.Done()
}) })
@ -230,16 +235,25 @@ func (r *Runner) Done() {
func (r *Runner) Outputting() { func (r *Runner) Outputting() {
go func() { go func() {
var outFunc func(*pkg.Baseline) var saveFunc func(*pkg.Baseline)
if r.OutputFile != nil {
saveFunc = func(bl *pkg.Baseline) {
r.OutputFile.SafeWrite(bl.Jsonify() + "\n")
r.OutputFile.SafeSync()
}
} else {
if len(r.Probes) > 0 { if len(r.Probes) > 0 {
outFunc = func(bl *pkg.Baseline) { saveFunc = func(bl *pkg.Baseline) {
logs.Log.Console("[+] " + bl.Format(r.Probes) + "\n") logs.Log.Console("[+] " + bl.Format(r.Probes) + "\n")
} }
} else { } else {
outFunc = func(bl *pkg.Baseline) { saveFunc = func(bl *pkg.Baseline) {
logs.Log.Console("[+] " + bl.String() + "\n") logs.Log.Console("[+] " + bl.String() + "\n")
} }
} }
}
for { for {
select { select {
@ -249,7 +263,7 @@ func (r *Runner) Outputting() {
} }
if bl.IsValid { if bl.IsValid {
outFunc(bl) saveFunc(bl)
} else { } else {
logs.Log.Debug(bl.String()) logs.Log.Debug(bl.String())
} }
@ -258,15 +272,27 @@ func (r *Runner) Outputting() {
}() }()
go func() { go func() {
var fuzzySaveFunc func(*pkg.Baseline)
if r.FuzzyFile != nil {
fuzzySaveFunc = func(bl *pkg.Baseline) {
r.FuzzyFile.SafeWrite(bl.Jsonify() + "\n")
r.FuzzyFile.SafeSync()
}
} else {
fuzzySaveFunc = func(bl *pkg.Baseline) {
if r.Fuzzy {
logs.Log.Console("[baseline.fuzzy] " + bl.String() + "\n")
}
}
}
for { for {
select { select {
case bl, ok := <-r.FuzzyCh: case bl, ok := <-r.FuzzyCh:
if !ok { if !ok {
return return
} }
if r.Fuzzy { fuzzySaveFunc(bl)
logs.Log.Console("[baseline.fuzzy] " + bl.String() + "\n")
}
} }
} }
}() }()