实现文件写入

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 (
"fmt"
"github.com/antonmedv/expr"
"github.com/chainreactors/files"
"github.com/chainreactors/logs"
"github.com/chainreactors/spray/pkg"
"github.com/chainreactors/words/mask"
@ -295,6 +296,21 @@ func (opt *Option) PrepareRunner() (*Runner, error) {
if 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
}

View File

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