From eb9fe85c753b2816e5b6cd97905bfd18ecb9a818 Mon Sep 17 00:00:00 2001
From: Daniel-GrunbergerCA <danielg@armosec.io>
Date: Wed, 20 Oct 2021 17:45:11 +0300
Subject: [PATCH] add error handling

---
 cautils/getter/getpoliciesutils.go |  1 +
 cautils/getter/loadpolicy.go       |  1 +
 clihandler/cmd/download.go         |  5 ++++-
 clihandler/cmd/scan.go             | 12 ++++++++++++
 4 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/cautils/getter/getpoliciesutils.go b/cautils/getter/getpoliciesutils.go
index 55a4bc2d..7949eb70 100644
--- a/cautils/getter/getpoliciesutils.go
+++ b/cautils/getter/getpoliciesutils.go
@@ -21,6 +21,7 @@ func GetDefaultPath(name string) string {
 	return defaultfilePath
 }
 
+// Save control as json in file
 func SaveControlInFile(control *reporthandling.Control, pathStr string) error {
 	encodedData, err := json.Marshal(control)
 	if err != nil {
diff --git a/cautils/getter/loadpolicy.go b/cautils/getter/loadpolicy.go
index f0553f45..c976a478 100644
--- a/cautils/getter/loadpolicy.go
+++ b/cautils/getter/loadpolicy.go
@@ -26,6 +26,7 @@ func NewLoadPolicy(filePath string) *LoadPolicy {
 	}
 }
 
+// Return control from file
 func (lp *LoadPolicy) GetControl(controlName string) (*reporthandling.Control, error) {
 
 	control := &reporthandling.Control{}
diff --git a/clihandler/cmd/download.go b/clihandler/cmd/download.go
index c0c3b6c9..6ed2dadf 100644
--- a/clihandler/cmd/download.go
+++ b/clihandler/cmd/download.go
@@ -18,7 +18,10 @@ var downloadCmd = &cobra.Command{
 	Long:  ``,
 	Args: func(cmd *cobra.Command, args []string) error {
 		if len(args) != 2 {
-			return fmt.Errorf("requires two arguments : framework <framework-name>")
+			return fmt.Errorf("requires two arguments : framework/control <framework-name>/<control-name>")
+		}
+		if !strings.EqualFold(args[0], "framework") && !strings.EqualFold(args[0], "control") {
+			return fmt.Errorf("invalid parameter '%s'. Supported parameters: framework, control", args[0])
 		}
 		return nil
 	},
diff --git a/clihandler/cmd/scan.go b/clihandler/cmd/scan.go
index 13d2da22..e86e5a02 100644
--- a/clihandler/cmd/scan.go
+++ b/clihandler/cmd/scan.go
@@ -1,6 +1,9 @@
 package cmd
 
 import (
+	"fmt"
+	"strings"
+
 	"github.com/armosec/kubescape/cautils"
 	"github.com/spf13/cobra"
 )
@@ -12,6 +15,15 @@ var scanCmd = &cobra.Command{
 	Use:   "scan <command>",
 	Short: "Scan the current running cluster or yaml files",
 	Long:  `The action you want to perform`,
+	Args: func(cmd *cobra.Command, args []string) error {
+		if len(args) == 0 {
+			return fmt.Errorf("requires one  argument: framework/control")
+		}
+		if !strings.EqualFold(args[0], "framework") && !strings.EqualFold(args[0], "control") {
+			return fmt.Errorf("invalid parameter '%s'. Supported parameters: framework, control", args[0])
+		}
+		return nil
+	},
 	Run: func(cmd *cobra.Command, args []string) {
 	},
 }