add key selector argument to replace (#74)

this adds a `-s`/`--key-selector` argument to the `replace` command. in non-regexp mode, it will specify a complete key to match against. in regexp mode, if the key selector yields any matches it will be considered be successful. in this case, replacements will be evaluated against that key-value pair. this pull also fixes two issues: - improper evaluation of the add/remove markup on a line diff during replace - indexes being generated for regexps (index+length) did not match the indexes being generated for substrings (start&end) resulting in a panic for matches longer than half the string
fishi0x01 · Feb 6, 2021 · d4c496e · d4c496e
1 parent 5ed1676
commit d4c496e
Show file tree

Hide file tree

Showing 5 changed files with 105 additions and 30 deletions.
diff --git a/cli/replace.go b/cli/replace.go
@@ -24,7 +24,8 @@ type ReplaceCommandArgs struct {
 	Search      string `arg:"positional,required"`
 	Replacement string `arg:"positional,required"`
 	Path        string `arg:"positional,required"`
-	Regexp      bool   `arg:"-e,--regexp" help:"Treat search string as a regexp"`
+	Regexp      bool   `arg:"-e,--regexp" help:"Treat search string and selector as a regexp"`
+	KeySelector string `arg:"-s,--key-selector" help:"Limit replacements to specified key" placeholder:"PATTERN"`
 	Keys        bool   `arg:"-k,--keys" help:"Match against keys (true if -v is not specified)"`
 	Values      bool   `arg:"-v,--values" help:"Match against values (true if -k is not specified)"`
 	Confirm     bool   `arg:"-y,--confirm" help:"Write results without prompt"`
@@ -68,10 +69,11 @@ func (cmd *ReplaceCommand) PrintUsage() {
 // GetSearchParams returns the search parameters the command was run with
 func (cmd *ReplaceCommand) GetSearchParams() SearchParameters {
 	return SearchParameters{
-		Search:      cmd.args.Search,
-		Replacement: &cmd.args.Replacement,
-		Mode:        cmd.Mode,
 		IsRegexp:    cmd.args.Regexp,
+		KeySelector: cmd.args.KeySelector,
+		Mode:        cmd.Mode,
+		Replacement: &cmd.args.Replacement,
+		Search:      cmd.args.Search,
 	}
 }
 

diff --git a/cli/search.go b/cli/search.go
@@ -21,6 +21,7 @@ type SearchingCommand interface {
 type SearchParameters struct {
 	Search      string
 	Replacement *string
+	KeySelector string
 	Mode        KeyValueMode
 	IsRegexp    bool
 }
@@ -44,13 +45,14 @@ type Match struct {
 // Searcher provides matching and replacement methods while maintaining references to the command
 // that provides an interface to search operations. Also maintains reference to a compiled regexp.
 type Searcher struct {
-	cmd    SearchingCommand
-	regexp *regexp.Regexp
+	cmd           SearchingCommand
+	regexp        *regexp.Regexp
+	keySelectorRe *regexp.Regexp
 }
 
 // NewSearcher creates a new Searcher container for performing search and optionally replace
 func NewSearcher(cmd SearchingCommand) (*Searcher, error) {
-	var re *regexp.Regexp
+	var re, keySelectorRe *regexp.Regexp
 	var err error
 	params := cmd.GetSearchParams()
 
@@ -60,8 +62,14 @@ func NewSearcher(cmd SearchingCommand) (*Searcher, error) {
 			return nil, fmt.Errorf("cannot parse regex pattern")
 		}
 	}
+	if params.KeySelector != "" && params.IsRegexp == true {
+		keySelectorRe, err = regexp.Compile(params.KeySelector)
+		if err != nil {
+			return nil, fmt.Errorf("key-selector: %s", err)
+		}
+	}
 
-	return &Searcher{cmd: cmd, regexp: re}, nil
+	return &Searcher{cmd: cmd, regexp: re, keySelectorRe: keySelectorRe}, nil
 }
 
 // IsMode returns true if the specified mode is enabled
@@ -74,14 +82,23 @@ func (s *Searcher) DoSearch(path string, k string, v string) (m []*Match) {
 	// Default to original strings
 	replacedKey, keyLineDiff := k, k
 	replacedValue, valueLineDiff := v, v
-	var keyMatchPairs, valueMatchPairs [][]int
+	var keyMatchPairs, valueMatchPairs, keySelectorMatches [][]int
 
+	if s.cmd.GetSearchParams().KeySelector != "" {
+		keySelectorMatches = s.keySelectorMatches(k)
+		if len(keySelectorMatches) == 0 {
+			return m
+		}
+	}
 	if s.IsMode(ModeKeys) {
-		keyMatchPairs, replacedKey, keyLineDiff = s.matchData(k, s.cmd.GetSearchParams().IsRegexp)
+		keyMatchPairs, replacedKey, keyLineDiff = s.matchData(k)
+	}
+	if len(keySelectorMatches) > 0 {
+		keyLineDiff = highlightMatches(keyLineDiff, s.keySelectorMatches(keyLineDiff))
 	}
 
 	if s.IsMode(ModeValues) {
-		valueMatchPairs, replacedValue, valueLineDiff = s.matchData(v, s.cmd.GetSearchParams().IsRegexp)
+		valueMatchPairs, replacedValue, valueLineDiff = s.matchData(v)
 	}
 
 	if len(keyMatchPairs) > 0 || len(valueMatchPairs) > 0 {
@@ -106,18 +123,28 @@ func (match *Match) print(out io.Writer, diff bool) {
 	if diff == true {
 		fmt.Fprintf(out, "%s> %s = %s\n", match.path, match.keyLineDiff, match.valueLineDiff)
 	} else {
-		fmt.Fprintf(out, "%s> %s = %s\n", match.path, match.highlightMatches(match.key, match.keyIndex), match.highlightMatches(match.value, match.valueIndex))
+		fmt.Fprintf(out, "%s> %s = %s\n", match.path, highlightMatches(match.key, match.keyIndex), highlightMatches(match.value, match.valueIndex))
 	}
 }
 
-// highlightMatches will take an array of index and lens and highlight them
-func (match *Match) highlightMatches(s string, matches [][]int) (result string) {
+// keySelectorMatches provides an array of start and end indexes of key selector matches
+func (s *Searcher) keySelectorMatches(k string) (matches [][]int) {
+	if s.cmd.GetSearchParams().IsRegexp == true {
+		return s.keySelectorRe.FindAllStringIndex(k, -1)
+	}
+	if k == s.cmd.GetSearchParams().KeySelector {
+		return [][]int{[]int{0, len(k)}}
+	}
+	return [][]int{}
+}
+
+// highlightMatches will take an array of start and end indexes and highlight them
+func highlightMatches(s string, matches [][]int) (result string) {
 	cur := 0
 	if len(matches) > 0 {
 		for _, pair := range matches {
 			next := pair[0]
-			length := pair[1]
-			end := next + length
+			end := pair[1]
 			result += s[cur:next]
 			result += color.New(color.FgYellow).SprintFunc()(s[next:end])
 			cur = end
@@ -138,14 +165,16 @@ func (s *Searcher) highlightLineDiff(d string) string {
 
 	for _, b := range []byte(d) {
 		buf = append(buf, b)
-		if string(buf) == "(~~" && !removeMode && !addMode {
+		if len(buf) >= 3 && string(buf[len(buf)-3:len(buf)]) == "(~~" && !removeMode && !addMode {
+			res = append(res, buf[0:len(buf)-3]...)
 			buf = make([]byte, 0)
 			removeMode = true
 		} else if len(buf) > 3 && string(buf[len(buf)-3:]) == "~~)" && removeMode {
 			res = append(res, removeColor.SprintFunc()(string(buf[0:len(buf)-3]))...)
 			buf = make([]byte, 0)
 			removeMode = false
-		} else if string(buf) == "(++" && !removeMode && !addMode {
+		} else if len(buf) >= 3 && string(buf[len(buf)-3:len(buf)]) == "(++" && !removeMode && !addMode {
+			res = append(res, buf[0:len(buf)-3]...)
 			buf = make([]byte, 0)
 			addMode = true
 		} else if len(buf) > 3 && string(buf[len(buf)-3:]) == "++)" && addMode {
@@ -157,29 +186,39 @@ func (s *Searcher) highlightLineDiff(d string) string {
 	return string(append(res, buf...))
 }
 
-func (s *Searcher) matchData(subject string, isRegexp bool) (matchPairs [][]int, replaced string, inlineDiff string) {
+func (s *Searcher) substrMatchData(subject string, search string) (matchPairs [][]int) {
+	index := suffixarray.New([]byte(subject))
+	matches := index.Lookup([]byte(search), -1)
+	sort.Ints(matches)
+	substrLength := len(search)
+	for _, offset := range matches {
+		matchPairs = append(matchPairs, []int{offset, offset + substrLength})
+	}
+	return matchPairs
+}
+
+func (s *Searcher) regexpMatchData(subject string, re *regexp.Regexp) (matchPairs [][]int) {
+	return re.FindAllStringIndex(subject, -1)
+}
+
+func (s *Searcher) matchData(subject string) (matchPairs [][]int, replaced string, inlineDiff string) {
 	replaced, inlineDiff = subject, subject
 	matchPairs = make([][]int, 0)
 
-	if isRegexp {
-		matchPairs = s.regexp.FindAllIndex([]byte(subject), -1)
+	if s.cmd.GetSearchParams().IsRegexp {
+		matchPairs = s.regexpMatchData(subject, s.regexp)
 	} else {
-		index := suffixarray.New([]byte(subject))
-		matches := index.Lookup([]byte(s.cmd.GetSearchParams().Search), -1)
-		sort.Ints(matches)
-		substrLength := len(s.cmd.GetSearchParams().Search)
-		for _, offset := range matches {
-			matchPairs = append(matchPairs, []int{offset, substrLength})
-		}
+		matchPairs = s.substrMatchData(subject, s.cmd.GetSearchParams().Search)
 	}
 
 	if s.cmd.GetSearchParams().Replacement != nil {
-		if isRegexp {
+		if s.cmd.GetSearchParams().IsRegexp {
 			replaced = s.regexp.ReplaceAllString(subject, *s.cmd.GetSearchParams().Replacement)
 		} else {
 			replaced = strings.ReplaceAll(subject, s.cmd.GetSearchParams().Search, *s.cmd.GetSearchParams().Replacement)
 		}
 		inlineDiff = s.highlightLineDiff(diff.CharacterDiff(subject, replaced))
 	}
+
 	return matchPairs, replaced, inlineDiff
 }
diff --git a/test/suites/commands/grep.bats b/test/suites/commands/grep.bats
@@ -47,6 +47,12 @@ load ../../bin/plugins/bats-assert/load
   assert_line --partial "cannot parse regex"
   assert_failure 1
 
+  #######################################
+  echo "==== case: regex pattern on a long value ===="
+  run ${APP_BIN} -c "grep -e 'value-for-testing' ${KV_BACKEND}/src/a/foo"
+  assert_line --partial this-is-a-really-long-value-for-testing
+  assert_success
+
   #######################################
   echo "==== case: pattern with spaces ===="
   run ${APP_BIN} -c "grep 'a spaced val' ${KV_BACKEND}/src/spaces"

diff --git a/test/suites/commands/replace.bats b/test/suites/commands/replace.bats
@@ -64,6 +64,16 @@ load ../../bin/plugins/bats-assert/load
   assert_line exhibit
   run get_vault_value "value" "${KV_BACKEND}/src/prod/all"
   assert_line all
+
+  #######################################
+  echo "==== case: replace value in single path with selector ===="
+  run ${APP_BIN} -c "replace -s 'produce' 'apple' 'orange' ${KV_BACKEND}/src/selector/1 -y"
+  assert_success
+  assert_line "Writing!"
+  run get_vault_value "produce" "${KV_BACKEND}/src/selector/1"
+  assert_line orange
+  run get_vault_value "fruit" "${KV_BACKEND}/src/selector/1"
+  assert_line apple
 }
 
 @test "vault-${VAULT_VERSION} ${KV_BACKEND} 'replace' regexp" {
@@ -110,4 +120,20 @@ load ../../bin/plugins/bats-assert/load
   assert_line testexhibit
   run get_vault_value "value" "${KV_BACKEND}/src/prod/all"
   assert_line all
+
+  #######################################
+  echo "==== case: replace value in single path with selector ===="
+  run ${APP_BIN} -c "replace -e -s 'prod.*' '^apple' 'orange' ${KV_BACKEND}/src/selector/1 -y"
+  assert_success
+  assert_line "Writing!"
+  run get_vault_value "produce" "${KV_BACKEND}/src/selector/1"
+  assert_line orange
+  run get_vault_value "fruit" "${KV_BACKEND}/src/selector/1"
+  assert_line apple
+
+  #######################################
+  echo "==== case: replace fails with bad regex selector ===="
+  run ${APP_BIN} -c "replace -e -s '][' '^apple' 'orange' ${KV_BACKEND}/src/selector/1 -y"
+  assert_failure
+  assert_line  --partial "key-selector: error parsing regexp"
 }
diff --git a/test/util/util.bash b/test/util/util.bash
@@ -53,7 +53,9 @@ setup() {
         vault_exec "vault kv put ${kv_backend}/src/tooling/v2 value=v2 drink=water key=C"
         vault_exec "vault kv put ${kv_backend}/src/ambivalence/1 value=1 fruit=apple"
         vault_exec "vault kv put ${kv_backend}/src/ambivalence/1/a value=2 fruit=banana"
-        vault_exec "vault kv put ${kv_backend}/src/a/foo value=1"
+        vault_exec "vault kv put ${kv_backend}/src/selector/1 value=1 fruit=apple produce=apple food=apple"
+        vault_exec "vault kv put ${kv_backend}/src/selector/2 value=2 fruit=banana produce=banana food=banana"
+        vault_exec "vault kv put ${kv_backend}/src/a/foo value=1 long=this-is-a-really-long-value-for-testing"
         vault_exec "vault kv put ${kv_backend}/src/a/foo/bar value=2"
         vault_exec "vault kv put ${kv_backend}/src/b/foo value=1"
         vault_exec "vault kv put ${kv_backend}/src/b/foo/bar value=2"