enhance probe output

chainreactors · Oct 13, 2024 · af82ae4 · af82ae4
1 parent e483bb4
commit af82ae4
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 7 deletions.
diff --git a/internal/format.go b/internal/format.go
@@ -39,17 +39,20 @@ func Format(opts Option) {
 		group[result.Url.Host] = append(group[result.Url.Host], &result)
 	}
 
-	// 分组
-
 	for _, results := range group {
 		for _, result := range results {
 			if !opts.Fuzzy && result.IsFuzzy {
 				continue
 			}
-			if !opts.NoColor {
-				logs.Log.Console(result.ColorString() + "\n")
+			if opts.OutputProbe == "" {
+				if !opts.NoColor {
+					logs.Log.Console(result.ColorString() + "\n")
+				} else {
+					logs.Log.Console(result.String() + "\n")
+				}
 			} else {
-				logs.Log.Console(result.String() + "\n")
+				probes := strings.Split(opts.OutputProbe, ",")
+				logs.Log.Console(result.ProbeOutput(probes) + "\n")
 			}
 		}
 	}

diff --git a/internal/runner.go b/internal/runner.go
@@ -364,7 +364,7 @@ func (r *Runner) Output(bl *pkg.Baseline) {
 	if r.Option.Json {
 		out = bl.ToJson()
 	} else if len(r.Probes) > 0 {
-		out = bl.Format(r.Probes)
+		out = bl.ProbeOutput(r.Probes)
 	} else if r.Color {
 		out = bl.ColorString()
 	} else {
@@ -385,7 +385,7 @@ func (r *Runner) Output(bl *pkg.Baseline) {
 		} else if r.FileOutput == "full" {
 			r.OutputFile.SafeWrite(bl.String() + "\n")
 		} else {
-			r.OutputFile.SafeWrite(bl.Format(strings.Split(r.FileOutput, ",")) + "\n")
+			r.OutputFile.SafeWrite(bl.ProbeOutput(strings.Split(r.FileOutput, ",")) + "\n")
 		}
 
 		r.OutputFile.SafeSync()

diff --git a/pkg/baseline.go b/pkg/baseline.go
@@ -235,6 +235,15 @@ func (bl *Baseline) Compare(other *Baseline) int {
 	return -1
 }
 
+func (bl *Baseline) ProbeOutput(format []string) string {
+	var s strings.Builder
+	for _, f := range format {
+		s.WriteString("\t")
+		s.WriteString(bl.Get(f))
+	}
+	return strings.TrimSpace(s.String())
+}
+
 var Distance uint8 = 5 // 数字越小越相似, 数字为0则为完全一致.
 
 func (bl *Baseline) FuzzyCompare(other *Baseline) bool {