Commit cd842f9

bryfry <bryon@fryer.io>
2025-09-17 12:27:10
task ranges
1 parent 7cd3074
Changed files (2)
cmd/task.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"mysh/pkg/mythic"
 	"mysh/pkg/mythic/api"
+	"sort"
 	"strconv"
 	"strings"
 
@@ -18,10 +19,10 @@ var (
 )
 
 var taskCmd = &cobra.Command{
-	Use:     "task <callback_id> <command> [args...]",
+	Use:     "task <callback_id(s)> <command> [args...]",
 	Aliases: []string{"t", "exec"},
-	Short:   "Send a command to an agent and wait for response",
-	Long:    "Send a command to the specified agent callback and wait for the response. This is like a single-shot version of interact.",
+	Short:   "Send a command to one or more agents",
+	Long:    "Send a command to specified agent callback(s) and wait for response. Supports single IDs, comma-separated lists, and ranges (e.g., '1,3-5,10').",
 	Args:    cobra.MinimumNArgs(2),
 	RunE:    runTask,
 }
@@ -33,14 +34,82 @@ func init() {
 	taskCmd.Flags().BoolVar(&taskNoWait, "no-wait", false, "Create task but don't wait for response")
 }
 
+// parseCallbackIDs parses a callback ID string that can contain:
+// - Single ID: "1"
+// - Comma-separated list: "1,2,3"
+// - Ranges: "1-5" (expands to 1,2,3,4,5)
+// - Mixed: "1,3-5,10" (expands to 1,3,4,5,10)
+func parseCallbackIDs(input string) ([]int, error) {
+	var callbackIDs []int
+	seen := make(map[int]bool)
+
+	parts := strings.Split(input, ",")
+	for _, part := range parts {
+		part = strings.TrimSpace(part)
+		if part == "" {
+			continue
+		}
+
+		if strings.Contains(part, "-") {
+			// Handle range (e.g., "3-5")
+			rangeParts := strings.Split(part, "-")
+			if len(rangeParts) != 2 {
+				return nil, fmt.Errorf("invalid range format: %s", part)
+			}
+
+			start, err := strconv.Atoi(strings.TrimSpace(rangeParts[0]))
+			if err != nil {
+				return nil, fmt.Errorf("invalid start of range: %s", rangeParts[0])
+			}
+
+			end, err := strconv.Atoi(strings.TrimSpace(rangeParts[1]))
+			if err != nil {
+				return nil, fmt.Errorf("invalid end of range: %s", rangeParts[1])
+			}
+
+			if start > end {
+				return nil, fmt.Errorf("invalid range: start (%d) is greater than end (%d)", start, end)
+			}
+
+			// Add all IDs in range
+			for i := start; i <= end; i++ {
+				if !seen[i] {
+					callbackIDs = append(callbackIDs, i)
+					seen[i] = true
+				}
+			}
+		} else {
+			// Handle single ID
+			id, err := strconv.Atoi(part)
+			if err != nil {
+				return nil, fmt.Errorf("invalid callback ID: %s", part)
+			}
+
+			if !seen[id] {
+				callbackIDs = append(callbackIDs, id)
+				seen[id] = true
+			}
+		}
+	}
+
+	if len(callbackIDs) == 0 {
+		return nil, fmt.Errorf("no valid callback IDs found")
+	}
+
+	// Sort the IDs for consistent output
+	sort.Ints(callbackIDs)
+	return callbackIDs, nil
+}
+
 func runTask(cmd *cobra.Command, args []string) error {
 	if err := validateConfig(); err != nil {
 		return err
 	}
 
-	callbackID, err := strconv.Atoi(args[0])
+	// Parse callback IDs (supports lists and ranges)
+	callbackIDs, err := parseCallbackIDs(args[0])
 	if err != nil {
-		return fmt.Errorf("invalid callback ID: %s", args[0])
+		return fmt.Errorf("invalid callback ID(s): %w", err)
 	}
 
 	command := args[1]
@@ -52,62 +121,112 @@ func runTask(cmd *cobra.Command, args []string) error {
 	client := mythic.NewClient(mythicURL, token, insecure, socksProxy)
 	ctx := context.Background()
 
-	// Verify callback exists and is active
-	targetCallback, err := api.FindActiveCallback(ctx, client, callbackID)
-	if err != nil {
-		return err
+	// Track results for multiple callbacks
+	type TaskResult struct {
+		CallbackID int
+		Callback   *mythic.Callback
+		Task       *mythic.Task
+		Error      error
 	}
 
-	if taskNoWait {
-		// Create task but don't wait for response
-		task, err := client.CreateTask(ctx, targetCallback.ID, command, params)
-		if err != nil {
-			return fmt.Errorf("failed to create task: %w", err)
-		}
+	var results []TaskResult
 
-		if !taskRawOutput {
-			fmt.Printf("Task %d created on callback %d (%s@%s) - not waiting for response\n",
-				task.DisplayID, callbackID, targetCallback.User, targetCallback.Host)
-		}
-		return nil
-	}
+	// For multiple callbacks, automatically use no-wait mode
+	useNoWait := taskNoWait || len(callbackIDs) > 1
 
-	// For raw output, don't show progress messages
-	if !taskRawOutput {
-		fmt.Printf("Executing '%s %s' on callback %d (%s@%s)\n",
-			command, params, callbackID, targetCallback.User, targetCallback.Host)
-	}
+	// Process each callback
+	for _, callbackID := range callbackIDs {
+		result := TaskResult{CallbackID: callbackID}
 
-	// Configure polling
-	config := api.TaskPollConfig{
-		TimeoutSeconds: taskWaitTime,
-		PollInterval:   api.DefaultPollInterval,
-		ShowProgress:   !taskRawOutput,
-		RawOutput:      taskRawOutput,
-	}
-
-	// Execute task and wait for response
-	updatedTask, err := api.ExecuteTaskAndWait(ctx, client, targetCallback.ID, command, params, config)
-	if err != nil {
-		if taskRawOutput {
-			// For raw output on timeout, don't show error message
-			return nil
+		// Verify callback exists and is active
+		targetCallback, err := api.FindActiveCallback(ctx, client, callbackID)
+		if err != nil {
+			result.Error = err
+			results = append(results, result)
+			continue
+		}
+		result.Callback = targetCallback
+
+		if useNoWait {
+			// Create task but don't wait for response
+			task, err := client.CreateTask(ctx, targetCallback.ID, command, params)
+			if err != nil {
+				result.Error = fmt.Errorf("failed to create task: %w", err)
+			} else {
+				result.Task = task
+			}
+			results = append(results, result)
+		} else {
+			// Configure polling (only for single callback)
+			config := api.TaskPollConfig{
+				TimeoutSeconds: taskWaitTime,
+				PollInterval:   api.DefaultPollInterval,
+				ShowProgress:   !taskRawOutput,
+				RawOutput:      taskRawOutput,
+			}
+
+			// Execute task and wait for response
+			task, err := api.ExecuteTaskAndWait(ctx, client, targetCallback.ID, command, params, config)
+			if err != nil {
+				result.Error = err
+			} else {
+				result.Task = task
+			}
+			results = append(results, result)
 		}
-		fmt.Printf("\n%v\n", err)
-		return nil
 	}
 
+	// Display results
 	if taskRawOutput {
-		// Raw output: only print response
-		if updatedTask.Response != "" {
-			fmt.Print(updatedTask.Response)
+		// Raw output: only print responses
+		for _, result := range results {
+			if result.Task != nil && result.Task.Response != "" {
+				fmt.Print(result.Task.Response)
+			}
 		}
 	} else {
-		// Normal output: show status and response
-		fmt.Printf("\nTask Status: %s\n", updatedTask.Status)
-		if updatedTask.Response != "" {
-			fmt.Printf("Response:\n%s\n", updatedTask.Response)
+		// Normal output: show detailed results
+		if len(callbackIDs) == 1 {
+			// Single callback: use original format
+			result := results[0]
+			if result.Error != nil {
+				return result.Error
+			}
+
+			if useNoWait {
+				fmt.Printf("Task %d created on callback %d (%s@%s) - not waiting for response\n",
+					result.Task.DisplayID, result.CallbackID, result.Callback.User, result.Callback.Host)
+			} else {
+				fmt.Printf("Task Status: %s\n", result.Task.Status)
+				if result.Task.Response != "" {
+					fmt.Printf("Response:\n%s\n", result.Task.Response)
+				}
+			}
+		} else {
+			// Multiple callbacks: show results for each
+			fmt.Printf("Executing '%s %s' on %d callbacks:\n\n", command, params, len(callbackIDs))
+
+			for _, result := range results {
+				fmt.Printf("=== Callback %d", result.CallbackID)
+				if result.Callback != nil {
+					fmt.Printf(" (%s@%s)", result.Callback.User, result.Callback.Host)
+				}
+				fmt.Printf(" ===\n")
+
+				if result.Error != nil {
+					fmt.Printf("Error: %v\n", result.Error)
+				} else if useNoWait {
+					fmt.Printf("Task %d created - not waiting for response\n", result.Task.DisplayID)
+				} else {
+					fmt.Printf("Status: %s\n", result.Task.Status)
+					if result.Task.Response != "" {
+						fmt.Printf("Response:\n%s\n", result.Task.Response)
+					}
+				}
+				fmt.Println()
+			}
 		}
 	}
+
 	return nil
 }
\ No newline at end of file
TODO.md
@@ -2,8 +2,10 @@
 - ✅ Renamed project from go-mythic to mysh (pronounced mɪʃh)
 - ✅ task-list --all (same as --limit 0), and -a
 - ✅ task-list without callback lists all tasks and shows callback column
-- task with --no-wait, doesn't wait for output
-- task [callback] arg make it accept lists, comma seperated, or ranges e.g. 1,2,3-5,10
+- ✅ task with --no-wait, doesn't wait for output
+- ✅ task [callback] arg make it accept lists, comma seperated, or ranges e.g. 1,2,3-5,10
+- make cb and callback aliases for callbacks
+- update callbacks to show last-checkin as 0s or 0m00s ago
 - invert cobra-cli nto functions (no global vars)
 - update output table formats to respect terminal width
 - create a .cache/mysh/ for task result cached values