gontrol/src/server/websocket/websocketServer.go

package websocketserver

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

	"github.com/gorilla/websocket"
)

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
// 	}
// 	defer c.Close()

// 	_, agentNameBytes, err := c.ReadMessage()
// 	if err != nil {
// 		log.Printf("Failed to read agent name: %s", err)
// 		return
// 	}

// 	agentName := string(agentNameBytes)
// 	agentSocketsMutex.Lock()
// 	agentSockets[agentName] = c
// 	agentSocketsMutex.Unlock()

// 	log.Printf("Agent registered: %s", agentName)

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

// 		log.Printf("Received message: %s from agent: %s", message, agentName)
// 		if err = c.WriteMessage(mt, message); err !=nil {
// 			log.Printf("Error writing the message: %s", err)
// 			break
// 		}
// 	}

// 	agentSocketsMutex.Lock()
// 	delete(agentSockets, agentName)
// 	agentSocketsMutex.UnLock()
// 	log.Printf("Agent disconnected: %s", agentName)
// }

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 {
		mt, 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 err = c.WriteMessage(mt, message); err  != nil {
			log.Printf("Error writing to agent %s: %v", agentName, err)
			break
		}
	}
}

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
	}

	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
	}

	w.WriteHeader(http.StatusOK)
	w.Write([]byte("Command sent successfully"))
}

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
}