1package api
 2
 3import (
 4	"context"
 5	"fmt"
 6	"mysh/pkg/mythic"
 7	"time"
 8)
 9
10// TaskPollConfig configures task polling behavior
11type TaskPollConfig struct {
12	TimeoutSeconds int
13	PollInterval   time.Duration
14	ShowProgress   bool
15	RawOutput      bool
16}
17
18// DefaultTaskPollConfig returns sensible defaults for task polling
19func DefaultTaskPollConfig() TaskPollConfig {
20	return TaskPollConfig{
21		TimeoutSeconds: int(DefaultTaskTimeout.Seconds()),
22		PollInterval:   DefaultPollInterval,
23		ShowProgress:   true,
24		RawOutput:      false,
25	}
26}
27
28// PollTaskResponse polls for a task response until completion or timeout
29func PollTaskResponse(ctx context.Context, client MythicClient, taskID int, config TaskPollConfig) (*mythic.Task, error) {
30	for i := 0; i < config.TimeoutSeconds; i++ {
31		time.Sleep(config.PollInterval)
32
33		updatedTask, err := client.GetTaskResponse(ctx, taskID)
34		if err != nil {
35			return nil, fmt.Errorf("failed to get task response: %w", err)
36		}
37
38		if updatedTask.Status == TaskStatusCompleted || updatedTask.Status == TaskStatusError {
39			return updatedTask, nil
40		}
41
42		// Show progress dots (but not for raw output)
43		if config.ShowProgress && !config.RawOutput && i > 0 && i%ProgressDotInterval == 0 {
44			fmt.Printf(".")
45		}
46	}
47
48	return nil, fmt.Errorf("timeout waiting for response after %d seconds", config.TimeoutSeconds)
49}
50
51// ExecuteTaskAndWait creates a task and waits for its response
52func ExecuteTaskAndWait(ctx context.Context, client MythicClient, callbackID int, command, params string, config TaskPollConfig) (*mythic.Task, error) {
53	// Create the task
54	task, err := client.CreateTask(ctx, callbackID, command, params)
55	if err != nil {
56		return nil, fmt.Errorf("failed to create task: %w", err)
57	}
58
59	if config.ShowProgress && !config.RawOutput {
60		fmt.Printf("Task %d created, waiting for response...\n", task.DisplayID)
61	}
62
63	// Poll for response
64	return PollTaskResponse(ctx, client, task.ID, config)
65}