Sigma-C2/operator/command_processor.go

package main

import (
	"crypto/tls"
	"fmt"
	"io"
	"slices"
	"strings"

	"github.com/chzyer/readline"
)

// Handles operator commands and context switching
func processOperatorCommand(command string, conn *tls.Conn, rl *readline.Instance, doneChan chan struct{}) bool {
	if command == "" {
		return false
	}

	// Handle exit and quit commands
	if command == "exit" {
		// Exit from agent context
		if currentAgentContext != "" {
			currentAgentContext = ""
			rl.SetPrompt("\033[31mSigma >\033[0m ")
			// Update tab completion after context change
			rl.Config.AutoComplete = getCompleter()
			return false
		} else {
			// Exit from application
			fmt.Println("Exiting...")
			conn.Write([]byte("exit"))
			close(doneChan) // signal graceful exit
			return true
		}
	} else if command == "quit" { // Exit from application and disregard context
		fmt.Println("Exiting...")
		conn.Write([]byte("exit"))
		close(doneChan) // signal graceful exit
		return true
	}

	// Receive help message
	if command == "help" {
		if currentAgentContext != "" {
			command = "context_help"
			conn.Write([]byte(command))
			return false

		} else {
			command = "general_help"
			conn.Write([]byte(command))
			return false
		}
	}

	// Handle interact command
	if strings.HasPrefix(command, "interact") {
		handleInteractCommand(command, rl)
		return false
	}

	// Handle command in agent context
	if currentAgentContext != "" && isControlCommand(command) {
		command = "CONTEXT:" + currentAgentContext + " " + command

		// Handle command in general context
	} else {
		// Double check maybe this is not agent context, but command has agent ID in the beginning
		cmd := strings.Fields(command)
		if slices.Contains(agents, cmd[0]) {
			command = "CONTEXT:" + command

			// General command
		} else {

			command = "GENERAL:" + command
		}
	}

	// Send command to server
	_, err := conn.Write([]byte(command))
	if err != nil {
		fmt.Printf("Error sending command: %v", err)
		return true
	}

	return false
}

// Executes the main command processing loop
func runCommandLoop(conn *tls.Conn, rl *readline.Instance, exitChan chan struct{}, doneChan chan struct{}) {
	for {
		select {
		case <-exitChan:
			fmt.Println("Exiting due to lost connection.")
			return
		default:
			line, err := rl.Readline()
			if err != nil {
				if err == readline.ErrInterrupt {
					if len(line) == 0 {
						conn.Write([]byte("exit"))
						fmt.Println("\nExiting...")
						close(doneChan) // signal graceful exit
						return
					}
					continue
				} else if err == io.EOF {
					fmt.Println("\nExiting...")
					return
				}
				fmt.Printf("Error reading line: %v", err)
				continue
			}

			command := strings.TrimSpace(line)

			// Print the command to terminal history before processing
			if command != "" {
				// Construct the prompt based on current context
				var currentPrompt string
				if currentAgentContext != "" {
					currentPrompt = fmt.Sprintf("\033[31mSigma [%s] >\033[0m ", currentAgentContext)
				} else {
					currentPrompt = "\033[31mSigma >\033[0m "
				}
				rl.Stdout().Write([]byte(currentPrompt + command + "\n"))
			}

			shouldExit := processOperatorCommand(command, conn, rl, doneChan)
			if shouldExit {
				return
			}
		}
	}
}

// Check if command is a task for agent
func isControlCommand(command string) bool {
	// Split the command into words
	parts := strings.Fields(command)
	if len(parts) == 0 {
		fmt.Println("Empty command :(")
		return false
	}

	// Check if the first word is a context command
	if slices.Contains(implantCommands, parts[0]) {
		return true
	} else {
		return false
	}
}

// Processes the interact command for agent context switching
func handleInteractCommand(command string, rl *readline.Instance) {
	cmd := strings.Fields(command)

	if len(cmd) != 2 {
		fmt.Println("Usage: interact <agent_id>")
		return
	}

	if !slices.Contains(agents, cmd[1]) {
		fmt.Println("Agent not found!")
		return
	}

	// Set new agent context
	currentAgentContext = cmd[1]

	// Update tab completion
	rl.Config.AutoComplete = getCompleter()

	// Update prompt with agent context
	newPrompt := fmt.Sprintf("\033[31mSigma [%s] >\033[0m ", currentAgentContext)
	rl.SetPrompt(newPrompt)
}