Sigma-C2/operator/operator.go

package main

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

	"github.com/chzyer/readline"
)

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

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

// Available commands fpr tab completion
var availableCommands = []struct {
	command     string
	subcommands map[string][]string // Changed to a map for better nesting
}{
	{"agents", nil},
	{"show agent", map[string][]string{
		"": {""}, // To be dynamically updated
	}},
	{"show tasks", map[string][]string{
		"": {""}, // To be dynamically updated
	}},
	{"clear tasks", map[string][]string{
		"": {""},
	}},
	// Tasks
	{"kill", map[string][]string{
		"": {""},
	}},
	{"cleanup", map[string][]string{
		"": {""},
	}},
	{"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{
		"": {""},
	}},
	{"keylogger", map[string][]string{
		"start": {""},
		"stop":  {""},
	}},
	{"persistence", map[string][]string{
		"add":    {""},
		"remove": {""},
	}},
	{"download", map[string][]string{
		"": {""},
	}},

	// Listeners
	{"listen", map[string][]string{
		"-t":          {"ssl", "tcp", "dns"},
		"--transport": {"ssl", "tcp", "dns"},
	}},
	{"show listeners", nil},
	{"stop listener", map[string][]string{
		"": {""},
	}},

	{"generate", map[string][]string{
		"agent":  {},
		"beacon": {},
	}},

	{"show modules", nil},
	{"exit", nil},
}

// Update available commands for "generate agent" and "generate beacon" and for "stop listeners"
func UpdateListenersSubCommands(listeners []string) {
	for i := range availableCommands {
		if availableCommands[i].command == "generate" {
			availableCommands[i].subcommands["agent"] = listeners
			availableCommands[i].subcommands["beacon"] = listeners
		} else if availableCommands[i].command == "stop listener" {
			availableCommands[i].subcommands[""] = listeners
		}
	}
}

// Update available commands
func UpdateAgentsSubCommands(agent []string) {
	for i := range availableCommands {
		if availableCommands[i].command == "show agent" {
			availableCommands[i].subcommands[""] = agent
		} else if availableCommands[i].command == "show tasks" {
			availableCommands[i].subcommands[""] = agent
		} else if availableCommands[i].command == "clear tasks" {
			availableCommands[i].subcommands[""] = agent
		} else if availableCommands[i].command == "kill" {
			availableCommands[i].subcommands[""] = agent
		} else if availableCommands[i].command == "cleanup" {
			availableCommands[i].subcommands[""] = agent
		} else if availableCommands[i].command == "sleep" {
			availableCommands[i].subcommands[""] = agent
		} else if availableCommands[i].command == "shell" {
			availableCommands[i].subcommands[""] = agent
		} else if availableCommands[i].command == "ps" {
			availableCommands[i].subcommands[""] = agent
		} else if availableCommands[i].command == "run" {
			availableCommands[i].subcommands[""] = agent
		} else if availableCommands[i].command == "artifacts" {
			availableCommands[i].subcommands[""] = agent
		} else if availableCommands[i].command == "files" {
			availableCommands[i].subcommands[""] = agent
		} else if availableCommands[i].command == "sysinfo" {
			availableCommands[i].subcommands[""] = agent
		} else if availableCommands[i].command == "download" {
			availableCommands[i].subcommands[""] = agent
		} else if availableCommands[i].command == "keylogger" {
			availableCommands[i].subcommands["start"] = agent
			availableCommands[i].subcommands["stop"] = agent
		} else if availableCommands[i].command == "persistence" {
			availableCommands[i].subcommands["add"] = agent
			availableCommands[i].subcommands["remove"] = agent
		}
	}
}

func getCompleter() *readline.PrefixCompleter {
	var items []readline.PrefixCompleterInterface

	for _, cmd := range availableCommands {
		if cmd.subcommands != nil {
			var subItems []readline.PrefixCompleterInterface
			for sub, subsub := range cmd.subcommands {
				if sub == "" { // Top-level subcommands
					for _, s := range subsub {
						subItems = append(subItems, readline.PcItem(s))
					}
				} else { // Nested subcommands
					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))
		}
	}

	return readline.NewPrefixCompleter(items...)
}

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")

	// Load CA certificate
	caCert, err := os.ReadFile(caCertPath)
	if err != nil {
		log.Fatalf("Failed to read CA certificate: %v", err)
	}
	caCertPool := x509.NewCertPool()
	if !caCertPool.AppendCertsFromPEM(caCert) {
		log.Fatal("Failed to append CA certificate to pool")
	}

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

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

	// Connect to the server
	conn, err := tls.Dial("tcp", *serverAddr, tlsConfig)
	if err != nil {
		fmt.Println("Error connecting to server:", err)
		return
	}
	defer conn.Close()
	fmt.Println("Connected to server")

	// Authentication
	fmt.Println("Starting authentication")
	authMessage := fmt.Sprintf("AUTH:%s:%s", *operatorID, *operatorPassword)
	_, err = conn.Write([]byte(authMessage))
	if err != nil {
		fmt.Println("Error sending credentials:", err)
		return
	}

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

	// Parse response message
	if response == "AUTH_OK" {
		fmt.Println("Auth successful. You can now issue commands.")
	} else {
		fmt.Println("Auth failed.")
		return
	}

	// Create readline config with more explicit settings
	rlConfig := &readline.Config{
		Prompt:          "\033[31mSigma >\033[0m ", // Colored prompt
		HistoryFile:     "/tmp/operator_history.txt",
		AutoComplete:    getCompleter(),
		InterruptPrompt: "^C",
		EOFPrompt:       "exit",

		// Enable tab completion settings
		DisableAutoSaveHistory: false,
		HistorySearchFold:      true,

		// Add filters for better handling of input
		UniqueEditLine: true,
	}

	rl, err := readline.NewEx(rlConfig)
	if err != nil {
		fmt.Printf("Error initializing readline: %v\n", err)
		return
	}
	defer rl.Close()

	// Set up better handling of terminal
	rl.SetVimMode(false)

	exitChan := make(chan struct{})

	// Goroutine for receiving messages from the server
	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])
			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)
				// message = fmt.Sprintf("Listeners list updated: %s", listeners)
				// fmt.Print("\r\033[K") // Clears the line and the prompt
				// fmt.Println("\n" + message + "\n")

				// Rebuild the tab completion
				rl.Config.AutoComplete = getCompleter()
				// Refresh readline input
				rl.Refresh()
			} else if strings.HasPrefix(message, "AGENTS_UPDATE:") {
				// Extract the listener names from the message
				agents := strings.Split(strings.TrimPrefix(message, "AGENTS_UPDATE:"), ",")

				// Update subcommands
				UpdateAgentsSubCommands(agents)
				// message = fmt.Sprintf("Agents list updated: %s", agents)
				// fmt.Print("\r\033[K") // Clears the line and the prompt
				// fmt.Println("\n" + message + "\n")

				// Rebuild the tab completion
				rl.Config.AutoComplete = getCompleter()
				// Refresh readline input
				rl.Refresh()
			} else {
				// Handle other messages
				fmt.Print("\r\033[K") // Clears the line and the prompt
				fmt.Println("\n" + message)
				rl.Refresh() // Refresh readline input
			}
		}
	}()

	// Interactive terminal loop
	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)
			if command == "" {
				continue
			}

			if command == "exit" {
				fmt.Println("Exiting...")
				_, err = conn.Write([]byte(command))
				if err != nil {
					fmt.Printf("Error sending command: %v", err)
				}
				return
			}

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