From 779c98487ce93870d0834590e9a031c3be306fef Mon Sep 17 00:00:00 2001
From: M09Ic <m09ician@gmail.com>
Date: Tue, 29 Nov 2022 15:08:10 +0800
Subject: [PATCH] =?UTF-8?q?=E5=AE=9E=E7=8E=B0=E6=96=87=E4=BB=B6=E5=86=99?=
 =?UTF-8?q?=E5=85=A5?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 internal/option.go | 16 ++++++++++++++++
 internal/runner.go | 46 ++++++++++++++++++++++++++++++++++++----------
 2 files changed, 52 insertions(+), 10 deletions(-)

diff --git a/internal/option.go b/internal/option.go
index 450e362..b7254fb 100644
--- a/internal/option.go
+++ b/internal/option.go
@@ -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
 }
 
diff --git a/internal/runner.go b/internal/runner.go
index 8b790f4..5976161 100644
--- a/internal/runner.go
+++ b/internal/runner.go
@@ -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)
 			}
 		}
 	}()