Merge pull request #18 from saschagrunert/command-filter

Add command.Filter API

command/command.go

@@ -21,6 +21,7 @@ import (
 	"io"
 	"os"
 	"os/exec"
+	"regexp"
 	"strings"
 	"sync"
 	"syscall"
@@ -35,6 +36,7 @@ type Command struct {
 	stdErrWriters, stdOutWriters []io.Writer
 	env                          []string
 	verbose                      bool
+	filter                       *filter
 }
 
 // The internal command representation
@@ -43,6 +45,12 @@ type command struct {
 	pipeWriter *io.PipeWriter
 }
 
+// filter is the internally used struct for filtering command output.
+type filter struct {
+	regex      *regexp.Regexp
+	replaceAll string
+}
+
 // A generic command exit status
 type Status struct {
 	waitStatus syscall.WaitStatus
@@ -123,6 +131,7 @@ func (c *Command) isVerbose() bool {
 func (c *Command) Add(cmd string, args ...string) Commands {
 	addCmd := NewWithWorkDir(c.cmds[0].Dir, cmd, args...)
 	addCmd.verbose = c.verbose
+	addCmd.filter = c.filter
 	return Commands{c, addCmd}
 }
 
@@ -149,6 +158,20 @@ func (c *Command) AddOutputWriter(writer io.Writer) *Command {
 	return c
 }
 
+// Filter adds an output filter regular expression to the command. Every output
+// will then be replaced with the string provided by replaceAll.
+func (c *Command) Filter(regex, replaceAll string) (*Command, error) {
+	filterRegex, err := regexp.Compile(regex)
+	if err != nil {
+		return nil, errors.Wrap(err, "compile regular expression")
+	}
+	c.filter = &filter{
+		regex:      filterRegex,
+		replaceAll: replaceAll,
+	}
+	return c, nil
+}
+
 // Run starts the command and waits for it to finish. It returns an error if
 // the command execution was not possible at all, otherwise the Status.
 // This method prints the commands output during execution
@@ -266,12 +289,31 @@ func (c *Command) run(printOutput bool) (res *Status, err error) {
 				wg := sync.WaitGroup{}
 
 				wg.Add(2)
+
+				filterCopy := func(read io.ReadCloser, write io.Writer) (err error) {
+					if c.filter != nil {
+						builder := &strings.Builder{}
+						_, err = io.Copy(builder, read)
+						if err != nil {
+							return err
+						}
+						str := c.filter.regex.ReplaceAllString(
+							builder.String(), c.filter.replaceAll,
+						)
+						_, err = io.Copy(write, strings.NewReader(str))
+					} else {
+						_, err = io.Copy(write, read)
+					}
+					return err
+				}
+
 				go func() {
-					_, stdoutErr = io.Copy(stdOutWriter, stdout)
+					stdoutErr = filterCopy(stdout, stdOutWriter)
 					wg.Done()
 				}()
+
 				go func() {
-					_, stderrErr = io.Copy(stdErrWriter, stderr)
+					stderrErr = filterCopy(stderr, stdErrWriter)
 					wg.Done()
 				}()
 
@@ -390,6 +432,7 @@ func Available(commands ...string) (ok bool) {
 func (c Commands) Add(cmd string, args ...string) Commands {
 	addCmd := NewWithWorkDir(c[0].cmds[0].Dir, cmd, args...)
 	addCmd.verbose = c[0].verbose
+	addCmd.filter = c[0].filter
 	return append(c, addCmd)
 }
 

command/command_test.go

@@ -335,3 +335,23 @@ func TestEnv(t *testing.T) {
 	require.Nil(t, err)
 	require.Equal(t, "123\nbar\ntest", res.OutputTrimNL())
 }
+
+func TestFilterStdout(t *testing.T) {
+	cmd, err := New("echo", "-n", "1 2 2 3").Filter("[25]", "0")
+	require.Nil(t, err)
+
+	res, err := cmd.Add("echo", "-n", "4 5 6 2 2").Run()
+	require.Nil(t, err)
+	require.True(t, res.Success())
+	require.Zero(t, res.ExitCode())
+	require.Equal(t, "\n1 0 0 3\n4 0 6 0 0", res.Output())
+}
+
+func TestFilterStderr(t *testing.T) {
+	res, err := New("bash", "-c", ">&2 echo -n my secret").Filter("secret", "***")
+	require.Nil(t, err)
+	out, err := res.RunSilentSuccessOutput()
+	require.Nil(t, err)
+	require.Equal(t, "my ***", out.Error())
+	require.Empty(t, out.Output())
+}