Add append and delete fields to edit step (#416)

Summary: Pull Request resolved: #416 Added `append` and `delete` fields to edit_step Reviewed By: d3sch41n Differential Revision: D51180590 fbshipit-source-id: 0a112844474f3d36c6069d8461644246b5c6fafc
facebookincubator · Nov 14, 2023 · 02c554d · 02c554d
1 parent fb58f05
commit 02c554d
Show file tree

Hide file tree

Showing 2 changed files with 270 additions and 8 deletions.
diff --git a/pkg/blocks/editstep.go b/pkg/blocks/editstep.go
@@ -30,6 +30,8 @@ import (
 type Edit struct {
 	Old    string `yaml:"old,omitempty"`
 	New    string `yaml:"new,omitempty"`
+	Append string `yaml:"append,omitempty"`
+	Delete string `yaml:"delete,omitempty"`
 	Regexp bool   `yaml:"regexp,omitempty"`
 
 	oldRegexp *regexp.Regexp
@@ -65,22 +67,45 @@ func (s *EditStep) Validate(execCtx TTPExecutionContext) error {
 		return fmt.Errorf("no edits specified")
 	}
 
-	// TODO: make this compatible with deleting lines
 	for editIdx, edit := range s.Edits {
-		if edit.Old == "" {
-			return fmt.Errorf("edit #%d is missing 'old:'", editIdx+1)
-		} else if edit.New == "" {
-			return fmt.Errorf("edit #%d is missing 'new:'", editIdx+1)
+
+		if edit.Append == "" && edit.Delete == "" {
+			if edit.Old == "" {
+				return fmt.Errorf("edit #%d is missing 'old:'", editIdx+1)
+			} else if edit.New == "" {
+				return fmt.Errorf("edit #%d is missing 'new:'", editIdx+1)
+			}
+		} else if edit.Append != "" {
+			if edit.Old != "" {
+				return fmt.Errorf("append is not to be used in conjunction with 'old:'")
+			} else if edit.New != "" {
+				return fmt.Errorf("append is not to be used in conjunction with 'new:'")
+
+			} else if edit.Regexp {
+				return fmt.Errorf("append is not to be used in conjunction with 'regexp:'")
+
+			}
+		} else if edit.Delete != "" {
+			if edit.Old != "" {
+				return fmt.Errorf("delete is not to be used in conjunction with 'old:'")
+			} else if edit.New != "" {
+				return fmt.Errorf("delete is not to be used in conjunction with 'new:'")
+
+			}
 		}
 
+		oldStr := edit.Old
+		if edit.Delete != "" {
+			oldStr = edit.Delete
+		}
 		var err error
 		if edit.Regexp {
-			edit.oldRegexp, err = regexp.Compile(edit.Old)
+			edit.oldRegexp, err = regexp.Compile(oldStr)
 			if err != nil {
 				return fmt.Errorf("edit #%d has invalid regex for 'old:'", editIdx+1)
 			}
 		} else {
-			edit.oldRegexp = regexp.MustCompile(regexp.QuoteMeta(edit.Old))
+			edit.oldRegexp = regexp.MustCompile(regexp.QuoteMeta(oldStr))
 		}
 	}
 	return nil
@@ -123,6 +148,12 @@ func (s *EditStep) Execute(execCtx TTPExecutionContext) (*ActResult, error) {
 	// but it's unlikely to be a performance issue in practice. If it is,
 	// we can optimize
 	for editIdx, edit := range s.Edits {
+
+		if edit.Append != "" {
+			contents += "\n" + edit.Append
+			continue
+		}
+
 		matches := edit.oldRegexp.FindAllStringIndex(contents, -1)
 		// we want to error here because otherwise ppl will be confused by silent
 		// failures if the format of the file they're trying to edit changes
@@ -135,7 +166,11 @@ func (s *EditStep) Execute(execCtx TTPExecutionContext) (*ActResult, error) {
 				s.FileToEdit,
 			)
 		}
-		contents = edit.oldRegexp.ReplaceAllString(contents, edit.New)
+		newStr := edit.New
+		if edit.Delete != "" {
+			newStr = ""
+		}
+		contents = edit.oldRegexp.ReplaceAllString(contents, newStr)
 	}
 
 	err = afero.WriteFile(fileSystem, targetPath, []byte(contents), 0644)

diff --git a/pkg/blocks/editstep_test.go b/pkg/blocks/editstep_test.go
@@ -298,3 +298,230 @@ edits:
 	require.Error(t, err, "not finding a search string should result in an error")
 	assert.Equal(t, "pattern 'not_going_to_find_this' from edit #1 was not found in file b.txt", err.Error())
 }
+
+func TestAppendOld(t *testing.T) {
+	content := `name: test append with old
+description: test should fail when they specify append and old at the same time
+steps:
+  - name: append_old
+    edit_file: yolo
+    edits:
+      - old: help
+        new: me
+        append: damn help`
+
+	var ttp TTP
+	err := yaml.Unmarshal([]byte(content), &ttp)
+	require.NoError(t, err)
+
+	err = ttp.Validate(TTPExecutionContext{})
+	require.Error(t, err)
+
+	assert.Equal(t, "append is not to be used in conjunction with 'old:'", err.Error())
+}
+
+func TestAppendNew(t *testing.T) {
+	content := `name: test append with new
+description: test should fail when they specify append and new at the same time
+steps:
+  - name: append_new
+    edit_file: yolo
+    edits:
+      - old:
+        new: me
+        append: damn help`
+
+	var ttp TTP
+	err := yaml.Unmarshal([]byte(content), &ttp)
+	require.NoError(t, err)
+
+	err = ttp.Validate(TTPExecutionContext{})
+	require.Error(t, err)
+
+	assert.Equal(t, "append is not to be used in conjunction with 'new:'", err.Error())
+}
+
+func TestAppendRegex(t *testing.T) {
+	content := `name: test append with regex
+description: test should fail when they specify append and regex at the same time
+steps:
+  - name: append_regex
+    edit_file: yolo
+    edits:
+      - old:
+        new:
+        append: damn help
+        regexp: true`
+
+	var ttp TTP
+	err := yaml.Unmarshal([]byte(content), &ttp)
+	require.NoError(t, err)
+
+	err = ttp.Validate(TTPExecutionContext{})
+	require.Error(t, err)
+
+	assert.Equal(t, "append is not to be used in conjunction with 'regexp:'", err.Error())
+}
+
+func TestDeletedOld(t *testing.T) {
+	content := `name: test delete with old
+description: test should fail when they specify delete and old at the same time
+steps:
+  - name: delete_old
+    edit_file: yolo
+    edits:
+      - old: help
+        new: me
+        delete: you`
+
+	var ttp TTP
+	err := yaml.Unmarshal([]byte(content), &ttp)
+	require.NoError(t, err)
+
+	err = ttp.Validate(TTPExecutionContext{})
+	require.Error(t, err)
+
+	assert.Equal(t, "delete is not to be used in conjunction with 'old:'", err.Error())
+}
+
+func TestDeletedNew(t *testing.T) {
+	content := `name: test delete with new
+description: test should fail when they specify delete and new at the same time
+steps:
+  - name: delete_new
+    edit_file: yolo
+    edits:
+      - old:
+        new: me
+        delete: you`
+
+	var ttp TTP
+	err := yaml.Unmarshal([]byte(content), &ttp)
+	require.NoError(t, err)
+
+	err = ttp.Validate(TTPExecutionContext{})
+	require.Error(t, err)
+
+	assert.Equal(t, "delete is not to be used in conjunction with 'new:'", err.Error())
+}
+
+func TestAppend(t *testing.T) {
+	content := `
+name: append_stuff
+edit_file: a.txt
+edits:
+  - append: put this at the end`
+
+	var step EditStep
+	err := yaml.Unmarshal([]byte(content), &step)
+	require.NoError(t, err)
+
+	testFs := afero.NewMemMapFs()
+	err = afero.WriteFile(testFs, "a.txt", []byte("foo\nanother"), 0644)
+	require.NoError(t, err)
+	step.FileSystem = testFs
+
+	var execCtx TTPExecutionContext
+	err = step.Validate(execCtx)
+	require.NoError(t, err)
+
+	_, err = step.Execute(execCtx)
+	require.NoError(t, err)
+
+	contents, err := afero.ReadFile(testFs, "a.txt")
+	require.NoError(t, err)
+
+	assert.Equal(t, "foo\nanother\nput this at the end", string(contents))
+}
+
+func TestMultipleAppend(t *testing.T) {
+	content := `
+name: append_stuff
+edit_file: a.txt
+edits:
+  - old: foo
+    new: bar
+  - append: baz
+  - old: baz
+    new: wut
+  - append: yo`
+
+	var step EditStep
+	err := yaml.Unmarshal([]byte(content), &step)
+	require.NoError(t, err)
+
+	testFs := afero.NewMemMapFs()
+	err = afero.WriteFile(testFs, "a.txt", []byte("foo\nanother"), 0644)
+	require.NoError(t, err)
+	step.FileSystem = testFs
+
+	var execCtx TTPExecutionContext
+	err = step.Validate(execCtx)
+	require.NoError(t, err)
+
+	_, err = step.Execute(execCtx)
+	require.NoError(t, err)
+
+	contents, err := afero.ReadFile(testFs, "a.txt")
+	require.NoError(t, err)
+
+	assert.Equal(t, "bar\nanother\nwut\nyo", string(contents))
+}
+
+func TestDelete(t *testing.T) {
+	content := `
+name: delete_stuff
+edit_file: a.txt
+edits:
+  - delete: foo`
+
+	var step EditStep
+	err := yaml.Unmarshal([]byte(content), &step)
+	require.NoError(t, err)
+
+	testFs := afero.NewMemMapFs()
+	err = afero.WriteFile(testFs, "a.txt", []byte("foo\nanother"), 0644)
+	require.NoError(t, err)
+	step.FileSystem = testFs
+
+	var execCtx TTPExecutionContext
+	err = step.Validate(execCtx)
+	require.NoError(t, err)
+
+	_, err = step.Execute(execCtx)
+	require.NoError(t, err)
+
+	contents, err := afero.ReadFile(testFs, "a.txt")
+	require.NoError(t, err)
+
+	assert.Equal(t, "\nanother", string(contents))
+}
+
+func TestDeleteLine(t *testing.T) {
+	content := `
+name: delete_entire_line
+edit_file: a.txt
+edits:
+  - delete: "foo\n"`
+
+	var step EditStep
+	err := yaml.Unmarshal([]byte(content), &step)
+	require.NoError(t, err)
+
+	testFs := afero.NewMemMapFs()
+	err = afero.WriteFile(testFs, "a.txt", []byte("foo\nanother"), 0644)
+	require.NoError(t, err)
+	step.FileSystem = testFs
+
+	var execCtx TTPExecutionContext
+	err = step.Validate(execCtx)
+	require.NoError(t, err)
+
+	_, err = step.Execute(execCtx)
+	require.NoError(t, err)
+
+	contents, err := afero.ReadFile(testFs, "a.txt")
+	require.NoError(t, err)
+
+	assert.Equal(t, "another", string(contents))
+}