From fa5e7fef23cb194ee415ecc563613bbf18b69f9b Mon Sep 17 00:00:00 2001 From: suhasgumma Date: Fri, 16 Dec 2022 03:25:13 +0530 Subject: [PATCH] Restructure Code to Improve Code Quality --- core/pkg/fixhandler/fixhandler.go | 117 +++++---- core/pkg/fixhandler/fixhandler_test.go | 247 ++++++++---------- .../fixed_yaml_scenario_1.yml | 2 +- .../fixed_yaml_scenario_1.yml | 2 +- .../fixed_yaml_scenario_2.yml | 2 +- .../fixed_yaml_scenario_3.yml | 2 +- .../fixed_yaml_scenario_5.yml | 2 +- .../fixed_yaml_scenario_6.yml | 2 +- .../fixed_yaml_scenario_1.yml | 2 +- .../fixed_yaml_scenario_2.yml | 1 - .../fixed_yaml_scenario_3.yml | 2 +- .../fixed_yaml_scenario_4.yml | 1 - .../fixed_yaml_scenario_1.yml | 2 +- .../fixed_yaml_scenario_2.yml | 2 +- core/pkg/fixhandler/yamlhandler.go | 113 ++++---- core/pkg/fixhandler/yamlhelper.go | 14 +- 16 files changed, 249 insertions(+), 264 deletions(-) diff --git a/core/pkg/fixhandler/fixhandler.go b/core/pkg/fixhandler/fixhandler.go index 3aabf011..7f61d4d5 100644 --- a/core/pkg/fixhandler/fixhandler.go +++ b/core/pkg/fixhandler/fixhandler.go @@ -186,39 +186,34 @@ func (h *FixHandler) PrintExpectedChanges(resourcesToFix []ResourceFixInfo) { func (h *FixHandler) ApplyChanges(resourcesToFix []ResourceFixInfo) (int, []error) { updatedFiles := make(map[string]bool) errors := make([]error, 0) - // Map with key as filepath - filePathFixInfo := make(map[string]*fileFixInfo) - for _, resourceToFix := range resourcesToFix { - singleExpression := reduceYamlExpressions(&resourceToFix) - resourceFilePath := resourceToFix.FilePath - if _, pathExistsInMap := filePathFixInfo[resourceFilePath]; !pathExistsInMap { - filePathFixInfo[resourceFilePath] = &fileFixInfo{ - contentToAdd: make([]contentToAdd, 0), - linesToRemove: make([]linesToRemove, 0), - } - } + fileYamlExpressions := h.getFileYamlExpressions(resourcesToFix) - contentsToAdd, linesToRemove, err := h.getResourceFileFix(resourceFilePath, singleExpression, resourceToFix.DocumentIndex) + for filepath, yamlExpression := range fileYamlExpressions { + fileAsString, err := getFileString(filepath) if err != nil { - errors = append(errors, - fmt.Errorf("failed to fix resource [Name: '%s', Kind: '%s'] in '%s': %w ", - resourceToFix.Resource.GetName(), - resourceToFix.Resource.GetKind(), - resourceToFix.FilePath, - err)) + logger.L().Error(err.Error()) + continue + } + + fixedYamlString, err := h.ApplyFix(fileAsString, yamlExpression) + + if err != nil { + errors = append(errors, fmt.Errorf("failed to fix file %s: %w ", filepath, err)) + continue } else { - h.addResourceFileFix(contentsToAdd, linesToRemove, filePathFixInfo[resourceFilePath]) - updatedFiles[resourceToFix.FilePath] = true + updatedFiles[filepath] = true + } + + err = writeFixesToFile(filepath, fixedYamlString) + + if err != nil { + logger.L().Error(fmt.Sprintf("Cannot Apply fixes to file %s, %v", filepath, err.Error())) + errors = append(errors, err) } } - err := h.applyFixToFiles(filePathFixInfo) - if err != nil { - logger.L().Fatal(fmt.Sprintf("Cannot Apply fixes to files, %v", err.Error())) - errors = append(errors, err) - } return len(updatedFiles), errors } @@ -236,42 +231,40 @@ func (h *FixHandler) getFilePathAndIndex(filePathWithIndex string) (filePath str } } -func (h *FixHandler) getResourceFileFix(filePath string, yamlExpression string, documentIdx int) (*[]contentToAdd, *[]linesToRemove, error) { - originalYamlNode := (*constructDecodedYaml(filePath))[documentIdx] - fixedYamlNodes, err := constructFixedYamlNodes(filePath, yamlExpression) +func (h *FixHandler) ApplyFix(yamlString, yamlExpression string) (fixedYamlString string, err error) { + yamlLines := strings.Split(yamlString, "\n") + + originalRootNodes := constructDecodedYaml(yamlString) + fixedRootNodes, err := constructFixedYamlNodes(yamlString, yamlExpression) + if err != nil { - return nil, nil, err + return "", err } - fixedYamlNode := (*fixedYamlNodes)[documentIdx] + contentsToAdd, linesToRemove := getFixInfo(originalRootNodes, fixedRootNodes) - originalList := constructDFSOrder(&originalYamlNode) - fixedList := constructDFSOrder(&fixedYamlNode) + fixedYamlLines := getFixedYamlLines(yamlLines, contentsToAdd, linesToRemove) - contentsToAdd, linesToRemove := getFixInfo(originalList, fixedList) - - return contentsToAdd, linesToRemove, nil + fixedYamlString = getFixedYamlString(fixedYamlLines) + return fixedYamlString, nil } -func (h *FixHandler) addResourceFileFix(contentToAdd *[]contentToAdd, linesToRemove *[]linesToRemove, fileFixInfo *fileFixInfo) { - for _, content := range *contentToAdd { - fileFixInfo.addContent(content) - } +func (h *FixHandler) getFileYamlExpressions(resourcesToFix []ResourceFixInfo) map[string]string { + fileYamlExpressions := make(map[string]string, 0) + for _, resourceToFix := range resourcesToFix { + singleExpression := reduceYamlExpressions(&resourceToFix) + resourceFilePath := resourceToFix.FilePath - for _, lines := range *linesToRemove { - fileFixInfo.addLinesToRemove(lines) - } -} - -func (h *FixHandler) applyFixToFiles(filePathFixInfo map[string]*fileFixInfo) error { - for filepath, fixInfo := range filePathFixInfo { - err := applyFixesToFile(filepath, &fixInfo.contentToAdd, &fixInfo.linesToRemove) - if err != nil { - return err + if _, pathExistsInMap := fileYamlExpressions[resourceFilePath]; !pathExistsInMap { + fileYamlExpressions[resourceFilePath] = singleExpression + } else { + fileYamlExpressions[resourceFilePath] = joinStrings(fileYamlExpressions[resourceFilePath], " | ", singleExpression) } + } - return nil + + return fileYamlExpressions } func (rfi *ResourceFixInfo) addYamlExpressionsFromResourceAssociatedControl(documentIndex int, ac *resourcesresults.ResourceAssociatedControl, skipUserValues bool) { @@ -304,6 +297,30 @@ func reduceYamlExpressions(resource *ResourceFixInfo) string { return strings.Join(expressions, " | ") } +func joinStrings(inputStrings ...string) string { + return strings.Join(inputStrings, "") +} + +func getFileString(filepath string) (string, error) { + bytes, err := ioutil.ReadFile(filepath) + + if err != nil { + return "", fmt.Errorf("Error reading file %s", filepath) + } + + return string(bytes), nil +} + +func writeFixesToFile(filepath, content string) error { + err := ioutil.WriteFile(filepath, []byte(content), 0644) + + if err != nil { + return fmt.Errorf("Error writing fixes to file: %w", err) + } + + return nil +} + func fixPathToValidYamlExpression(fixPath, value string, documentIndexInYaml int) string { isStringValue := true if _, err := strconv.ParseBool(value); err == nil { diff --git a/core/pkg/fixhandler/fixhandler_test.go b/core/pkg/fixhandler/fixhandler_test.go index b7ef9e03..8fadad4c 100644 --- a/core/pkg/fixhandler/fixhandler_test.go +++ b/core/pkg/fixhandler/fixhandler_test.go @@ -1,14 +1,10 @@ package fixhandler import ( - "fmt" - "io/ioutil" "os" "path/filepath" "testing" - "github.com/stretchr/testify/assert" - logger "github.com/kubescape/go-logger" metav1 "github.com/kubescape/kubescape/v2/core/meta/datastructures/v1" reporthandlingv2 "github.com/kubescape/opa-utils/reporthandling/v2" @@ -16,6 +12,12 @@ import ( "gopkg.in/op/go-logging.v1" ) +type indentationTestCase struct { + inputFile string + yamlExpression string + expectedFile string +} + func NewFixHandlerMock() (*FixHandler, error) { backendLoggerLeveled := logging.AddModuleLevel(logging.NewLogBackend(logger.L().GetWriter(), "", 0)) backendLoggerLeveled.SetLevel(logging.ERROR, "") @@ -33,151 +35,128 @@ func getTestdataPath() string { return filepath.Join(currentDir, "testdata") } -func testDirectoryApplyFixHelper(t *testing.T, yamlExpressions *[][]string, directoryPath string) { +func getTestCases() []indentationTestCase { + indentationTestCases := []indentationTestCase{ + // Insertion Scenarios + { + "insert_scenarios/original_yaml_scenario_1.yml", + "select(di==0).spec.containers[0].securityContext.allowPrivilegeEscalation |= false", + "insert_scenarios/fixed_yaml_scenario_1.yml", + }, + { + "insert_scenarios/original_yaml_scenario_2.yml", + "select(di==0).spec.containers[0].securityContext.capabilities.drop += [\"NET_RAW\"]", + "insert_scenarios/fixed_yaml_scenario_2.yml", + }, + { + "insert_scenarios/original_yaml_scenario_3.yml", + "select(di==0).spec.containers[0].securityContext.capabilities.drop += [\"SYS_ADM\"]", + "insert_scenarios/fixed_yaml_scenario_3.yml", + }, + { + "insert_scenarios/original_yaml_scenario_4.yml", - scenarioCount := len(*yamlExpressions) + `select(di==0).spec.template.spec.securityContext.allowPrivilegeEscalation |= false | + select(di==0).spec.template.spec.containers[0].securityContext.capabilities.drop += ["NET_RAW"] | + select(di==0).spec.template.spec.containers[0].securityContext.seccompProfile.type |= "RuntimeDefault" | + select(di==0).spec.template.spec.containers[0].securityContext.allowPrivilegeEscalation |= false | + select(di==0).spec.template.spec.containers[0].securityContext.readOnlyRootFilesystem |= true`, - for scenario := 1; scenario <= scenarioCount; scenario++ { - originalFile := fmt.Sprintf("original_yaml_scenario_%d.yml", scenario) - fixedFile := fmt.Sprintf("fixed_yaml_scenario_%d.yml", scenario) + "insert_scenarios/fixed_yaml_scenario_4.yml", + }, + { + "insert_scenarios/original_yaml_scenario_5.yml", + "select(di==0).spec.containers[0].securityContext.allowPrivilegeEscalation |= false", + "insert_scenarios/fixed_yaml_scenario_5.yml", + }, + { + "insert_scenarios/original_yaml_scenario_6.yml", + "select(di==0).spec.containers[0].securityContext.capabilities.drop += [\"SYS_ADM\"]", + "insert_scenarios/fixed_yaml_scenario_6.yml", + }, + { + "insert_scenarios/original_yaml_scenario_7.yml", - originalFilePath := filepath.Join(directoryPath, originalFile) - fixedFilePath := filepath.Join(directoryPath, fixedFile) + `select(di==0).spec.containers[0].securityContext.allowPrivilegeEscalation |= false | + select(di==1).spec.containers[0].securityContext.allowPrivilegeEscalation |= false`, - // create temp file - tempFile, err := ioutil.TempFile("", originalFile) - if err != nil { - panic(err) - } - defer os.Remove(tempFile.Name()) + "insert_scenarios/fixed_yaml_scenario_7.yml", + }, - // read original file - originalFileContent, err := ioutil.ReadFile(originalFilePath) - if err != nil { - panic(err) + // Removal Scenarios + + { + "remove_scenarios/original_yaml_scenario_1.yml", + "del(select(di==0).spec.containers[0].securityContext)", + "remove_scenarios/fixed_yaml_scenario_1.yml", + }, + { + "remove_scenarios/original_yaml_scenario_2.yml", + "del(select(di==0).spec.containers[1])", + "remove_scenarios/fixed_yaml_scenario_2.yml", + }, + { + "remove_scenarios/original_yaml_scenario_3.yml", + "del(select(di==0).spec.containers[0].securityContext.capabilities.drop[1])", + "remove_scenarios/fixed_yaml_scenario_3.yml", + }, + { + "remove_scenarios/original_yaml_scenario_4.yml", + `del(select(di==0).spec.containers[0].securityContext) | + del(select(di==1).spec.containers[1])`, + "remove_scenarios/fixed_yaml_scenario_4.yml", + }, + + // Replace Scenarios + { + "replace_scenarios/original_yaml_scenario_1.yml", + "select(di==0).spec.containers[0].securityContext.runAsRoot |= false", + "replace_scenarios/fixed_yaml_scenario_1.yml", + }, + { + "replace_scenarios/original_yaml_scenario_2.yml", + `select(di==0).spec.containers[0].securityContext.capabilities.drop[0] |= "SYS_ADM" | + select(di==0).spec.containers[0].securityContext.capabilities.add[0] |= "NET_RAW"`, + "replace_scenarios/fixed_yaml_scenario_2.yml", + }, + + // Hybrid Scenarios + { + "hybrid_scenarios/original_yaml_scenario_1.yml", + `del(select(di==0).spec.containers[0].securityContext) | + select(di==0).spec.securityContext.runAsRoot |= false`, + "hybrid_scenarios/fixed_yaml_scenario_1.yml", + }, + } + + return indentationTestCases +} + +func TestApplyFixKeepsIndentation(t *testing.T) { + testCases := getTestCases() + + for _, tc := range testCases { + getTestDataPath := func(filename string) string { + currentDir, _ := os.Getwd() + currentFile := "testdata/" + filename + return filepath.Join(currentDir, currentFile) } - // write original file contents to temp file - err = ioutil.WriteFile(tempFile.Name(), originalFileContent, 0644) - if err != nil { - panic(err) - } + input, _ := os.ReadFile(getTestDataPath(tc.inputFile)) + want, _ := os.ReadFile(getTestDataPath(tc.expectedFile)) + expression := tc.yamlExpression - // make changes to temp file h, _ := NewFixHandlerMock() - filePathFixInfo := make(map[string]*fileFixInfo) - filePath := tempFile.Name() - filePathFixInfo[filePath] = &fileFixInfo{ - contentToAdd: make([]contentToAdd, 0), - linesToRemove: make([]linesToRemove, 0), + got, _ := h.ApplyFix(string(input), expression) + + if got != string(want) { + t.Errorf("Fixed file does not match the expected.\nGot: <%s>\nWant:<%s>", got, want) } - fixInfo := filePathFixInfo[filePath] - - for idx, yamlExpression := range (*yamlExpressions)[scenario-1] { - contentToAdd, linesToRemove, err := h.getResourceFileFix(filePath, yamlExpression, idx) - - if err == nil { - h.addResourceFileFix(contentToAdd, linesToRemove, fixInfo) - } - - } - - err = h.applyFixToFiles(filePathFixInfo) - assert.NoError(t, err) - - // Check temp file contents - tempFileContent, err := ioutil.ReadFile(tempFile.Name()) - if err != nil { - panic(err) - } - - // Get fixed Yaml file content and check if it is equal to tempFileContent - fixedFileContent, err := ioutil.ReadFile(fixedFilePath) - - errorMessage := fmt.Sprintf("Content of fixed %s doesn't match content of %s in %s", originalFile, fixedFile, directoryPath) - - assert.Equal(t, string(fixedFileContent), string(tempFileContent), errorMessage) } } -func testDirectoryApplyFix(t *testing.T, directory string) { - directoryPath := filepath.Join(getTestdataPath(), directory) - var yamlExpressions [][]string - - switch directory { - case "insert_scenarios": - yamlExpressions = [][]string{ - {"select(di==0).spec.containers[0].securityContext.allowPrivilegeEscalation |= false"}, - - {"select(di==0).spec.containers[0].securityContext.capabilities.drop += [\"NET_RAW\"]"}, - - {"select(di==0).spec.containers[0].securityContext.capabilities.drop += [\"SYS_ADM\"]"}, - - {`select(di==0).spec.template.spec.securityContext.allowPrivilegeEscalation |= false | - select(di==0).spec.template.spec.containers[0].securityContext.capabilities.drop += ["NET_RAW"] | - select(di==0).spec.template.spec.containers[0].securityContext.seccompProfile.type |= "RuntimeDefault" | - select(di==0).spec.template.spec.containers[0].securityContext.allowPrivilegeEscalation |= false | - select(di==0).spec.template.spec.containers[0].securityContext.readOnlyRootFilesystem |= true`}, - - {"select(di==0).spec.containers[0].securityContext.allowPrivilegeEscalation |= false"}, - - {"select(di==0).spec.containers[0].securityContext.capabilities.drop += [\"SYS_ADM\"]"}, - - { - "select(di==0).spec.containers[0].securityContext.allowPrivilegeEscalation |= false", - "select(di==1).spec.containers[0].securityContext.allowPrivilegeEscalation |= false", - }, - } - - case "remove_scenarios": - yamlExpressions = [][]string{ - {"del(select(di==0).spec.containers[0].securityContext)"}, - - {"del(select(di==0).spec.containers[1])"}, - - {"del(select(di==0).spec.containers[0].securityContext.capabilities.drop[1])"}, - - { - "del(select(di==0).spec.containers[0].securityContext)", - "del(select(di==1).spec.containers[1])", - }, - } - - case "replace_scenarios": - yamlExpressions = [][]string{ - {"select(di==0).spec.containers[0].securityContext.runAsRoot |= false"}, - - {`select(di==0).spec.containers[0].securityContext.capabilities.drop[0] |= "SYS_ADM" | - select(di==0).spec.containers[0].securityContext.capabilities.add[0] |= "NET_RAW"`}, - } - - case "hybrid_scenarios": - yamlExpressions = [][]string{ - {`del(select(di==0).spec.containers[0].securityContext) | - select(di==0).spec.securityContext.runAsRoot |= false`}, - } - } - - testDirectoryApplyFixHelper(t, &yamlExpressions, directoryPath) -} - -func TestFixHandler_applyFixToFile(t *testing.T) { - // Tests for Insert scenarios - testDirectoryApplyFix(t, "insert_scenarios") - - // Tests for Removal scenarios - testDirectoryApplyFix(t, "remove_scenarios") - - // Tests for Replace scenarios - testDirectoryApplyFix(t, "replace_scenarios") - - // Tests for Hybrid Scenarios - testDirectoryApplyFix(t, "hybrid_scenarios") - -} - func Test_fixPathToValidYamlExpression(t *testing.T) { type args struct { fixPath string diff --git a/core/pkg/fixhandler/testdata/hybrid_scenarios/fixed_yaml_scenario_1.yml b/core/pkg/fixhandler/testdata/hybrid_scenarios/fixed_yaml_scenario_1.yml index b74c6cdc..5ab9f08c 100644 --- a/core/pkg/fixhandler/testdata/hybrid_scenarios/fixed_yaml_scenario_1.yml +++ b/core/pkg/fixhandler/testdata/hybrid_scenarios/fixed_yaml_scenario_1.yml @@ -16,4 +16,4 @@ spec: - name: nginx_container image: nginx securityContext: - runAsRoot: false + runAsRoot: false \ No newline at end of file diff --git a/core/pkg/fixhandler/testdata/insert_scenarios/fixed_yaml_scenario_1.yml b/core/pkg/fixhandler/testdata/insert_scenarios/fixed_yaml_scenario_1.yml index dbe3c81e..55f3aeb5 100644 --- a/core/pkg/fixhandler/testdata/insert_scenarios/fixed_yaml_scenario_1.yml +++ b/core/pkg/fixhandler/testdata/insert_scenarios/fixed_yaml_scenario_1.yml @@ -11,4 +11,4 @@ spec: - name: nginx_container image: nginx securityContext: - allowPrivilegeEscalation: false + allowPrivilegeEscalation: false \ No newline at end of file diff --git a/core/pkg/fixhandler/testdata/insert_scenarios/fixed_yaml_scenario_2.yml b/core/pkg/fixhandler/testdata/insert_scenarios/fixed_yaml_scenario_2.yml index f0d77a9f..ba084424 100644 --- a/core/pkg/fixhandler/testdata/insert_scenarios/fixed_yaml_scenario_2.yml +++ b/core/pkg/fixhandler/testdata/insert_scenarios/fixed_yaml_scenario_2.yml @@ -12,4 +12,4 @@ spec: securityContext: capabilities: drop: - - NET_RAW + - NET_RAW \ No newline at end of file diff --git a/core/pkg/fixhandler/testdata/insert_scenarios/fixed_yaml_scenario_3.yml b/core/pkg/fixhandler/testdata/insert_scenarios/fixed_yaml_scenario_3.yml index 60b4f0ed..e673574d 100644 --- a/core/pkg/fixhandler/testdata/insert_scenarios/fixed_yaml_scenario_3.yml +++ b/core/pkg/fixhandler/testdata/insert_scenarios/fixed_yaml_scenario_3.yml @@ -13,4 +13,4 @@ spec: capabilities: drop: - NET_RAW - - SYS_ADM + - SYS_ADM \ No newline at end of file diff --git a/core/pkg/fixhandler/testdata/insert_scenarios/fixed_yaml_scenario_5.yml b/core/pkg/fixhandler/testdata/insert_scenarios/fixed_yaml_scenario_5.yml index 6b66f26b..c9917375 100644 --- a/core/pkg/fixhandler/testdata/insert_scenarios/fixed_yaml_scenario_5.yml +++ b/core/pkg/fixhandler/testdata/insert_scenarios/fixed_yaml_scenario_5.yml @@ -15,4 +15,4 @@ spec: # Testing if comments are retained as intended securityContext: - runAsRoot: false + runAsRoot: false \ No newline at end of file diff --git a/core/pkg/fixhandler/testdata/insert_scenarios/fixed_yaml_scenario_6.yml b/core/pkg/fixhandler/testdata/insert_scenarios/fixed_yaml_scenario_6.yml index b0091e62..d64fecd3 100644 --- a/core/pkg/fixhandler/testdata/insert_scenarios/fixed_yaml_scenario_6.yml +++ b/core/pkg/fixhandler/testdata/insert_scenarios/fixed_yaml_scenario_6.yml @@ -11,4 +11,4 @@ spec: image: nginx securityContext: capabilities: - drop: [NET_RAW, SYS_ADM] + drop: [NET_RAW, SYS_ADM] \ No newline at end of file diff --git a/core/pkg/fixhandler/testdata/remove_scenarios/fixed_yaml_scenario_1.yml b/core/pkg/fixhandler/testdata/remove_scenarios/fixed_yaml_scenario_1.yml index 5c5f896c..8ecc8724 100644 --- a/core/pkg/fixhandler/testdata/remove_scenarios/fixed_yaml_scenario_1.yml +++ b/core/pkg/fixhandler/testdata/remove_scenarios/fixed_yaml_scenario_1.yml @@ -9,4 +9,4 @@ metadata: spec: containers: - name: nginx_container - image: nginx + image: nginx \ No newline at end of file diff --git a/core/pkg/fixhandler/testdata/remove_scenarios/fixed_yaml_scenario_2.yml b/core/pkg/fixhandler/testdata/remove_scenarios/fixed_yaml_scenario_2.yml index 43b80fa1..56fc35f8 100644 --- a/core/pkg/fixhandler/testdata/remove_scenarios/fixed_yaml_scenario_2.yml +++ b/core/pkg/fixhandler/testdata/remove_scenarios/fixed_yaml_scenario_2.yml @@ -10,4 +10,3 @@ spec: containers: - name: nginx_container image: nginx - diff --git a/core/pkg/fixhandler/testdata/remove_scenarios/fixed_yaml_scenario_3.yml b/core/pkg/fixhandler/testdata/remove_scenarios/fixed_yaml_scenario_3.yml index bc2225bd..b0899e2c 100644 --- a/core/pkg/fixhandler/testdata/remove_scenarios/fixed_yaml_scenario_3.yml +++ b/core/pkg/fixhandler/testdata/remove_scenarios/fixed_yaml_scenario_3.yml @@ -11,4 +11,4 @@ spec: image: nginx securityContext: capabilities: - drop: ["NET_RAW"] + drop: ["NET_RAW"] \ No newline at end of file diff --git a/core/pkg/fixhandler/testdata/remove_scenarios/fixed_yaml_scenario_4.yml b/core/pkg/fixhandler/testdata/remove_scenarios/fixed_yaml_scenario_4.yml index 5bfa9988..fd24bbc9 100644 --- a/core/pkg/fixhandler/testdata/remove_scenarios/fixed_yaml_scenario_4.yml +++ b/core/pkg/fixhandler/testdata/remove_scenarios/fixed_yaml_scenario_4.yml @@ -25,4 +25,3 @@ spec: containers: - name: nginx_container image: nginx - diff --git a/core/pkg/fixhandler/testdata/replace_scenarios/fixed_yaml_scenario_1.yml b/core/pkg/fixhandler/testdata/replace_scenarios/fixed_yaml_scenario_1.yml index 067565ca..4a4eda33 100644 --- a/core/pkg/fixhandler/testdata/replace_scenarios/fixed_yaml_scenario_1.yml +++ b/core/pkg/fixhandler/testdata/replace_scenarios/fixed_yaml_scenario_1.yml @@ -11,4 +11,4 @@ spec: - name: nginx_container image: nginx securityContext: - runAsRoot: false + runAsRoot: false \ No newline at end of file diff --git a/core/pkg/fixhandler/testdata/replace_scenarios/fixed_yaml_scenario_2.yml b/core/pkg/fixhandler/testdata/replace_scenarios/fixed_yaml_scenario_2.yml index 2c8b1da4..4ea9b48c 100644 --- a/core/pkg/fixhandler/testdata/replace_scenarios/fixed_yaml_scenario_2.yml +++ b/core/pkg/fixhandler/testdata/replace_scenarios/fixed_yaml_scenario_2.yml @@ -15,4 +15,4 @@ spec: capabilities: drop: - "SYS_ADM" - add: ["NET_RAW"] + add: ["NET_RAW"] \ No newline at end of file diff --git a/core/pkg/fixhandler/yamlhandler.go b/core/pkg/fixhandler/yamlhandler.go index 39e85832..1c22f24f 100644 --- a/core/pkg/fixhandler/yamlhandler.go +++ b/core/pkg/fixhandler/yamlhandler.go @@ -1,14 +1,11 @@ package fixhandler import ( - "bufio" - "bytes" "container/list" "errors" "fmt" "io" - "io/ioutil" - "os" + "strings" logger "github.com/kubescape/go-logger" "github.com/mikefarah/yq/v4/pkg/yqlib" @@ -16,44 +13,37 @@ import ( "gopkg.in/yaml.v3" ) -func constructDecodedYaml(filepath string) *[]yaml.Node { - file, err := ioutil.ReadFile(filepath) - if err != nil { - logger.L().Fatal("Cannot read file") - } - fileReader := bytes.NewReader(file) +func constructDecodedYaml(yamlString string) *[]yaml.Node { + fileReader := strings.NewReader(yamlString) dec := yaml.NewDecoder(fileReader) nodes := make([]yaml.Node, 0) for { var node yaml.Node - err = dec.Decode(&node) + err := dec.Decode(&node) + nodes = append(nodes, node) - // break the loop in case of EOF if errors.Is(err, io.EOF) { break } if err != nil { - panic(err) + logger.L().Fatal("Cannot decode given document") } } return &nodes } -func constructFixedYamlNodes(filePath, yamlExpression string) (*[]yaml.Node, error) { +func constructFixedYamlNodes(yamlString, yamlExpression string) (*[]yaml.Node, error) { preferences := yqlib.ConfiguredYamlPreferences preferences.EvaluateTogether = true decoder := yqlib.NewYamlDecoder(preferences) var allDocuments = list.New() - reader, err := constructNewReader(filePath) - if err != nil { - return nil, err - } + reader := strings.NewReader(yamlString) - fileDocuments, err := readDocuments(reader, filePath, 0, decoder) + fileDocuments, err := readDocuments(reader, decoder) if err != nil { return nil, err } @@ -103,7 +93,22 @@ func matchNodes(nodeOne, nodeTwo *yaml.Node) int { } } -func getFixInfo(originalList, fixedList *[]nodeInfo) (*[]contentToAdd, *[]linesToRemove) { +func getFixInfo(originalRootNodes, fixedRootNodes *[]yaml.Node) (*[]contentToAdd, *[]linesToRemove) { + contentToAdd := make([]contentToAdd, 0) + linesToRemove := make([]linesToRemove, 0) + + for idx, _ := range *fixedRootNodes { + originalList := constructDFSOrder(&(*originalRootNodes)[idx]) + fixedList := constructDFSOrder(&(*fixedRootNodes)[idx]) + nodeContentToAdd, nodeLinesToRemove := getFixInfoHelper(originalList, fixedList) + contentToAdd = append(contentToAdd, *nodeContentToAdd...) + linesToRemove = append(linesToRemove, *nodeLinesToRemove...) + } + + return &contentToAdd, &linesToRemove +} + +func getFixInfoHelper(originalList, fixedList *[]nodeInfo) (*[]contentToAdd, *[]linesToRemove) { // While obtaining fixedYamlNode, comments and empty lines at the top are ignored. // This causes a difference in Line numbers across the tree structure. In order to @@ -232,71 +237,53 @@ func updateLinesToReplace(fixInfoMetadata *fixInfoMetadata) (int, int) { return updatedOriginalTracker, updatedFixedTracker } -func applyFixesToFile(filePath string, contentToAdd *[]contentToAdd, linesToRemove *[]linesToRemove) error { - // Read contents of the file line by line and store in a list - linesSlice, err := getLinesSlice(filePath) - - if err != nil { - return err - } - - // Determining last line required lineSlice. The placeholder for last line is replaced with the real last line - assignLastLine(contentToAdd, linesToRemove, &linesSlice) - - // Clear the current content of file - if err := os.Truncate(filePath, 0); err != nil { - return err - } - - file, err := os.OpenFile(filePath, os.O_RDWR, 0644) - if err != nil { - return err - } - - defer func() error { - if err := file.Close(); err != nil { - return err +func removeNewLinesAtTheEnd(yamlLines []string) []string { + for idx := 1; idx < len(yamlLines); idx++ { + if yamlLines[len(yamlLines)-idx] != "\n" { + yamlLines = yamlLines[:len(yamlLines)-idx+1] + break } - return nil - }() + } + return yamlLines +} - removeLines(linesToRemove, &linesSlice) +func getFixedYamlLines(yamlLines []string, contentToAdd *[]contentToAdd, linesToRemove *[]linesToRemove) (fixedYamlLines []string) { - writer := bufio.NewWriter(file) + // Determining last line requires original yaml lines slice. The placeholder for last line is replaced with the real last line + assignLastLine(contentToAdd, linesToRemove, &yamlLines) + + removeLines(linesToRemove, &yamlLines) + + fixedYamlLines = make([]string, 0) lineIdx, lineToAddIdx := 1, 0 // Ideally, new node is inserted at line before the next node in DFS order. But, when the previous line contains a // comment or empty line, we need to insert new nodes before them. - adjustContentLines(contentToAdd, &linesSlice) + adjustContentLines(contentToAdd, &yamlLines) for lineToAddIdx < len(*contentToAdd) { for lineIdx <= (*contentToAdd)[lineToAddIdx].line { // Check if the current line is not removed - if linesSlice[lineIdx-1] != "*" { - _, err := writer.WriteString(linesSlice[lineIdx-1] + "\n") - if err != nil { - return err - } + if yamlLines[lineIdx-1] != "*" { + fixedYamlLines = append(fixedYamlLines, yamlLines[lineIdx-1]) } lineIdx += 1 } content := (*contentToAdd)[lineToAddIdx].content - writer.WriteString(content) + fixedYamlLines = append(fixedYamlLines, content) lineToAddIdx += 1 } - for lineIdx <= len(linesSlice) { - if linesSlice[lineIdx-1] != "*" { - _, err := writer.WriteString(linesSlice[lineIdx-1] + "\n") - if err != nil { - return err - } + for lineIdx <= len(yamlLines) { + if yamlLines[lineIdx-1] != "*" { + fixedYamlLines = append(fixedYamlLines, yamlLines[lineIdx-1]) } lineIdx += 1 } - writer.Flush() - return nil + fixedYamlLines = removeNewLinesAtTheEnd(fixedYamlLines) + + return fixedYamlLines } diff --git a/core/pkg/fixhandler/yamlhelper.go b/core/pkg/fixhandler/yamlhelper.go index e817e503..d566ae9a 100644 --- a/core/pkg/fixhandler/yamlhelper.go +++ b/core/pkg/fixhandler/yamlhelper.go @@ -129,7 +129,7 @@ func constructContent(parentNode *yaml.Node, nodeList *[]nodeInfo, tracker int) content = indentContent(content, indentationSpaces) - return content + return strings.TrimSuffix(content, "\n") } func indentContent(content string, indentationSpaces int) string { @@ -306,12 +306,13 @@ func isEmptyLineOrComment(lineContent string) bool { return false } -func readDocuments(reader io.Reader, filename string, fileIndex int, decoder yqlib.Decoder) (*list.List, error) { +func readDocuments(reader io.Reader, decoder yqlib.Decoder) (*list.List, error) { err := decoder.Init(reader) if err != nil { return nil, err } inputList := list.New() + var currentIndex uint for { @@ -324,11 +325,10 @@ func readDocuments(reader io.Reader, filename string, fileIndex int, decoder yql } return inputList, nil } else if errorReading != nil { - return nil, fmt.Errorf("bad file '%v': %w", filename, errorReading) + return nil, fmt.Errorf("Error Decoding YAML file") } + candidateNode.Document = currentIndex - candidateNode.Filename = filename - candidateNode.FileIndex = fileIndex candidateNode.EvaluateTogether = true inputList.PushBack(candidateNode) @@ -484,3 +484,7 @@ func updateTracker(nodeList *[]nodeInfo, tracker int) int { return updatedTracker } + +func getFixedYamlString(yamlLines []string) (fixedYamlString string) { + return strings.Join(yamlLines, "\n") +}