gontrol/src/server/websocket/websocketServer.go

package websocketserver

import (
	"encoding/json"
	"log"
	"net/http"
	"sync"
	"time"

	"github.com/gorilla/websocket"
)

var responseChannels sync.Map // Key: agentName, Value: chan string

type webSocketHandler struct {
	upgrader websocket.Upgrader
}

var agentSockets = make(map[string]*websocket.Conn)
var agentSocketsMutex sync.Mutex

var getAgentNames http.HandlerFunc = func(w http.ResponseWriter, r *http.Request) {
	agentSocketsMutex.Lock()
	agentNames := make([]string, 0, len(agentSockets))
	for agentName := range agentSockets {
		agentNames = append(agentNames, agentName)
	}
	agentSocketsMutex.Unlock()

	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(agentNames)

}

func (wsh webSocketHandler) ServeHTTP(w http.ResponseWriter, r *http.Request){
	c, err := wsh.upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Printf("Error %s when upgrading connection to websocket", err)
		return
	}

	agentName := r.URL.Query().Get("agentName")
	agentIP := r.URL.Query().Get("IPv4Address")
	if agentName == "" || agentIP == "" {
		log.Printf("Missing agentName or IPv4Address in query parameters")
		c.Close()
		return
	}

	log.Printf("Agent connected: %s (%s)", agentName, agentIP)

	agentSocketsMutex.Lock()
	agentSockets[agentName] = c
	agentSocketsMutex.Unlock()

	defer func() {
		agentSocketsMutex.Lock()
		delete(agentSockets, agentName)
		agentSocketsMutex.Unlock()
		c.Close()
		log.Printf("Agent disconnected: %s (%s)", agentName, agentIP)
	}()

	for {
		_, message, err := c.ReadMessage()
		if err != nil {
			log.Printf("Error reading from agent %s: %v", agentName, err)
			break
		}
		log.Printf("Message from agent %s: %s", agentName, message)


		if ch, ok := responseChannels.Load(agentName); ok {
			responseChan := ch.(chan string)
			responseChan <- string(message)
		}
	}
}

type Message struct {
	Type 	string `json:"type"`
	Payload string `json:"payload"`
}

var executeCommand http.HandlerFunc = func(w http.ResponseWriter, r *http.Request){
	err := r.ParseForm()
	if err != nil {
		http.Error(w, "Invalid form data", http.StatusBadRequest)
		return
	}

	agentName := r.FormValue("agentName")
	command := r.FormValue("command")

	agentSocketsMutex.Lock()
	conn, ok := agentSockets[agentName]
	agentSocketsMutex.Unlock()

	if !ok {
		http.Error(w, "Agent not connected", http.StatusNotFound)
		return
	}

	responseChan := make(chan string, 1)
	responseChannels.Store(agentName, responseChan)
	defer responseChannels.Delete(agentName)

	message := Message {
		Type: "command",
		Payload: command,
	}

	messageBytes, _ := json.Marshal(message)

	err = conn.WriteMessage(websocket.TextMessage, messageBytes)
	if err != nil {
		http.Error(w, "Failed to send command to the agent", http.StatusInternalServerError)
		return
	}

	select {
	case response := <-responseChan:
		var parsedResponse map[string]string
		if err := json.Unmarshal([]byte(response), &parsedResponse); err != nil {
			http.Error(w, "Failed to parse response", http.StatusInternalServerError)
			return
		}
		payload, ok := parsedResponse["payload"]
		if !ok {
			http.Error(w, "Invalid response structure", http.StatusInternalServerError)
			return
		}
		w.WriteHeader(http.StatusOK)
		w.Header().Set("Content-Type", "text/plain")
		w.Write([]byte(payload))
	case <- time.After(10 * time.Second):
			http.Error(w, "Agent repsonse timed out", http.StatusGatewayTimeout)
		}
}

func Server() (*http.Server) {


	webSocketHandler := webSocketHandler {
		upgrader: websocket.Upgrader{
			CheckOrigin: func(r *http.Request) bool {
				return true
			},
		},
	}

	corsMiddleware := func(next http.Handler) http.Handler {
        return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
            w.Header().Set("Access-Control-Allow-Origin", "*") // Allow the WebUI origin
            w.Header().Set("Access-Control-Allow-Methods", "POST, GET, OPTIONS")
            w.Header().Set("Access-Control-Allow-Headers", "Content-Type, HX-Current-URL, HX-Request, HX-Target, HX-Trigger, HX-Trigger-Name")
            if r.Method == "OPTIONS" {
                // Handle preflight requests
                w.WriteHeader(http.StatusOK)
                return
            }
            next.ServeHTTP(w, r)
        })
	}

	webSocketMux := http.NewServeMux()
	webSocketMux.Handle("/data", webSocketHandler)
	webSocketMux.Handle("/executeCommand", corsMiddleware(http.HandlerFunc(executeCommand)))
	webSocketMux.Handle("/agentNames", corsMiddleware(http.HandlerFunc(getAgentNames)))
	websocketServer := &http.Server{
		Addr: ":5555",
		Handler: webSocketMux,
	}

	return websocketServer
}