Merge pull request #76 from lacework/afiune/cleanup

style: avoid mixing duties between api/ and cli/
lacework · Apr 27, 2020 · 88e92a8 · 88e92a8
2 parents fc8a90f + fb7b7c2
commit 88e92a8
Show file tree

Hide file tree

Showing 2 changed files with 53 additions and 46 deletions.
diff --git a/api/vulnerabilities.go b/api/vulnerabilities.go
@@ -21,6 +21,8 @@ package api
 import (
 	"fmt"
 	"strings"
+
+	"github.com/lacework/go-sdk/internal/array"
 )
 
 // VulnerabilitiesService is a service that interacts with the vulnerabilities
@@ -29,6 +31,9 @@ type VulnerabilitiesService struct {
 	client *Client
 }
 
+// ValidVulSeverities is a list of all valid severities in a vulnerability report
+var ValidVulSeverities = []string{"critical", "high", "medium", "low", "info"}
+
 // Scan triggers a vulnerability scan to the provider registry, repository, and
 // tag provided. This function calls the underlaying API endpoint that assumes
 // that the container repository has been already integrated with the platform.
@@ -147,7 +152,7 @@ type VulContainerReport struct {
 	InfoVulnerabilities     int32             `json:"info_vulnerabilities"`
 	FixableVulnerabilities  int32             `json:"fixable_vulnerabilities"`
 	LastEvaluationTime      string            `json:"last_evaluation_time"`
-	Image                   vulContainerImage `json:"image"`
+	Image                   VulContainerImage `json:"image"`
 
 	// @afiune these two parameters, Status and Message will appear when
 	// the vulnerability scan is still running. ugh. why?
@@ -162,24 +167,11 @@ type VulContainerReport struct {
 	//LastEvaluationTime      time.Time         `json:"last_evaluation_time"`
 }
 
-func (report *VulContainerReport) VulCountsTable() [][]string {
-	return [][]string{
-		[]string{"Critical", fmt.Sprint(report.CriticalVulnerabilities),
-			fmt.Sprint(report.VulFixableCount("critical"))},
-		[]string{"High", fmt.Sprint(report.HighVulnerabilities),
-			fmt.Sprint(report.VulFixableCount("high"))},
-		[]string{"Medium", fmt.Sprint(report.MediumVulnerabilities),
-			fmt.Sprint(report.VulFixableCount("medium"))},
-		[]string{"Low", fmt.Sprint(report.LowVulnerabilities),
-			fmt.Sprint(report.VulFixableCount("low"))},
-		[]string{"Info", fmt.Sprint(report.InfoVulnerabilities),
-			fmt.Sprint(report.VulFixableCount("info"))},
-	}
-}
-
 func (report *VulContainerReport) VulFixableCount(severity string) int32 {
-	// @afiune check valid severity
 	severity = strings.ToLower(severity)
+	if !array.ContainsStr(ValidVulSeverities, severity) {
+		return 0
+	}
 
 	if len(report.Image.ImageLayers) == 0 {
 		return 0
@@ -198,37 +190,11 @@ func (report *VulContainerReport) VulFixableCount(severity string) int32 {
 	return fixable
 }
 
-type vulContainerImage struct {
+type VulContainerImage struct {
 	ImageInfo   vulContainerImageInfo    `json:"image_info"`
 	ImageLayers []vulContainerImageLayer `json:"image_layers"`
 }
 
-func (image *vulContainerImage) Table() [][]string {
-	info := image.ImageInfo
-	return [][]string{
-		[]string{"ID", info.ImageID},
-		[]string{"Digest", info.ImageDigest},
-		[]string{"Registry", info.Registry},
-		[]string{"Repository", info.Repository},
-		[]string{"Size", ByteCountBinary(info.Size)},
-		[]string{"Created At", info.CreatedTime},
-		[]string{"Tags", strings.Join(info.Tags, ",")},
-	}
-}
-
-func ByteCountBinary(b int64) string {
-	const unit = 1024
-	if b < unit {
-		return fmt.Sprintf("%d B", b)
-	}
-	div, exp := int64(unit), 0
-	for n := b / unit; n >= unit; n /= unit {
-		div *= unit
-		exp++
-	}
-	return fmt.Sprintf("%.1f %cB", float64(b)/float64(div), "KMGTPE"[exp])
-}
-
 type vulContainerImageInfo struct {
 	ImageDigest string   `json:"image_digest"`
 	ImageID     string   `json:"image_id"`

diff --git a/cli/cmd/vulnerability.go b/cli/cmd/vulnerability.go
@@ -340,7 +340,7 @@ func buildVulnerabilityReport(report *api.VulContainerReport) string {
 	t = tablewriter.NewWriter(imageDetailsTable)
 	t.SetBorder(false)
 	t.SetColumnSeparator("")
-	t.AppendBulk(report.Image.Table())
+	t.AppendBulk(vulContainerImageToTable(&report.Image))
 	t.Render()
 
 	t = tablewriter.NewWriter(vulCountsTable)
@@ -349,7 +349,7 @@ func buildVulnerabilityReport(report *api.VulContainerReport) string {
 	t.SetHeader([]string{
 		"Severity", "Count", "Fixable",
 	})
-	t.AppendBulk(report.VulCountsTable())
+	t.AppendBulk(vulContainerReportToCountsTable(report))
 	t.Render()
 
 	t = tablewriter.NewWriter(mainReport)
@@ -367,3 +367,44 @@ func buildVulnerabilityReport(report *api.VulContainerReport) string {
 
 	return mainReport.String()
 }
+
+func vulContainerReportToCountsTable(report *api.VulContainerReport) [][]string {
+	return [][]string{
+		[]string{"Critical", fmt.Sprint(report.CriticalVulnerabilities),
+			fmt.Sprint(report.VulFixableCount("critical"))},
+		[]string{"High", fmt.Sprint(report.HighVulnerabilities),
+			fmt.Sprint(report.VulFixableCount("high"))},
+		[]string{"Medium", fmt.Sprint(report.MediumVulnerabilities),
+			fmt.Sprint(report.VulFixableCount("medium"))},
+		[]string{"Low", fmt.Sprint(report.LowVulnerabilities),
+			fmt.Sprint(report.VulFixableCount("low"))},
+		[]string{"Info", fmt.Sprint(report.InfoVulnerabilities),
+			fmt.Sprint(report.VulFixableCount("info"))},
+	}
+}
+
+func vulContainerImageToTable(image *api.VulContainerImage) [][]string {
+	info := image.ImageInfo
+	return [][]string{
+		[]string{"ID", info.ImageID},
+		[]string{"Digest", info.ImageDigest},
+		[]string{"Registry", info.Registry},
+		[]string{"Repository", info.Repository},
+		[]string{"Size", byteCountBinary(info.Size)},
+		[]string{"Created At", info.CreatedTime},
+		[]string{"Tags", strings.Join(info.Tags, ",")},
+	}
+}
+
+func byteCountBinary(b int64) string {
+	const unit = 1024
+	if b < unit {
+		return fmt.Sprintf("%d B", b)
+	}
+	div, exp := int64(unit), 0
+	for n := b / unit; n >= unit; n /= unit {
+		div *= unit
+		exp++
+	}
+	return fmt.Sprintf("%.1f %cB", float64(b)/float64(div), "KMGTPE"[exp])
+}