From 92f0b6e551dfd69486ba79bc45108ff4f032e7be Mon Sep 17 00:00:00 2001
From: Barnabas Makonda <6409210+makoscafee@users.noreply.github.com>
Date: Thu, 12 May 2022 16:45:46 +0300
Subject: [PATCH] fix issue when the files-path is actual file path instead of
 directory (#761)

* fix issue when the files-path is actual path

* remove check for file extension

* do not panic when sub command is not there

* remove panic and fail graciously
---
 cmd/polaris/fix.go | 48 ++++++++++++++++++++++++++++++++--------------
 1 file changed, 34 insertions(+), 14 deletions(-)

diff --git a/cmd/polaris/fix.go b/cmd/polaris/fix.go
index a5f7fe09..d2193378 100644
--- a/cmd/polaris/fix.go
+++ b/cmd/polaris/fix.go
@@ -22,6 +22,7 @@ import (
 	"io/ioutil"
 	"os"
 	"path/filepath"
+	"strings"
 
 	"github.com/fairwindsops/polaris/pkg/kube"
 	"github.com/fairwindsops/polaris/pkg/mutation"
@@ -47,15 +48,28 @@ var fixCommand = &cobra.Command{
 		logrus.Debug("Setting up controller manager")
 
 		if filesPath == "" {
-			logrus.Error("Please specify a sub-command.")
-			err := cmd.Help()
-			panic(err)
+			logrus.Error("Please specify a file-path flag")
+			cmd.Help()
+			os.Exit(1)
 		}
-
-		baseDir := filesPath + "/"
-		yamlFiles, err := getYamlFiles(baseDir)
+		var yamlFiles []string
+		fileInfo, err := os.Stat(filesPath)
 		if err != nil {
-			panic(err)
+			logrus.Error(err)
+			os.Exit(1)
+		}
+		if fileInfo.IsDir() {
+			baseDir := filesPath
+			if !strings.HasSuffix(filesPath, "/") {
+				baseDir = baseDir + "/"
+			}
+			yamlFiles, err = getYamlFiles(baseDir)
+			if err != nil {
+				logrus.Error(err)
+				os.Exit(1)
+			}
+		} else {
+			yamlFiles = append(yamlFiles, filesPath)
 		}
 		var contentStr string
 		isFirstResource := true
@@ -63,7 +77,8 @@ var fixCommand = &cobra.Command{
 
 			yamlFile, err := ioutil.ReadFile(fullFilePath)
 			if err != nil {
-				panic(err)
+				logrus.Errorf("Error reading file with file path %s: %v", fullFilePath, err)
+				os.Exit(1)
 			}
 
 			dec := yamlV3.NewDecoder(bytes.NewReader(yamlFile))
@@ -80,16 +95,19 @@ var fixCommand = &cobra.Command{
 					break
 				}
 				if err != nil {
-					panic(err)
+					logrus.Errorf("Error decoding data for file with file path %s: %v", fullFilePath, err)
+					os.Exit(1)
 				}
 				yamlContent, err := yamlV3.Marshal(data)
 				if err != nil {
-					panic(err)
+					logrus.Errorf("Error marshalling %s: %v", fullFilePath, err)
+					os.Exit(1)
 				}
 				kubeResources := kube.CreateResourceProviderFromYaml(string(yamlContent))
 				results, err := validator.ApplyAllSchemaChecksToResourceProvider(&config, kubeResources)
 				if err != nil {
-					panic(err)
+					logrus.Errorf("Error applying schema check to the resources %s: %v", fullFilePath, err)
+					os.Exit(1)
 				}
 				comments, allMutations := mutation.GetMutationsAndCommentsFromResults(results)
 				updatedYamlContent := string(yamlContent)
@@ -99,11 +117,13 @@ var fixCommand = &cobra.Command{
 						mutations := allMutations[key]
 						mutated, err := mutation.ApplyAllSchemaMutations(&config, kubeResources, resources[0], mutations)
 						if err != nil {
-							panic(err)
+							logrus.Errorf("Error applying schema mutations to the resources: %v", err)
+							os.Exit(1)
 						}
 						mutatedYamlContent, err := yaml.JSONToYAML(mutated.OriginalObjectJSON)
 						if err != nil {
-							panic(err)
+							logrus.Errorf("Error converting JSON to Yaml : %v", err)
+							os.Exit(1)
 						}
 						updatedYamlContent = mutation.UpdateMutatedContentWithComments(string(mutatedYamlContent), comments)
 					}
@@ -137,7 +157,7 @@ func getYamlFiles(rootpath string) ([]string, error) {
 		if info.IsDir() {
 			return nil
 		}
-		if filepath.Ext(path) == ".yaml" {
+		if filepath.Ext(path) == ".yaml" || filepath.Ext(path) == ".yml" {
 			list = append(list, path)
 		}
 		return nil