fix: only show severity total from filter flags (#559)

* Add filter table output by severity flags * Simplify filter output * Fix unit test * add filter unique severities * remove wrong comment * Remove utils and update logic * chore(mod): tidy Co-authored-by: knqyf263 <knqyf263@gmail.com>
aquasecurity · Jul 29, 2020 · 469c0b4 · 469c0b4
1 parent 4a34f72
commit 469c0b4
Show file tree

Hide file tree

Showing 5 changed files with 22 additions and 13 deletions.
diff --git a/internal/artifact/run.go b/internal/artifact/run.go
@@ -95,7 +95,7 @@ func run(c config.Config, initializeScanner InitializeScanner) error {
 		results[i].Vulnerabilities = vulns
 	}
 
-	if err = report.WriteResults(c.Format, c.Output, results, c.Template, c.Light); err != nil {
+	if err = report.WriteResults(c.Format, c.Output, c.Severities, results, c.Template, c.Light); err != nil {
 		return xerrors.Errorf("unable to write results: %w", err)
 	}
 

diff --git a/internal/client/run.go b/internal/client/run.go
@@ -88,7 +88,7 @@ func run(c config.Config) (err error) {
 		results[i].Vulnerabilities = vulns
 	}
 
-	if err = report.WriteResults(c.Format, c.Output, results, c.Template, false); err != nil {
+	if err = report.WriteResults(c.Format, c.Output, c.Severities, results, c.Template, false); err != nil {
 		return xerrors.Errorf("unable to write results: %w", err)
 	}
 

diff --git a/pkg/report/writer.go b/pkg/report/writer.go
@@ -15,7 +15,7 @@ import (
 
 	dbTypes "github.com/aquasecurity/trivy-db/pkg/types"
 	"github.com/aquasecurity/trivy/pkg/types"
-
+	"github.com/aquasecurity/trivy/pkg/utils"
 	"github.com/olekukonko/tablewriter"
 )
 
@@ -27,11 +27,11 @@ type Result struct {
 	Vulnerabilities []types.DetectedVulnerability `json:"Vulnerabilities"`
 }
 
-func WriteResults(format string, output io.Writer, results Results, outputTemplate string, light bool) error {
+func WriteResults(format string, output io.Writer, severities []dbTypes.Severity, results Results, outputTemplate string, light bool) error {
 	var writer Writer
 	switch format {
 	case "table":
-		writer = &TableWriter{Output: output, Light: light}
+		writer = &TableWriter{Output: output, Light: light, Severities: severities}
 	case "json":
 		writer = &JsonWriter{Output: output}
 	case "template":
@@ -80,8 +80,9 @@ type Writer interface {
 }
 
 type TableWriter struct {
-	Output io.Writer
-	Light  bool
+	Severities []dbTypes.Severity
+	Output     io.Writer
+	Light      bool
 }
 
 func (tw TableWriter) Write(results Results) error {
@@ -125,7 +126,16 @@ func (tw TableWriter) write(result Result) {
 	}
 
 	var results []string
+
+	var severities []string
+	for _, sev := range tw.Severities {
+		severities = append(severities, sev.String())
+	}
+
 	for _, severity := range dbTypes.SeverityNames {
+		if !utils.StringInSlice(severity, severities) {
+			continue
+		}
 		r := fmt.Sprintf("%s: %d", severity, severityCount[severity])
 		results = append(results, r)
 	}

diff --git a/pkg/report/writer_test.go b/pkg/report/writer_test.go
@@ -122,7 +122,7 @@ func TestReportWriter_Table(t *testing.T) {
 				},
 			}
 			tableWritten := bytes.Buffer{}
-			assert.NoError(t, report.WriteResults("table", &tableWritten, inputResults, "", tc.light), tc.name)
+			assert.NoError(t, report.WriteResults("table", &tableWritten, nil, inputResults, "", tc.light), tc.name)
 			assert.Equal(t, tc.expectedOutput, tableWritten.String(), tc.name)
 		})
 	}
@@ -184,7 +184,7 @@ func TestReportWriter_JSON(t *testing.T) {
 				},
 			}
 
-			assert.NoError(t, report.WriteResults("json", &jsonWritten, inputResults, "", false), tc.name)
+			assert.NoError(t, report.WriteResults("json", &jsonWritten, nil, inputResults, "", false), tc.name)
 
 			writtenResults := report.Results{}
 			errJson := json.Unmarshal([]byte(jsonWritten.String()), &writtenResults)
@@ -308,7 +308,7 @@ func TestReportWriter_Template(t *testing.T) {
 				},
 			}
 
-			assert.NoError(t, report.WriteResults("template", &tmplWritten, inputResults, tc.template, false))
+			assert.NoError(t, report.WriteResults("template", &tmplWritten, nil, inputResults, tc.template, false))
 			assert.Equal(t, tc.expected, tmplWritten.String())
 		})
 	}

diff --git a/pkg/utils/utils_test.go b/pkg/utils/utils_test.go
@@ -9,11 +9,10 @@ import (
 	"strings"
 	"testing"
 
-	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
-
 	"github.com/aquasecurity/trivy/pkg/log"
 	"github.com/kylelemons/godebug/pretty"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func touch(t *testing.T, name string) {