refactor(report): move colorize function from trivy-db (#2122)

pkg/k8s/summary.go

@@ -9,8 +9,7 @@ import (
 
 	"github.com/aquasecurity/table"
 	dbTypes "github.com/aquasecurity/trivy-db/pkg/types"
-
-	"github.com/liamg/tml"
+	pkgReport "github.com/aquasecurity/trivy/pkg/report"
 )
 
 type SummaryWriter struct {
@@ -62,7 +61,7 @@ func (s SummaryWriter) Write(report Report) error {
 
 	keyParts := []string{"Severities:"}
 	for _, s := range s.Severities {
-		keyParts = append(keyParts, fmt.Sprintf("%s=%s", s[:1], colourSeverityValue(s, s)))
+		keyParts = append(keyParts, fmt.Sprintf("%s=%s", s[:1], pkgReport.ColorizeSeverity(s, s)))
 	}
 
 	_, _ = fmt.Fprintln(s.Output, strings.Join(keyParts, " "))
@@ -75,7 +74,7 @@ func (s SummaryWriter) generateSummary(sevCount map[string]int) []string {
 
 	for _, sev := range s.Severities {
 		if count, ok := sevCount[sev]; ok {
-			parts = append(parts, colourSeverityValue(strconv.Itoa(count), sev))
+			parts = append(parts, pkgReport.ColorizeSeverity(strconv.Itoa(count), sev))
 		} else {
 			parts = append(parts, " ")
 		}
@@ -142,18 +141,3 @@ func configureHeader(s SummaryWriter, t *table.Table) {
 	t.SetAutoMergeHeaders(true)
 	t.SetHeaderColSpans(0, 1, 1, sevCount, sevCount, sevCount)
 }
-
-func colourSeverityValue(value string, severity string) string {
-	switch severity {
-	case "CRITICAL":
-		return tml.Sprintf("<bold><red>%s</red></bold>", value)
-	case "HIGH":
-		return tml.Sprintf("<red>%s</red>", value)
-	case "MEDIUM":
-		return tml.Sprintf("<yellow>%s</yellow>", value)
-	case "UNKNOWN":
-		return tml.Sprintf("<blue>%s</blue>", value)
-	default:
-		return tml.Sprintf("%s", value)
-	}
-}

pkg/report/table.go

@@ -8,18 +8,25 @@ import (
 	"strings"
 	"sync"
 
-	ftypes "github.com/aquasecurity/fanal/types"
-
-	"github.com/aquasecurity/table"
-
+	"github.com/fatih/color"
 	"github.com/liamg/tml"
+	"golang.org/x/exp/slices"
 
+	ftypes "github.com/aquasecurity/fanal/types"
+	"github.com/aquasecurity/table"
 	dbTypes "github.com/aquasecurity/trivy-db/pkg/types"
 	"github.com/aquasecurity/trivy/pkg/log"
 	"github.com/aquasecurity/trivy/pkg/types"
+)
 
-	"github.com/fatih/color"
-	"golang.org/x/exp/slices"
+var (
+	SeverityColor = []func(a ...interface{}) string{
+		color.New(color.FgCyan).SprintFunc(),   // UNKNOWN
+		color.New(color.FgBlue).SprintFunc(),   // LOW
+		color.New(color.FgYellow).SprintFunc(), // MEDIUM
+		color.New(color.FgHiRed).SprintFunc(),  // HIGH
+		color.New(color.FgRed).SprintFunc(),    // CRITICAL
+	}
 )
 
 // TableWriter implements Writer and output in tabular form
@@ -175,7 +182,7 @@ func (tw TableWriter) setVulnerabilityRows(tableWriter *table.Table, vulns []typ
 
 		var row []string
 		if tw.isOutputToTerminal() {
-			row = []string{lib, v.VulnerabilityID, dbTypes.ColorizeSeverity(v.Severity),
+			row = []string{lib, v.VulnerabilityID, ColorizeSeverity(v.Severity, v.Severity),
 				v.InstalledVersion, v.FixedVersion, strings.TrimSpace(title)}
 		} else {
 			row = []string{lib, v.VulnerabilityID, v.Severity, v.InstalledVersion, v.FixedVersion, strings.TrimSpace(title)}
@@ -227,7 +234,7 @@ func (tw TableWriter) setSecretRows(tableWriter *table.Table, secrets []ftypes.S
 	for _, secret := range secrets {
 		severity := secret.Severity
 		if tw.isOutputToTerminal() {
-			severity = dbTypes.ColorizeSeverity(severity)
+			severity = ColorizeSeverity(severity, severity)
 		}
 		row := []string{string(secret.Category), secret.Title, severity,
 			fmt.Sprint(secret.StartLine), // multi-line is not supported for now.
@@ -257,3 +264,12 @@ func (tw TableWriter) countSeverities(result types.Result) map[string]int {
 	}
 	return severityCount
 }
+
+func ColorizeSeverity(value, severity string) string {
+	for i, name := range dbTypes.SeverityNames {
+		if severity == name {
+			return SeverityColor[i](value)
+		}
+	}
+	return color.New(color.FgBlue).SprintFunc()(severity)
+}