From f41d68d0fd26bbf7b5fcd61f18ebf04c3abce7c1 Mon Sep 17 00:00:00 2001
From: Patazerty <github@pata.ovh>
Date: Wed, 15 Jan 2020 15:44:35 +0100
Subject: [PATCH] Feat: Added --version, --values and --set helm flags

Version bump to 0.1.2
---
 README.md   | 14 ++++++++++----
 main.go     | 40 ++++++++++++++++++++++++++--------------
 plugin.yaml |  2 +-
 3 files changed, 37 insertions(+), 19 deletions(-)

diff --git a/README.md b/README.md
index 0b898f5..c219987 100644
--- a/README.md
+++ b/README.md
@@ -20,14 +20,20 @@ Usage: helm trivy [options] <helm chart>
 Example: helm trivy -json stable/mariadb
 
 Options:
-  -debug
+  --debug
     	Enable debug logging
-  -json
+  --json
     	Enable JSON output
-  -nopull
+  --nopull
     	Don't pull latest trivy image
-  -trivyargs string
+  --set string
+    	Values to set for helm chart, format: 'key1=value1,key2=value2'
+  --trivyargs string
     	CLI args to passthrough to trivy
+  --values string
+    	Specify chart values in a YAML file or a URL
+  --version string
+    	Specify chart version
 ```
 
 Some examples:
diff --git a/main.go b/main.go
index 7b3bf30..b6f5217 100644
--- a/main.go
+++ b/main.go
@@ -21,9 +21,21 @@ import (
 
 var debug = false
 
-func getChartImages(chart string) (error, []string) {
+func getChartImages(chart string, set string, values string, version string) (error, []string) {
 	images := []string{}
-	out, err := exec.Command("helm", "template", chart).Output()
+	cmd := []string{"template"}
+	if len(set) > 0 {
+		cmd = append(cmd, "--set", set)
+	}
+	if len(values) > 0 {
+		cmd = append(cmd, "--values", values)
+	}
+	if len(version) > 0 {
+		cmd = append(cmd, "--version", version)
+	}
+	cmd = append(cmd, chart)
+	log.Debugf("Running helm cmd: helm %v", cmd)
+	out, err := exec.Command("helm", cmd...).Output()
 	if err != nil {
 		return err, images
 	}
@@ -91,10 +103,10 @@ func scanImage(image string, ctx context.Context, cli *client.Client, cacheDir s
 	return string(outputContent)
 }
 
-func scanChart(chart string, json bool, ctx context.Context, cli *client.Client, cacheDir string, trivyOpts string) {
+func scanChart(chart string, json bool, ctx context.Context, cli *client.Client, cacheDir string, trivyOpts string, templateSet string, templateValues string, chartversion string) {
 	log.Infof("Scanning chart %s", chart)
 	jsonOutput := ""
-	if err, images := getChartImages(chart); err != nil {
+	if err, images := getChartImages(chart, templateSet, templateValues, chartversion); err != nil {
 		log.Fatalf("Could not find images for chart %v: %v. Did you run 'helm repo update' ?", chart, err)
 	} else {
 		if len(images) == 0 {
@@ -120,6 +132,9 @@ func main() {
 	var jsonOutput bool
 	var noPull bool
 	var chart string = ""
+	var templateSet = ""
+	var templateValues = ""
+	var chartVersion = ""
 	var trivyArgs = ""
 
 	flag.Usage = func() {
@@ -131,26 +146,23 @@ func main() {
 
 	flag.BoolVar(&jsonOutput, "json", false, "Enable JSON output")
 	flag.BoolVar(&debug, "debug", false, "Enable debug logging")
-
 	flag.BoolVar(&noPull, "nopull", false, "Don't pull latest trivy image")
 	flag.StringVar(&trivyArgs, "trivyargs", "", "CLI args to passthrough to trivy")
+	flag.StringVar(&templateSet, "set", "", "Values to set for helm chart, format: 'key1=value1,key2=value2'")
+	flag.StringVar(&templateValues, "values", "", "Specify chart values in a YAML file or a URL")
+	flag.StringVar(&chartVersion, "version", "", "Specify chart version")
 	flag.Parse()
 
 	if debug {
 		log.SetLevel(log.DebugLevel)
 	}
 
-	for _, v := range os.Args[1:] {
-		if strings.HasPrefix(v, "-") {
-			continue
-		}
-		chart = v
-		break
-	}
-	if chart == "" {
+	if len(flag.Args()) == 0 {
 		fmt.Fprintf(os.Stderr, "Error: No chart specified.\n")
 		flag.Usage()
 		os.Exit(2)
+	} else {
+		chart = flag.Args()[0]
 	}
 
 	ctx := context.Background()
@@ -183,5 +195,5 @@ func main() {
 		os.Exit(0)
 	}(cacheDir)
 
-	scanChart(chart, jsonOutput, ctx, cli, cacheDir, trivyArgs)
+	scanChart(chart, jsonOutput, ctx, cli, cacheDir, trivyArgs, templateSet, templateValues, chartVersion)
 }
diff --git a/plugin.yaml b/plugin.yaml
index b41901a..e982634 100644
--- a/plugin.yaml
+++ b/plugin.yaml
@@ -1,6 +1,6 @@
 ---
 name: "trivy"
-version: "0.1.1"
+version: "0.1.2"
 usage: "Check images in your charts for vulnerabilities"
 description: |-
   "Test your helm charts' docker images with trivy"