Sigma-C2/operator/operator.go

package main

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

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

// 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
}

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

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{})
	doneChan := make(chan struct{}) // new: signals graceful exit

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

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