diff --git a/server/fleet/fleet_vars.go b/server/fleet/fleet_vars.go new file mode 100644 index 0000000000..03ac6de747 --- /dev/null +++ b/server/fleet/fleet_vars.go @@ -0,0 +1,109 @@ +package fleet + +import ( + "os" + "strings" +) + +// ContainsPrefixVars scans a string for variables in the form of $VAR +// and ${VAR} that begin with prefix, and return an array of those +// variables with the prefix removed. +func ContainsPrefixVars(text, prefix string) []string { + vars := []string{} + gather := func(variable string) string { + if strings.HasPrefix(variable, prefix) { + vars = append(vars, strings.TrimPrefix(variable, prefix)) + } + return "" + } + os.Expand(text, gather) + + return vars +} + +// MaybeExpand conditionally replaces ${var} or $var in the string based on the mapping function. +// Only repalces the variable with the mapper string if it returns true. +// The mapper returning false will leave the original variable unchanged. +// Based on os.Expand +func MaybeExpand(s string, mapping func(string) (string, bool)) string { + var buf []byte + // ${} is all ASCII, so bytes are fine for this operation. + i := 0 + for j := 0; j < len(s); j++ { + if s[j] == '$' && j+1 < len(s) { + if buf == nil { + buf = make([]byte, 0, 2*len(s)) + } + buf = append(buf, s[i:j]...) + name, w := getShellName(s[j+1:]) + if name == "" { + // Unlike the original Expand, don't + // eat invalid syntax, just leave it + // and pass over. + w = 0 + buf = append(buf, s[j]) + } else { + replacement, shouldReplace := mapping(name) + if shouldReplace { + buf = append(buf, replacement...) + } else { + // We aren't replacing this + // reference, pass over it. + w = 0 + buf = append(buf, s[j]) + } + } + j += w + i = j + 1 + } + } + if buf == nil { + return s + } + return string(buf) + s[i:] +} + +// isShellSpecialVar reports whether the character identifies a special +// shell variable such as $*. +func isShellSpecialVar(c uint8) bool { + switch c { + case '*', '#', '$', '@', '!', '?', '-', '0', '1', '2', '3', '4', '5', '6', '7', '8', '9': + return true + } + return false +} + +// isAlphaNum reports whether the byte is an ASCII letter, number, or underscore. +func isAlphaNum(c uint8) bool { + return c == '_' || '0' <= c && c <= '9' || 'a' <= c && c <= 'z' || 'A' <= c && c <= 'Z' +} + +// getShellName returns the name that begins the string and the number of bytes +// consumed to extract it. If the name is enclosed in {}, it's part of a ${} +// expansion and two more bytes are needed than the length of the name. +func getShellName(s string) (string, int) { + switch { + case s[0] == '{': + if len(s) > 2 && isShellSpecialVar(s[1]) && s[2] == '}' { + return s[1:2], 3 + } + // Scan to closing brace + for i := 1; i < len(s); i++ { + if s[i] == '}' { + if i == 1 { + return "", 2 // Bad syntax; eat "${}" + } + return s[1:i], i + 1 + } + } + return "", 1 // Bad syntax; eat "${" + case isShellSpecialVar(s[0]): + return s[0:1], 1 + } + // Scan alphanumerics. + var i int + // nolint:revive + for i = 0; i < len(s) && isAlphaNum(s[i]); i++ { + } + return s[:i], i +} diff --git a/server/fleet/fleet_vars_test.go b/server/fleet/fleet_vars_test.go new file mode 100644 index 0000000000..0876799ec2 --- /dev/null +++ b/server/fleet/fleet_vars_test.go @@ -0,0 +1,51 @@ +package fleet + +import ( + "strings" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestContainsPrefixVars(t *testing.T) { + script := ` +#!/bin/sh + +echo $FLEET_SECRET_FOO is the secret +echo words${FLEET_SECRET_BAR}words +$FLEET_SECRET_BAZ +${FLEET_SECRET_QUX} +` + secrets := ContainsPrefixVars(script, FLEET_SECRET_PREFIX) + require.Contains(t, secrets, "FOO") + require.Contains(t, secrets, "BAR") + require.Contains(t, secrets, "BAZ") + require.Contains(t, secrets, "QUX") +} + +func TestMaybeExpand(t *testing.T) { + script := ` +This is $OTHER_VAR, $ $$ $* ${} in a sentence with${ALSO_OTHER_VAR}in the middle. +We want to remember $FLEET_SECRET_BANANA and also${FLEET_SECRET_STRAWBERRY}are important. +` + expected := ` +This is $OTHER_VAR, $ $$ $* ${} in a sentence with${ALSO_OTHER_VAR}in the middle. +We want to remember BREAD and alsoSHORTCAKEare important. +` + + mapping := map[string]string{ + "BANANA": "BREAD", + "STRAWBERRY": "SHORTCAKE", + } + + mapper := func(s string) (string, bool) { + if strings.HasPrefix(s, FLEET_SECRET_PREFIX) { + return mapping[strings.TrimPrefix(s, FLEET_SECRET_PREFIX)], true + } + return "", false + } + + expanded := MaybeExpand(script, mapper) + + require.Equal(t, expected, expanded) +} diff --git a/server/fleet/secrets.go b/server/fleet/secrets.go new file mode 100644 index 0000000000..8f5b5d5c0b --- /dev/null +++ b/server/fleet/secrets.go @@ -0,0 +1,3 @@ +package fleet + +const FLEET_SECRET_PREFIX = "FLEET_SECRET_"