fix: misc

Author: h3zh1

client/command/common/completer.go

@@ -360,6 +360,25 @@ func SessionModuleCompleter(con *core.Console) carapace.Action {
 	return carapace.ActionCallback(callback)
 }
 
+func SessionBundleCompleter(con *core.Console) carapace.Action {
+	callback := func(c carapace.Context) carapace.Action {
+		results := make([]string, 0)
+		sess := con.GetInteractive()
+		if sess == nil || sess.Data == nil || sess.Data.BundleMap == nil {
+			return carapace.ActionValuesDescribed(results...).Tag("session bundles")
+		}
+		seen := make(map[string]bool)
+		for _, bundle := range sess.Data.BundleMap {
+			if bundle != "" && !seen[bundle] {
+				seen[bundle] = true
+				results = append(results, bundle, "")
+			}
+		}
+		return carapace.ActionValuesDescribed(results...).Tag("session bundles")
+	}
+	return carapace.ActionCallback(callback)
+}
+
 func ModulesCompleter() carapace.Action {
 	callback := func(c carapace.Context) carapace.Action {
 		results := make([]string, 0)

client/command/common/output.go

@@ -10,37 +10,47 @@ import (
 	"github.com/spf13/pflag"
 )
 
-// BindOutputFlags adds --file and --output flags to a command
+// BindOutputFlags adds -f/--file flag to a command for saving output to a file.
 func BindOutputFlags(cmd *cobra.Command) {
 	Bind("output", false, cmd, func(f *pflag.FlagSet) {
-		f.BoolP("file", "f", false, "save output to file")
-		f.StringP("output", "o", "", "output file path")
+		f.StringP("file", "f", "", "save output to file path")
 	})
 }
 
-// HandleTaskOutput waits for task completion and optionally saves to file
+// HandleTaskOutput optionally saves task output to file based on -f flag,
+// then always displays output to console.
 func HandleTaskOutput(cmd *cobra.Command, con *core.Console, task *clientpb.Task) error {
-	toFile, _ := cmd.Flags().GetBool("file")
-	outputPath, _ := cmd.Flags().GetString("output")
+	outputPath, _ := cmd.Flags().GetString("file")
 
-	if !toFile && outputPath == "" {
-		// 正常流程：异步处理
+	if outputPath == "" {
 		con.GetInteractive().Console(task, string(*con.App.Shell().Line()))
 		return nil
 	}
 
-	// 需要保存文件：同步等待结果
-	tasksContext, err := con.Rpc.GetAllTaskContent(con.GetInteractive().Context(), &clientpb.Task{
+	if task == nil {
+		return fmt.Errorf("task is nil")
+	}
+
+	session := con.GetInteractive()
+	if session == nil {
+		return fmt.Errorf("no active session")
+	}
+
+	taskReq := &clientpb.Task{
 		SessionId: task.SessionId,
 		TaskId:    task.TaskId,
 		Need:      -1,
-	})
-	if err != nil {
-		return fmt.Errorf("failed to get task content: %w", err)
 	}
 
-	if outputPath == "" {
-		outputPath = fmt.Sprintf("task_%d.txt", task.TaskId)
+	// Task content may arrive asynchronously. Wait for completion first so the
+	// server-side cache/disk state is ready before collecting all chunks.
+	if _, err := con.Rpc.WaitTaskFinish(session.Context(), taskReq); err != nil {
+		return fmt.Errorf("failed to wait task finish: %w", err)
+	}
+
+	tasksContext, err := con.Rpc.GetAllTaskContent(session.Context(), taskReq)
+	if err != nil {
+		return fmt.Errorf("failed to get task content: %w", err)
 	}
 
 	var rendered []byte
@@ -54,6 +64,9 @@ func HandleTaskOutput(cmd *cobra.Command, con *core.Console, task *clientpb.Task
 		if err != nil {
 			return fmt.Errorf("failed to render task output: %w", err)
 		}
+		if text == "" {
+			continue
+		}
 		rendered = append(rendered, []byte(text+"\n")...)
 	}
 
@@ -63,7 +76,13 @@ func HandleTaskOutput(cmd *cobra.Command, con *core.Console, task *clientpb.Task
 
 	con.Log.Infof("Task output saved to: %s\n", outputPath)
 
-	// 同时也显示到控制台
-	con.GetInteractive().Console(task, string(*con.App.Shell().Line()))
+	// File output already waited for completion, so a user-supplied --wait would
+	// otherwise trigger a second wait/render pass in PersistentPostRunE.
+	if waitFlag := cmd.Flags().Lookup("wait"); waitFlag != nil {
+		_ = cmd.Flags().Set("wait", "false")
+	}
+
+	// Also display to console
+	session.Console(task, string(*con.App.Shell().Line()))
 	return nil
 }

client/command/exec/commands.go

@@ -35,6 +35,7 @@ run gogo.exe -- -i 127.0.0.1 -p http
 		carapace.ActionValues().Usage("command to execute"),
 		carapace.ActionValues().Usage("arguments to the command"),
 	)
+	common.BindOutputFlags(runCmd)
 
 	executeCmd := &cobra.Command{
 		Use:   consts.ModuleAliasExecute + " [cmdline]",
@@ -60,6 +61,7 @@ execute gogo.exe -- -i 127.0.0.1 -p http
 	common.BindArgCompletions(executeCmd, nil,
 		carapace.ActionValues().Usage("command to execute"),
 	)
+	common.BindOutputFlags(executeCmd)
 
 	shellCmd := &cobra.Command{
 		Use:     consts.ModuleAliasShell + " [cmdline]",
@@ -83,6 +85,7 @@ execute gogo.exe -- -i 127.0.0.1 -p http
 	common.BindFlag(shellCmd, func(f *pflag.FlagSet) {
 		f.BoolP("quiet", "q", false, "disable output")
 	})
+	common.BindOutputFlags(shellCmd)
 
 	execLocalCmd := &cobra.Command{
 		Use:   consts.ModuleExecuteLocal + " [local_exe]",
@@ -108,6 +111,7 @@ execute_local local_exe --ppid 1234 --block_dll --etw --argue "argue"
 		f.StringP("process", "n", "", "custom process path")
 		f.BoolP("output", "o", false, "disable output")
 	})
+	common.BindOutputFlags(execLocalCmd)
 
 	inlineLocalCmd := &cobra.Command{
 		Use:   consts.ModuleInlineLocal + " [local_exe]",
@@ -133,6 +137,7 @@ inline_local whoami
 		f.StringP("process", "n", "", "custom process path")
 		f.BoolP("output", "o", false, "disable output")
 	})
+	common.BindOutputFlags(inlineLocalCmd)
 
 	powershellCmd := &cobra.Command{
 		Use:   consts.ModuleAliasPowershell + " [cmdline]",
@@ -160,6 +165,7 @@ powershell dir
 	common.BindFlag(powershellCmd, func(f *pflag.FlagSet) {
 		f.BoolP("quiet", "q", false, "disable output")
 	})
+	common.BindOutputFlags(powershellCmd)
 
 	execAssemblyCmd := &cobra.Command{
 		Use:   consts.ModuleExecuteAssembly + " [file]",
@@ -185,6 +191,7 @@ execute_assembly potato.exe "whoami"
 		carapace.ActionValues().Usage("arguments to pass to the assembly args"))
 
 	common.BindFlag(execAssemblyCmd, common.SacrificeFlagSet, common.CLRFlagSet)
+	common.BindOutputFlags(execAssemblyCmd)
 
 	inlineAssemblyCmd := &cobra.Command{
 		Use:   consts.ModuleInlineAssembly + " [file]",
@@ -216,6 +223,7 @@ inline_assembly --amsi potato.exe -- cmd /c whoami
 		carapace.ActionValues().Usage("arguments to pass to the assembly args"))
 
 	common.BindFlag(inlineAssemblyCmd, common.CLRFlagSet)
+	common.BindOutputFlags(inlineAssemblyCmd)
 
 	execShellcodeCmd := &cobra.Command{
 		Use:   consts.ModuleExecuteShellcode + " [shellcode_file]",
@@ -242,6 +250,7 @@ execute_shellcode example.bin
 		carapace.ActionValues().Usage("arguments to pass to the assembly entrypoint"))
 
 	common.BindFlag(execShellcodeCmd, common.ExecuteFlagSet, common.SacrificeFlagSet)
+	common.BindOutputFlags(execShellcodeCmd)
 
 	inlineShellcodeCmd := &cobra.Command{
 		Use:   consts.ModuleAliasInlineShellcode + " [shellcode_file]",
@@ -269,6 +278,7 @@ inline_shellcode example.bin
 	common.BindArgCompletions(inlineShellcodeCmd, nil,
 		carapace.ActionFiles().Usage("path the shellcode file"))
 	common.BindFlag(inlineShellcodeCmd, common.ExecuteFlagSet)
+	common.BindOutputFlags(inlineShellcodeCmd)
 
 	execDLLCmd := &cobra.Command{
 		Use:   consts.ModuleExecuteDll + " [dll]",
@@ -304,6 +314,7 @@ execute_dll example.dll -e entrypoint -- arg1 arg2
 		f.StringP("entrypoint", "e", "", "custom entrypoint")
 		f.StringP("binPath", "", "", "custom process path")
 	})
+	common.BindOutputFlags(execDLLCmd)
 
 	common.BindFlagCompletions(execDLLCmd, func(comp carapace.ActionMap) {
 		comp["binPath"] = carapace.ActionFiles()
@@ -334,6 +345,7 @@ dllspawn example.dll
 		f.StringP("entrypoint", "e", "", "custom entrypoint")
 		f.StringP("binPath", "", "", "custom process path")
 	})
+	common.BindOutputFlags(execDLLSpawnCmd)
 
 	common.BindFlagCompletions(execDLLSpawnCmd, func(comp carapace.ActionMap) {
 		comp["binPath"] = carapace.ActionFiles()
@@ -371,6 +383,7 @@ inline_dll example.dll -e RunFunction -- arg1 arg2
 	common.BindFlag(inlineDLLCmd, common.ExecuteFlagSet, func(f *pflag.FlagSet) {
 		f.StringP("entrypoint", "e", "", "entrypoint")
 	})
+	common.BindOutputFlags(inlineDLLCmd)
 
 	execExeCmd := &cobra.Command{
 		Use:   consts.ModuleExecuteExe + " [exe]",
@@ -395,6 +408,7 @@ execute_exe gogo.exe -- -i 123.123.123.123 -p top2
 		carapace.ActionValues().Usage("arguments to pass to the assembly entrypoint"))
 
 	common.BindFlag(execExeCmd, common.ExecuteFlagSet, common.SacrificeFlagSet)
+	common.BindOutputFlags(execExeCmd)
 
 	inlinePECmd := &cobra.Command{
 		Use:   consts.ModuleAliasInlineExe + " [exe]",
@@ -421,6 +435,7 @@ inline_exe hackbrowserdata.exe -- -h
 `,
 	}
 	common.BindFlag(inlinePECmd, common.ExecuteFlagSet)
+	common.BindOutputFlags(inlinePECmd)
 	common.BindArgCompletions(inlinePECmd, nil,
 		carapace.ActionFiles().Usage("path the PE file"))
 
@@ -456,6 +471,7 @@ bof dir.x64.o -- wstr:"C:\\Windows\\System32"
 	common.BindArgCompletions(execBofCmd, nil,
 		carapace.ActionFiles().Usage("path the BOF file"),
 		carapace.ActionValues().Usage("arguments to pass to the assembly entrypoint"))
+	common.BindOutputFlags(execBofCmd)
 
 	powerpickCmd := &cobra.Command{
 		Use:   consts.ModulePowerpick + " [args]",
@@ -483,6 +499,7 @@ powerpick -s powerview.ps1 -- Get-NetUser
 	common.BindFlagCompletions(powerpickCmd, func(comp carapace.ActionMap) {
 		comp["script"] = carapace.ActionFiles()
 	})
+	common.BindOutputFlags(powerpickCmd)
 
 	return []*cobra.Command{
 		runCmd,

client/command/exec/dllspwan.go

@@ -29,8 +29,7 @@ func ExecuteDLLSpawnCmd(cmd *cobra.Command, con *core.Console) error {
 	if err != nil {
 		return err
 	}
-	session.Console(task, string(*con.App.Shell().Line()))
-	return nil
+	return common.HandleTaskOutput(cmd, con, task)
 }
 
 func ExecuteDLLSpawn(rpc clientrpc.MaliceRPCClient, sess *client.Session, dllPath string, entrypoint string, data string, binPath string, out bool, timeout uint32, arch string, process string, sac *implantpb.SacrificeProcess) (*clientpb.Task, error) {

client/command/exec/exec_test.go

@@ -1,9 +1,16 @@
 package exec_test
 
 import (
+	"context"
+	"errors"
+	"fmt"
+	"os"
+	"path/filepath"
+	"strings"
 	"testing"
 
 	"github.com/chainreactors/IoM-go/consts"
+	"github.com/chainreactors/IoM-go/proto/client/clientpb"
 	implantpb "github.com/chainreactors/IoM-go/proto/implant/implantpb"
 	"github.com/chainreactors/malice-network/client/command/testsupport"
 	"google.golang.org/grpc/metadata"
@@ -105,6 +112,87 @@ func TestExecCommandConformance(t *testing.T) {
 	})
 }
 
+func TestShellFileOutputWaitsForAsyncTask(t *testing.T) {
+	h := testsupport.NewHarness(t)
+	outputPath := filepath.Join(t.TempDir(), "shell-output.txt")
+
+	waitCalled := false
+	h.Recorder.OnTaskContext("WaitTaskFinish", func(_ context.Context, request any) (*clientpb.TaskContext, error) {
+		waitCalled = true
+		task, ok := request.(*clientpb.Task)
+		if !ok {
+			return nil, fmt.Errorf("wait request type = %T, want *clientpb.Task", request)
+		}
+		return &clientpb.TaskContext{
+			Task: &clientpb.Task{
+				TaskId:    task.GetTaskId(),
+				SessionId: task.GetSessionId(),
+				Type:      consts.ModuleExecute,
+				Cur:       2,
+				Total:     2,
+				Finished:  true,
+			},
+			Session: h.Session.Session,
+			Spite: &implantpb.Spite{
+				Error: 6,
+			},
+		}, nil
+	})
+	h.Recorder.OnTaskContexts("GetAllTaskContent", func(_ context.Context, request any) (*clientpb.TaskContexts, error) {
+		if !waitCalled {
+			return nil, errors.New("GetAllTaskContent called before WaitTaskFinish")
+		}
+		task, ok := request.(*clientpb.Task)
+		if !ok {
+			return nil, fmt.Errorf("content request type = %T, want *clientpb.Task", request)
+		}
+		return &clientpb.TaskContexts{
+			Task: &clientpb.Task{
+				TaskId:    task.GetTaskId(),
+				SessionId: task.GetSessionId(),
+				Type:      consts.ModuleExecute,
+				Cur:       2,
+				Total:     2,
+				Finished:  true,
+			},
+			Session: h.Session.Session,
+			Spites: []*implantpb.Spite{
+				{Error: 0},
+				{Error: 6},
+			},
+		}, nil
+	})
+
+	if err := h.Execute(consts.ModuleAliasShell, "-f", outputPath, "whoami"); err != nil {
+		t.Fatalf("execute shell with file output failed: %v", err)
+	}
+
+	calls := h.Recorder.Calls()
+	if len(calls) != 3 {
+		t.Fatalf("primary call count = %d, want 3 (Execute, WaitTaskFinish, GetAllTaskContent)", len(calls))
+	}
+	if calls[0].Method != "Execute" || calls[1].Method != "WaitTaskFinish" || calls[2].Method != "GetAllTaskContent" {
+		t.Fatalf("primary call order = %#v, want [Execute WaitTaskFinish GetAllTaskContent]", []string{calls[0].Method, calls[1].Method, calls[2].Method})
+	}
+
+	taskReq, ok := calls[1].Request.(*clientpb.Task)
+	if !ok {
+		t.Fatalf("wait request type = %T, want *clientpb.Task", calls[1].Request)
+	}
+
+	data, err := os.ReadFile(outputPath)
+	if err != nil {
+		t.Fatalf("read output file failed: %v", err)
+	}
+	text := string(data)
+	if !strings.Contains(text, fmt.Sprintf("task: %d true", taskReq.GetTaskId())) {
+		t.Fatalf("output file = %q, want first aggregated chunk", text)
+	}
+	if !strings.Contains(text, fmt.Sprintf("task: %d false", taskReq.GetTaskId())) {
+		t.Fatalf("output file = %q, want second aggregated chunk", text)
+	}
+}
+
 func assertExecTaskEvent(t testing.TB, h *testsupport.Harness, md metadata.MD, wantType string) {
 	t.Helper()
 

client/command/exec/execute-assembly.go

@@ -21,8 +21,7 @@ func ExecuteAssemblyCmd(cmd *cobra.Command, con *core.Console) error {
 	if err != nil {
 		return err
 	}
-	session.Console(task, string(*con.App.Shell().Line()))
-	return nil
+	return common.HandleTaskOutput(cmd, con, task)
 }
 
 func ExecuteAssembly(rpc clientrpc.MaliceRPCClient, sess *client.Session, path string, args []string, out bool, param map[string]string, sac *implantpb.SacrificeProcess) (*clientpb.Task, error) {
@@ -46,8 +45,7 @@ func InlineAssemblyCmd(cmd *cobra.Command, con *core.Console) error {
 	if err != nil {
 		return err
 	}
-	session.Console(task, string(*con.App.Shell().Line()))
-	return nil
+	return common.HandleTaskOutput(cmd, con, task)
 }
 
 func InlineAssembly(rpc clientrpc.MaliceRPCClient, sess *client.Session, path string, args []string, out bool, param map[string]string) (*clientpb.Task, error) {

client/command/exec/execute-bof.go

@@ -21,8 +21,7 @@ func ExecuteBofCmd(cmd *cobra.Command, con *core.Console) error {
 	if err != nil {
 		return err
 	}
-	con.GetInteractive().Console(task, string(*con.App.Shell().Line()))
-	return nil
+	return common.HandleTaskOutput(cmd, con, task)
 }
 
 func ExecBof(rpc clientrpc.MaliceRPCClient, sess *client.Session, bofPath string, args []string, out bool) (*clientpb.Task, error) {