Implementation of the capture_local() Starlark function

This patch implements the Go code for starlark builtin function capture_local(). This
function allows Crashd script to execute a commands on the local machine and  capture
the result in a local file.

This patch does the followings:
- Adds Go function to support starlark builtin func for capture_local
- Adds and updates tests for capture_local

Signed-off-by: Vladimir Vivien <vivienv@vmware.com>

starlark/capture.go

@@ -218,10 +218,8 @@ func captureOutput(source io.Reader, filePath, desc string) error {
 		return fmt.Errorf("source reader is nill")
 	}
 
-	logrus.Debugf("%s: capturing command output: %s", identifiers.capture, filePath)
 	file, err := os.Create(filePath)
 	if err != nil {
-		logrus.Errorf("%s output failed to create file: %s", identifiers.capture, err)
 		return err
 	}
 	defer file.Close()
@@ -233,11 +231,8 @@ func captureOutput(source io.Reader, filePath, desc string) error {
 	}
 
 	if _, err := io.Copy(file, source); err != nil {
-		logrus.Errorf("%s output failed to write file: %s", identifiers.capture, err)
 		return err
 	}
 
-	logrus.Debugf("%s output saved in %s", identifiers.capture, filePath)
-
 	return nil
 }

starlark/capture_local.go

@@ -0,0 +1,56 @@
+// Copyright (c) 2020 VMware, Inc. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+package starlark
+
+import (
+	"fmt"
+	"os"
+	"path/filepath"
+
+	"github.com/vladimirvivien/echo"
+	"go.starlark.net/starlark"
+)
+
+// captureLocalFunc is a built-in starlark function that runs a provided command on the local machine.
+// The output of the command is stored in a file at a specified location under the workdir directory.
+// Starlark format: run_local(cmd=<command> [,workdir=path][,file_name=name][,desc=description])
+func captureLocalFunc(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
+	var cmdStr, workdir, fileName, desc string
+	if err := starlark.UnpackArgs(
+		identifiers.captureLocal, args, kwargs,
+		"cmd", &cmdStr,
+		"workdir?", &workdir,
+		"file_name?", &fileName,
+		"desc?", &desc,
+	); err != nil {
+		return starlark.None, err
+	}
+
+	if len(workdir) == 0 {
+		dir, err := getWorkdirFromThread(thread)
+		if err != nil {
+			return starlark.None, err
+		}
+		workdir = dir
+	}
+	if len(fileName) == 0 {
+		fileName = fmt.Sprintf("%s.txt", sanitizeStr(cmdStr))
+	}
+
+	filePath := filepath.Join(workdir, fileName)
+	if err := os.MkdirAll(workdir, 0744); err != nil && !os.IsExist(err) {
+		return starlark.None, fmt.Errorf("%s: %s", identifiers.captureLocal, err)
+	}
+
+	p := echo.New().RunProc(cmdStr)
+	if p.Err() != nil {
+		return starlark.None, fmt.Errorf("%s: %s", identifiers.captureLocal, p.Err())
+	}
+
+	if err := captureOutput(p.Out(), filePath, desc); err != nil {
+		return starlark.None, fmt.Errorf("%s: %s", identifiers.captureLocal, err)
+	}
+
+	return starlark.String(filePath), nil
+}

starlark/capture_local_test.go

@@ -0,0 +1,229 @@
+// Copyright (c) 2020 VMware, Inc. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+package starlark
+
+import (
+	"bytes"
+	"fmt"
+	"io"
+	"os"
+	"path/filepath"
+	"strings"
+	"testing"
+
+	"go.starlark.net/starlark"
+)
+
+func TestCaptureLocalFunc(t *testing.T) {
+	tests := []struct {
+		name string
+		args func(t *testing.T) []starlark.Tuple
+		eval func(t *testing.T, kwargs []starlark.Tuple)
+	}{
+		{
+			name: "capture with defaults",
+			args: func(t *testing.T) []starlark.Tuple {
+				return []starlark.Tuple{{starlark.String("cmd"), starlark.String("echo 'Hello World!'")}}
+			},
+			eval: func(t *testing.T, kwargs []starlark.Tuple) {
+				val, err := captureLocalFunc(newTestThreadLocal(t), nil, nil, kwargs)
+				if err != nil {
+					t.Fatal(err)
+				}
+				expected := filepath.Join(defaults.workdir, fmt.Sprintf("%s.txt", sanitizeStr("echo 'Hello World!'")))
+				result := ""
+				if r, ok := val.(starlark.String); ok {
+					result = string(r)
+				}
+				defer func() {
+					os.RemoveAll(result)
+					os.RemoveAll(defaults.workdir)
+				}()
+
+				if result != expected {
+					t.Errorf("unexpected result: %s", result)
+				}
+
+				file, err := os.Open(result)
+				if err != nil {
+					t.Fatal(err)
+				}
+				buf := new(bytes.Buffer)
+				if _, err := io.Copy(buf, file); err != nil {
+					t.Fatal(err)
+				}
+				expected = strings.TrimSpace(buf.String())
+				if expected != "Hello World!" {
+					t.Errorf("unexpected content captured: %s", expected)
+				}
+				if err := file.Close(); err != nil {
+					t.Error(err)
+				}
+			},
+		},
+		{
+			name: "capture with args",
+			args: func(t *testing.T) []starlark.Tuple {
+				return []starlark.Tuple{
+					{starlark.String("cmd"), starlark.String("echo 'Hello World!'")},
+					{starlark.String("workdir"), starlark.String("/tmp/capturecrashd")},
+					{starlark.String("file_name"), starlark.String("echo.txt")},
+					{starlark.String("desc"), starlark.String("echo command")},
+				}
+			},
+			eval: func(t *testing.T, kwargs []starlark.Tuple) {
+				val, err := captureLocalFunc(newTestThreadLocal(t), nil, nil, kwargs)
+				if err != nil {
+					t.Fatal(err)
+				}
+				expected := filepath.Join("/tmp/capturecrashd", "echo.txt")
+				result := ""
+				if r, ok := val.(starlark.String); ok {
+					result = string(r)
+				}
+				defer func() {
+					os.RemoveAll(result)
+					os.RemoveAll(defaults.workdir)
+				}()
+
+				if result != expected {
+					t.Errorf("unexpected result: %s", result)
+				}
+
+				file, err := os.Open(result)
+				if err != nil {
+					t.Fatal(err)
+				}
+				buf := new(bytes.Buffer)
+				if _, err := io.Copy(buf, file); err != nil {
+					t.Fatal(err)
+				}
+				expected = strings.TrimSpace(buf.String())
+				if expected != "echo command\nHello World!" {
+					t.Errorf("unexpected content captured: %s", expected)
+				}
+				if err := file.Close(); err != nil {
+					t.Error(err)
+				}
+			},
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			test.eval(t, test.args(t))
+		})
+	}
+}
+
+func TestCaptureLocalScript(t *testing.T) {
+	tests := []struct {
+		name   string
+		script string
+		eval   func(t *testing.T, script string)
+	}{
+		{
+			name: "capture local defaults",
+			script: `
+result = capture_local("echo 'Hello World!'")
+`,
+			eval: func(t *testing.T, script string) {
+				exe := New()
+				if err := exe.Exec("test.star", strings.NewReader(script)); err != nil {
+					t.Fatal(err)
+				}
+
+				expected := filepath.Join(defaults.workdir, fmt.Sprintf("%s.txt", sanitizeStr("echo 'Hello World!'")))
+				var result string
+				resultVal := exe.result["result"]
+				if resultVal == nil {
+					t.Fatal("capture_local() should be assigned to a variable for test")
+				}
+				res, ok := resultVal.(starlark.String)
+				if !ok {
+					t.Fatal("capture_local() should return a string")
+				}
+				result = string(res)
+				defer func() {
+					os.RemoveAll(result)
+					os.RemoveAll(defaults.workdir)
+				}()
+
+				if result != expected {
+					t.Errorf("unexpected result: %s", result)
+				}
+
+				file, err := os.Open(result)
+				if err != nil {
+					t.Fatal(err)
+				}
+				buf := new(bytes.Buffer)
+				if _, err := io.Copy(buf, file); err != nil {
+					t.Fatal(err)
+				}
+				expected = strings.TrimSpace(buf.String())
+				if expected != "Hello World!" {
+					t.Errorf("unexpected content captured: %s", expected)
+				}
+				if err := file.Close(); err != nil {
+					t.Error(err)
+				}
+			},
+		},
+		{
+			name: "capture local with args",
+			script: `
+result = capture_local(cmd="echo 'Hello World!'", workdir="/tmp/capturecrash", file_name="echo_out.txt", desc="output command")
+`,
+			eval: func(t *testing.T, script string) {
+				exe := New()
+				if err := exe.Exec("test.star", strings.NewReader(script)); err != nil {
+					t.Fatal(err)
+				}
+
+				expected := filepath.Join("/tmp/capturecrash", "echo_out.txt")
+				var result string
+				resultVal := exe.result["result"]
+				if resultVal == nil {
+					t.Fatal("capture_local() should be assigned to a variable for test")
+				}
+				res, ok := resultVal.(starlark.String)
+				if !ok {
+					t.Fatal("capture_local() should return a string")
+				}
+				result = string(res)
+				defer func() {
+					os.RemoveAll(result)
+					os.RemoveAll(defaults.workdir)
+				}()
+
+				if result != expected {
+					t.Errorf("unexpected result: %s", result)
+				}
+
+				file, err := os.Open(result)
+				if err != nil {
+					t.Fatal(err)
+				}
+				buf := new(bytes.Buffer)
+				if _, err := io.Copy(buf, file); err != nil {
+					t.Fatal(err)
+				}
+				expected = strings.TrimSpace(buf.String())
+				if expected != "output command\nHello World!" {
+					t.Errorf("unexpected content captured: %s", expected)
+				}
+				if err := file.Close(); err != nil {
+					t.Error(err)
+				}
+			},
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			test.eval(t, test.script)
+		})
+	}
+}

starlark/starlark_exec.go

@@ -77,6 +77,7 @@ func newPredeclareds() starlark.StringDict {
 		identifiers.run:              starlark.NewBuiltin(identifiers.run, runFunc),
 		identifiers.runLocal:         starlark.NewBuiltin(identifiers.runLocal, runLocalFunc),
 		identifiers.capture:          starlark.NewBuiltin(identifiers.capture, captureFunc),
+		identifiers.captureLocal:     starlark.NewBuiltin(identifiers.capture, captureLocalFunc),
 		identifiers.copyFrom:         starlark.NewBuiltin(identifiers.copyFrom, copyFromFunc),
 		identifiers.kubeCfg:          starlark.NewBuiltin(identifiers.kubeCfg, kubeConfigFn),
 		identifiers.kubeCapture:      starlark.NewBuiltin(identifiers.kubeGet, KubeCaptureFn),

starlark/support.go

@@ -36,6 +36,7 @@ var (
 		run              string
 		runLocal         string
 		capture          string
+		captureLocal     string
 		copyFrom         string
 
 		kubeCapture string
@@ -58,6 +59,7 @@ var (
 		run:              "run",
 		runLocal:         "run_local",
 		capture:          "capture",
+		captureLocal:     "capture_local",
 		copyFrom:         "copy_from",
 
 		kubeCapture: "kube_capture",