From de2dc98dcbbcef10a6fd12feb4d03ac66c9ce6ab Mon Sep 17 00:00:00 2001 From: thinkmo Date: Mon, 21 Dec 2020 20:29:45 +0800 Subject: [PATCH] add parameter to capture_local --- docs/README.md | 3 ++- starlark/capture.go | 15 +++++++++---- starlark/capture_local.go | 6 ++++-- starlark/capture_local_test.go | 39 +++++++++++++++++++++++++++++----- 4 files changed, 51 insertions(+), 12 deletions(-) diff --git a/docs/README.md b/docs/README.md index 2c0ff347..a4462dce 100644 --- a/docs/README.md +++ b/docs/README.md @@ -480,6 +480,7 @@ This function runs a command locally on the machine running the script. It then | `workdir`|A parent directory where captured files will be saved|No, defaults to `crashd_config.workdir`| | `file_name`|The path/name of the generated file|No, auto-generated based on command string, if omitted| | `desc`|A short description added at the start of the file|No| +| `append` | boolean indicator to append to a file if it already exists or not |No, defaults to `False`| #### Output `capture_local()` returns the full path of the capured output file. @@ -767,4 +768,4 @@ bloop blah # this will result in a warning message with foo=bar as the only pair pairs to be passed to the .crsh file $ crash run diagnsotics.crsh --args-file /tmp/script.args WARN[0000] unknown entry in args file: blooop blah -``` \ No newline at end of file +``` diff --git a/starlark/capture.go b/starlark/capture.go index 132e84b1..7a526f1f 100644 --- a/starlark/capture.go +++ b/starlark/capture.go @@ -163,13 +163,13 @@ func execCaptureSSH(host, cmdStr, rootDir, fileName, desc string, agent ssh.Agen reader, err := ssh.RunRead(args, agent, cmdStr) if err != nil { logrus.Errorf("%s failed: %s", identifiers.capture, err) - if err := captureOutput(strings.NewReader(err.Error()), filePath, fmt.Sprintf("%s: failed", cmdStr)); err != nil { + if err := captureOutput(strings.NewReader(err.Error()), filePath, fmt.Sprintf("%s: failed", cmdStr), false); err != nil { logrus.Errorf("%s output failed: %s", identifiers.capture, err) return commandResult{resource: args.Host, result: filePath, err: err}, err } } - if err := captureOutput(reader, filePath, desc); err != nil { + if err := captureOutput(reader, filePath, desc, false); err != nil { logrus.Errorf("%s output failed: %s", identifiers.capture, err) return commandResult{resource: args.Host, result: filePath, err: err}, err } @@ -177,12 +177,19 @@ func execCaptureSSH(host, cmdStr, rootDir, fileName, desc string, agent ssh.Agen return commandResult{resource: args.Host, result: filePath, err: err}, nil } -func captureOutput(source io.Reader, filePath, desc string) error { +func captureOutput(source io.Reader, filePath, desc string, append bool) error { if source == nil { return fmt.Errorf("source reader is nill") } - file, err := os.Create(filePath) + flag := os.O_CREATE | os.O_WRONLY + if append { + flag |= os.O_APPEND + } else { + flag |= os.O_TRUNC + } + + file, err := os.OpenFile(filePath, flag, 0644) if err != nil { return err } diff --git a/starlark/capture_local.go b/starlark/capture_local.go index 71fcb3b0..39e97568 100644 --- a/starlark/capture_local.go +++ b/starlark/capture_local.go @@ -14,15 +14,17 @@ import ( // 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= [,workdir=path][,file_name=name][,desc=description]) +// Starlark format: run_local(cmd= [,workdir=path][,file_name=name][,desc=description][,append=append]) func captureLocalFunc(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { var cmdStr, workdir, fileName, desc string + var append bool if err := starlark.UnpackArgs( identifiers.captureLocal, args, kwargs, "cmd", &cmdStr, "workdir?", &workdir, "file_name?", &fileName, "desc?", &desc, + "append?", &append, ); err != nil { return starlark.None, fmt.Errorf("%s: %s", identifiers.captureLocal, err) } @@ -48,7 +50,7 @@ func captureLocalFunc(thread *starlark.Thread, b *starlark.Builtin, args starlar return starlark.None, fmt.Errorf("%s: %s", identifiers.captureLocal, p.Err()) } - if err := captureOutput(p.Out(), filePath, desc); err != nil { + if err := captureOutput(p.Out(), filePath, desc, append); err != nil { return starlark.None, fmt.Errorf("%s: %s", identifiers.captureLocal, err) } diff --git a/starlark/capture_local_test.go b/starlark/capture_local_test.go index 50174e49..62ee82d0 100644 --- a/starlark/capture_local_test.go +++ b/starlark/capture_local_test.go @@ -70,14 +70,29 @@ func TestCaptureLocalFunc(t *testing.T) { {starlark.String("workdir"), starlark.String("/tmp/capturecrashd")}, {starlark.String("file_name"), starlark.String("echo.txt")}, {starlark.String("desc"), starlark.String("echo command")}, + {starlark.String("append"), starlark.True}, } }, eval: func(t *testing.T, kwargs []starlark.Tuple) { + expected := filepath.Join("/tmp/capturecrashd", "echo.txt") + err := os.MkdirAll("/tmp/capturecrashd", 0777) + if err != nil { + t.Fatal(err) + } + f, err := os.OpenFile(expected, os.O_RDWR|os.O_CREATE, 0644) + if err != nil { + t.Fatal(err) + } + if _, err := f.Write([]byte("Hello World!\n")); err != nil { + t.Fatal(err) + } + if err := f.Close(); err != nil { + t.Fatal(err) + } 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) @@ -100,7 +115,7 @@ func TestCaptureLocalFunc(t *testing.T) { t.Fatal(err) } expected = strings.TrimSpace(buf.String()) - if expected != "echo command\nHello World!" { + if expected != "Hello World!\necho command\nHello World!" { t.Errorf("unexpected content captured: %s", expected) } if err := file.Close(); err != nil { @@ -174,15 +189,29 @@ result = capture_local("echo 'Hello World!'") { name: "capture local with args", script: ` -result = capture_local(cmd="echo 'Hello World!'", workdir="/tmp/capturecrash", file_name="echo_out.txt", desc="output command") +result = capture_local(cmd="echo 'Hello World!'", workdir="/tmp/capturecrash", file_name="echo_out.txt", desc="output command", append=True) `, eval: func(t *testing.T, script string) { + expected := filepath.Join("/tmp/capturecrash", "echo_out.txt") + err := os.MkdirAll("/tmp/capturecrash", 0777) + if err != nil { + t.Fatal(err) + } + f, err := os.OpenFile(expected, os.O_RDWR|os.O_CREATE, 0644) + if err != nil { + t.Fatal(err) + } + if _, err := f.Write([]byte("Hello World!\n")); err != nil { + t.Fatal(err) + } + if err := f.Close(); err != nil { + t.Fatal(err) + } 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 { @@ -211,7 +240,7 @@ result = capture_local(cmd="echo 'Hello World!'", workdir="/tmp/capturecrash", f t.Fatal(err) } expected = strings.TrimSpace(buf.String()) - if expected != "output command\nHello World!" { + if expected != "Hello World!\noutput command\nHello World!" { t.Errorf("unexpected content captured: %s", expected) } if err := file.Close(); err != nil {