fixed websocket messages, which were plaintext or just bytes. A JSON object was needed to differentiate between commands and what was returned to the server as a response. The random naming function now has lower case suffixes
This commit is contained in:
		
							parent
							
								
									68f135fb7b
								
							
						
					
					
						commit
						3e4722e3ca
					
				
							
								
								
									
										6
									
								
								.env
								
								
								
								
							
							
						
						
									
										6
									
								
								.env
								
								
								
								
							| 
						 | 
					@ -1,7 +1,7 @@
 | 
				
			||||||
#!/usr/bin/env sh
 | 
					#!/usr/bin/env sh
 | 
				
			||||||
 | 
					
 | 
				
			||||||
export DB_HOST="127.0.0.1"
 | 
					export DB_HOST="172.17.0.2"
 | 
				
			||||||
export DB_PORT=3306
 | 
					export DB_PORT=3306
 | 
				
			||||||
export DB_USERNAME="mysql"
 | 
					export DB_USERNAME="root"
 | 
				
			||||||
export DB_PASSWORD="mysql"
 | 
					export DB_PASSWORD="root"
 | 
				
			||||||
export DB_NAME="gomatic"
 | 
					export DB_NAME="gomatic"
 | 
				
			||||||
| 
						 | 
					
 | 
				
			||||||
| 
						 | 
					@ -2,7 +2,7 @@ package main
 | 
				
			||||||
 | 
					
 | 
				
			||||||
import (
 | 
					import (
 | 
				
			||||||
	// "bytes"
 | 
						// "bytes"
 | 
				
			||||||
	// "encoding/json"
 | 
						"encoding/json"
 | 
				
			||||||
	"fmt"
 | 
						"fmt"
 | 
				
			||||||
	"log"
 | 
						"log"
 | 
				
			||||||
	"net/http"
 | 
						"net/http"
 | 
				
			||||||
| 
						 | 
					@ -32,6 +32,11 @@ type Agent struct {
 | 
				
			||||||
	LastContact     string `json:"lastContact"`
 | 
						LastContact     string `json:"lastContact"`
 | 
				
			||||||
}
 | 
					}
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
					type Message struct {
 | 
				
			||||||
 | 
						Type 	string `json:"type"`
 | 
				
			||||||
 | 
						Payload string `json:"payload"`
 | 
				
			||||||
 | 
					}
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
					
 | 
				
			||||||
var conn *websocket.Conn
 | 
					var conn *websocket.Conn
 | 
				
			||||||
 | 
					
 | 
				
			||||||
| 
						 | 
					@ -85,22 +90,40 @@ func listenForCommands() {
 | 
				
			||||||
	defer conn.Close()
 | 
						defer conn.Close()
 | 
				
			||||||
 | 
					
 | 
				
			||||||
	for {
 | 
						for {
 | 
				
			||||||
		_, message, err := conn.ReadMessage()
 | 
							_, rawMessage, err := conn.ReadMessage()
 | 
				
			||||||
		if err != nil {
 | 
							if err != nil {
 | 
				
			||||||
			log.Printf("Error reading message: %v", err)
 | 
								log.Printf("Error reading message: %v", err)
 | 
				
			||||||
			break
 | 
								break
 | 
				
			||||||
		}
 | 
							}
 | 
				
			||||||
 | 
					
 | 
				
			||||||
		command := string(message)
 | 
							var message Message
 | 
				
			||||||
 | 
							if err := json.Unmarshal(rawMessage, &message); err != nil {
 | 
				
			||||||
 | 
								log.Printf("Error unmarshalling message: %v", err)
 | 
				
			||||||
 | 
								continue
 | 
				
			||||||
 | 
							}
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
							if message.Type != "command" {
 | 
				
			||||||
 | 
								log.Printf("Ignoring non-command message: %v", message)
 | 
				
			||||||
 | 
								continue
 | 
				
			||||||
 | 
							}
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
							command := message.Payload
 | 
				
			||||||
		log.Printf("Received command: %s", command)
 | 
							log.Printf("Received command: %s", command)
 | 
				
			||||||
 | 
					
 | 
				
			||||||
		cmd := exec.Command("bash", "-c", command)
 | 
							cmd := exec.Command("bash", "-c", command)
 | 
				
			||||||
		output, err := cmd.CombinedOutput()
 | 
							output, err := cmd.CombinedOutput()
 | 
				
			||||||
		if err != nil {
 | 
					
 | 
				
			||||||
			output = append(output, []byte(fmt.Sprintf("\n Error executing command: %v", err))...)
 | 
							response := Message{
 | 
				
			||||||
 | 
								Type: "response",
 | 
				
			||||||
 | 
								Payload: string(output),
 | 
				
			||||||
		}
 | 
							}
 | 
				
			||||||
 | 
					
 | 
				
			||||||
		if err := conn.WriteMessage(websocket.TextMessage, output); err != nil {
 | 
							if err != nil {
 | 
				
			||||||
 | 
								response.Payload += fmt.Sprintf("\n Error executing command: %v", err)
 | 
				
			||||||
 | 
							}
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
							responseBytes, _ := json.Marshal(response)
 | 
				
			||||||
 | 
							if err := conn.WriteMessage(websocket.TextMessage, responseBytes); err != nil {
 | 
				
			||||||
			log.Printf("Error sending output: %v", err)
 | 
								log.Printf("Error sending output: %v", err)
 | 
				
			||||||
			break
 | 
								break
 | 
				
			||||||
		}
 | 
							}
 | 
				
			||||||
| 
						 | 
					
 | 
				
			||||||
| 
						 | 
					@ -13,200 +13,200 @@ func GenerateRandomName() string {
 | 
				
			||||||
 | 
					
 | 
				
			||||||
var (
 | 
					var (
 | 
				
			||||||
	suffixes = []string {
 | 
						suffixes = []string {
 | 
				
			||||||
		"Ai",
 | 
							"ai",
 | 
				
			||||||
		"Albjofe",
 | 
							"albjofe",
 | 
				
			||||||
		"Alfr",
 | 
							"alfr",
 | 
				
			||||||
		"Amras",
 | 
							"amras",
 | 
				
			||||||
		"Amrod",
 | 
							"amrod",
 | 
				
			||||||
		"An",
 | 
							"an",
 | 
				
			||||||
		"Anarr",
 | 
							"anarr",
 | 
				
			||||||
		"Anduin",
 | 
							"anduin",
 | 
				
			||||||
		"Andvari",
 | 
							"andvari",
 | 
				
			||||||
		"Aragorn",
 | 
							"aragorn",
 | 
				
			||||||
		"Arathorn",
 | 
							"arathorn",
 | 
				
			||||||
		"Aredhel",
 | 
							"aredhel",
 | 
				
			||||||
		"Artanis",
 | 
							"artanis",
 | 
				
			||||||
		"Arwen",
 | 
							"arwen",
 | 
				
			||||||
		"Austri",
 | 
							"austri",
 | 
				
			||||||
		"Bafurr",
 | 
							"bafurr",
 | 
				
			||||||
		"Beleg",
 | 
							"beleg",
 | 
				
			||||||
		"Beorn",
 | 
							"beorn",
 | 
				
			||||||
		"Beren",
 | 
							"beren",
 | 
				
			||||||
		"Bilbo",
 | 
							"bilbo",
 | 
				
			||||||
		"Billingr",
 | 
							"billingr",
 | 
				
			||||||
		"Bomburr",
 | 
							"bomburr",
 | 
				
			||||||
		"Boromir",
 | 
							"boromir",
 | 
				
			||||||
		"Bruni",
 | 
							"bruni",
 | 
				
			||||||
		"Bufurr",
 | 
							"bufurr",
 | 
				
			||||||
		"Buldr",
 | 
							"buldr",
 | 
				
			||||||
		"Buri",
 | 
							"buri",
 | 
				
			||||||
		"Caranthir",
 | 
							"caranthir",
 | 
				
			||||||
		"Celeborn",
 | 
							"celeborn",
 | 
				
			||||||
		"Celebrian",
 | 
							"celebrian",
 | 
				
			||||||
		"Celebrían",
 | 
							"celebrían",
 | 
				
			||||||
		"Celegorm",
 | 
							"celegorm",
 | 
				
			||||||
		"Cirdan",
 | 
							"cirdan",
 | 
				
			||||||
		"Círdan",
 | 
							"círdan",
 | 
				
			||||||
		"Cotton",
 | 
							"cotton",
 | 
				
			||||||
		"Curufin",
 | 
							"curufin",
 | 
				
			||||||
		"Dainn",
 | 
							"dainn",
 | 
				
			||||||
		"Damrod",
 | 
							"damrod",
 | 
				
			||||||
		"Denethor",
 | 
							"denethor",
 | 
				
			||||||
		"Dilgbrasir",
 | 
							"dilgbrasir",
 | 
				
			||||||
		"Dori",
 | 
							"dori",
 | 
				
			||||||
		"Draupnir",
 | 
							"draupnir",
 | 
				
			||||||
		"Dufr",
 | 
							"dufr",
 | 
				
			||||||
		"Duilin",
 | 
							"duilin",
 | 
				
			||||||
		"Durinn",
 | 
							"durinn",
 | 
				
			||||||
		"Dvalinn",
 | 
							"dvalinn",
 | 
				
			||||||
		"Eärendil",
 | 
							"eärendil",
 | 
				
			||||||
		"Eärwen",
 | 
							"eärwen",
 | 
				
			||||||
		"Ecthelion",
 | 
							"ecthelion",
 | 
				
			||||||
		"Eikinskjaldi",
 | 
							"eikinskjaldi",
 | 
				
			||||||
		"Eldarion",
 | 
							"eldarion",
 | 
				
			||||||
		"Elendil",
 | 
							"elendil",
 | 
				
			||||||
		"Elessar",
 | 
							"elessar",
 | 
				
			||||||
		"Elladan",
 | 
							"elladan",
 | 
				
			||||||
		"Elrohir",
 | 
							"elrohir",
 | 
				
			||||||
		"Elrond",
 | 
							"elrond",
 | 
				
			||||||
		"Eluréd",
 | 
							"eluréd",
 | 
				
			||||||
		"Elurín",
 | 
							"elurín",
 | 
				
			||||||
		"Elwing",
 | 
							"elwing",
 | 
				
			||||||
		"Envinyatar",
 | 
							"envinyatar",
 | 
				
			||||||
		"Eöl",
 | 
							"eöl",
 | 
				
			||||||
		"Eomer",
 | 
							"eomer",
 | 
				
			||||||
		"Eowyn",
 | 
							"eowyn",
 | 
				
			||||||
		"Éowyn",
 | 
							"éowyn",
 | 
				
			||||||
		"Erestor",
 | 
							"erestor",
 | 
				
			||||||
		"Faniel",
 | 
							"faniel",
 | 
				
			||||||
		"Faramir",
 | 
							"faramir",
 | 
				
			||||||
		"Fëanor",
 | 
							"fëanor",
 | 
				
			||||||
		"Fili",
 | 
							"fili",
 | 
				
			||||||
		"Finduilas",
 | 
							"finduilas",
 | 
				
			||||||
		"Fingolfin",
 | 
							"fingolfin",
 | 
				
			||||||
		"Fingon",
 | 
							"fingon",
 | 
				
			||||||
		"Finnr",
 | 
							"finnr",
 | 
				
			||||||
		"Finrod",
 | 
							"finrod",
 | 
				
			||||||
		"Flalarr",
 | 
							"flalarr",
 | 
				
			||||||
		"Frar",
 | 
							"frar",
 | 
				
			||||||
		"Frawgr",
 | 
							"frawgr",
 | 
				
			||||||
		"Frodo",
 | 
							"frodo",
 | 
				
			||||||
		"Frosti",
 | 
							"frosti",
 | 
				
			||||||
		"Fundinn",
 | 
							"fundinn",
 | 
				
			||||||
		"Galadriel",
 | 
							"galadriel",
 | 
				
			||||||
		"Gandalf",
 | 
							"gandalf",
 | 
				
			||||||
		"Gandalfr",
 | 
							"gandalfr",
 | 
				
			||||||
		"Gildor",
 | 
							"gildor",
 | 
				
			||||||
		"Gil-galad",
 | 
							"gil-galad",
 | 
				
			||||||
		"Gilwen",
 | 
							"gilwen",
 | 
				
			||||||
		"Gimli",
 | 
							"gimli",
 | 
				
			||||||
		"Ginnarr",
 | 
							"ginnarr",
 | 
				
			||||||
		"Gloinn",
 | 
							"gloinn",
 | 
				
			||||||
		"Glorfindel",
 | 
							"glorfindel",
 | 
				
			||||||
		"Goldberry",
 | 
							"goldberry",
 | 
				
			||||||
		"Gollum",
 | 
							"gollum",
 | 
				
			||||||
		"Haldir",
 | 
							"haldir",
 | 
				
			||||||
		"Hanarr",
 | 
							"hanarr",
 | 
				
			||||||
		"Har",
 | 
							"har",
 | 
				
			||||||
		"Haugspori",
 | 
							"haugspori",
 | 
				
			||||||
		"Hepti",
 | 
							"hepti",
 | 
				
			||||||
		"Hirnbori",
 | 
							"hirnbori",
 | 
				
			||||||
		"Hlevangr",
 | 
							"hlevangr",
 | 
				
			||||||
		"Húrin",
 | 
							"húrin",
 | 
				
			||||||
		"Idril",
 | 
							"idril",
 | 
				
			||||||
		"Imladris",
 | 
							"imladris",
 | 
				
			||||||
		"Indis",
 | 
							"indis",
 | 
				
			||||||
		"Isildur",
 | 
							"isildur",
 | 
				
			||||||
		"Isilme",
 | 
							"isilme",
 | 
				
			||||||
		"Istarien",
 | 
							"istarien",
 | 
				
			||||||
		"Itarillë",
 | 
							"itarillë",
 | 
				
			||||||
		"Idril",
 | 
							"idril",
 | 
				
			||||||
		"Jari",
 | 
							"jari",
 | 
				
			||||||
		"Kili",
 | 
							"kili",
 | 
				
			||||||
		"Lalaith",
 | 
							"lalaith",
 | 
				
			||||||
		"Legolas",
 | 
							"legolas",
 | 
				
			||||||
		"Lindir",
 | 
							"lindir",
 | 
				
			||||||
		"Lissë",
 | 
							"lissë",
 | 
				
			||||||
		"Litr",
 | 
							"litr",
 | 
				
			||||||
		"Lofarr",
 | 
							"lofarr",
 | 
				
			||||||
		"Loni",
 | 
							"loni",
 | 
				
			||||||
		"Lothíriel",
 | 
							"lothíriel",
 | 
				
			||||||
		"Luthien",
 | 
							"luthien",
 | 
				
			||||||
		"Lúthien",
 | 
							"lúthien",
 | 
				
			||||||
		"Maedhros",
 | 
							"maedhros",
 | 
				
			||||||
		"Maeglin",
 | 
							"maeglin",
 | 
				
			||||||
		"Maglor",
 | 
							"maglor",
 | 
				
			||||||
		"Melian",
 | 
							"melian",
 | 
				
			||||||
		"Meneldor",
 | 
							"meneldor",
 | 
				
			||||||
		"Merry",
 | 
							"merry",
 | 
				
			||||||
		"Brandybuck",
 | 
							"brandybuck",
 | 
				
			||||||
		"Míriel",
 | 
							"míriel",
 | 
				
			||||||
		"Morwen",
 | 
							"morwen",
 | 
				
			||||||
		"Motsognir",
 | 
							"motsognir",
 | 
				
			||||||
		"Nainn",
 | 
							"nainn",
 | 
				
			||||||
		"Nali",
 | 
							"nali",
 | 
				
			||||||
		"Nar",
 | 
							"nar",
 | 
				
			||||||
		"Narr",
 | 
							"narr",
 | 
				
			||||||
		"Nellas",
 | 
							"nellas",
 | 
				
			||||||
		"Nerdanel",
 | 
							"nerdanel",
 | 
				
			||||||
		"Nessa",
 | 
							"nessa",
 | 
				
			||||||
		"Nienna",
 | 
							"nienna",
 | 
				
			||||||
		"Nienor",
 | 
							"nienor",
 | 
				
			||||||
		"Nimloth",
 | 
							"nimloth",
 | 
				
			||||||
		"Nimrodel",
 | 
							"nimrodel",
 | 
				
			||||||
		"Níniel",
 | 
							"níniel",
 | 
				
			||||||
		"Nioi",
 | 
							"nioi",
 | 
				
			||||||
		"Nipingr",
 | 
							"nipingr",
 | 
				
			||||||
		"Nori",
 | 
							"nori",
 | 
				
			||||||
		"Norori",
 | 
							"norori",
 | 
				
			||||||
		"Nyraor",
 | 
							"nyraor",
 | 
				
			||||||
		"Nyu",
 | 
							"nyu",
 | 
				
			||||||
		"Ori",
 | 
							"ori",
 | 
				
			||||||
		"Orodreth",
 | 
							"orodreth",
 | 
				
			||||||
		"Orophin",
 | 
							"orophin",
 | 
				
			||||||
		"Pekkr",
 | 
							"pekkr",
 | 
				
			||||||
		"Pengolodh",
 | 
							"pengolodh",
 | 
				
			||||||
		"Pippin",
 | 
							"pippin",
 | 
				
			||||||
		"Took",
 | 
							"took",
 | 
				
			||||||
		"Porinn",
 | 
							"porinn",
 | 
				
			||||||
		"Prainn",
 | 
							"prainn",
 | 
				
			||||||
		"Pror",
 | 
							"pror",
 | 
				
			||||||
		"Qurvangr",
 | 
							"qurvangr",
 | 
				
			||||||
		"Raosvidr",
 | 
							"raosvidr",
 | 
				
			||||||
		"Reginn",
 | 
							"reginn",
 | 
				
			||||||
		"Rosamunda",
 | 
							"rosamunda",
 | 
				
			||||||
		"Rúmil",
 | 
							"rúmil",
 | 
				
			||||||
		"Samwise",
 | 
							"samwise",
 | 
				
			||||||
		"Gamgee",
 | 
							"gamgee",
 | 
				
			||||||
		"Saruman",
 | 
							"saruman",
 | 
				
			||||||
		"Sauron",
 | 
							"sauron",
 | 
				
			||||||
		"Skafidr",
 | 
							"skafidr",
 | 
				
			||||||
		"Skirvir",
 | 
							"skirvir",
 | 
				
			||||||
		"Suori",
 | 
							"suori",
 | 
				
			||||||
		"Sviorr",
 | 
							"sviorr",
 | 
				
			||||||
		"Tauriel",
 | 
							"tauriel",
 | 
				
			||||||
		"Theoden",
 | 
							"theoden",
 | 
				
			||||||
		"Théoden",
 | 
							"théoden",
 | 
				
			||||||
		"Thingol",
 | 
							"thingol",
 | 
				
			||||||
		"Thorin",
 | 
							"thorin",
 | 
				
			||||||
		"Thranduil",
 | 
							"thranduil",
 | 
				
			||||||
		"Tinúviel",
 | 
							"tinúviel",
 | 
				
			||||||
		"Treebeard",
 | 
							"treebeard",
 | 
				
			||||||
		"Tuor",
 | 
							"tuor",
 | 
				
			||||||
		"Turambar",
 | 
							"turambar",
 | 
				
			||||||
		"Turgon",
 | 
							"turgon",
 | 
				
			||||||
		"Urwen",
 | 
							"urwen",
 | 
				
			||||||
		"Vairë",
 | 
							"vairë",
 | 
				
			||||||
		"Varda",
 | 
							"varda",
 | 
				
			||||||
		"Veigr",
 | 
							"veigr",
 | 
				
			||||||
		"Vestri",
 | 
							"vestri",
 | 
				
			||||||
		"Vievir",
 | 
							"vievir",
 | 
				
			||||||
		"Vili",
 | 
							"vili",
 | 
				
			||||||
		"Vindalfr",
 | 
							"vindalfr",
 | 
				
			||||||
		"Vitr",
 | 
							"vitr",
 | 
				
			||||||
		"Wormtongue",
 | 
							"wormtongue",
 | 
				
			||||||
		"Yavanna",
 | 
							"yavanna",
 | 
				
			||||||
		"Yngvi",
 | 
							"yngvi",
 | 
				
			||||||
	}
 | 
						}
 | 
				
			||||||
 | 
					
 | 
				
			||||||
	prefixes = []string{
 | 
						prefixes = []string{
 | 
				
			||||||
| 
						 | 
					
 | 
				
			||||||
| 
						 | 
					@ -1,46 +1,128 @@
 | 
				
			||||||
package websocketserver
 | 
					package websocketserver
 | 
				
			||||||
 | 
					
 | 
				
			||||||
import (
 | 
					import (
 | 
				
			||||||
	"net/http"
 | 
						"encoding/json"
 | 
				
			||||||
	"log"
 | 
						"log"
 | 
				
			||||||
 | 
						"net/http"
 | 
				
			||||||
	"sync"
 | 
						"sync"
 | 
				
			||||||
	"github.com/gorilla/websocket"
 | 
					 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
						"github.com/gorilla/websocket"
 | 
				
			||||||
)
 | 
					)
 | 
				
			||||||
 | 
					
 | 
				
			||||||
type webSocketHandler struct {
 | 
					type webSocketHandler struct {
 | 
				
			||||||
	upgrader websocket.Upgrader
 | 
						upgrader websocket.Upgrader
 | 
				
			||||||
}
 | 
					}
 | 
				
			||||||
 | 
					
 | 
				
			||||||
func (wsh webSocketHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 | 
					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)
 | 
						c, err := wsh.upgrader.Upgrade(w, r, nil)
 | 
				
			||||||
	if err != nil {
 | 
						if err != nil {
 | 
				
			||||||
		log.Printf("Error %s when upgrading connection to websocket", err)
 | 
							log.Printf("Error %s when upgrading connection to websocket", err)
 | 
				
			||||||
		return
 | 
							return
 | 
				
			||||||
	}
 | 
						}
 | 
				
			||||||
	defer c.Close()
 | 
					
 | 
				
			||||||
 | 
						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 {
 | 
						for {
 | 
				
			||||||
		mt, message, err := c.ReadMessage()
 | 
							mt, message, err := c.ReadMessage()
 | 
				
			||||||
		if err != nil {
 | 
							if err != nil {
 | 
				
			||||||
			log.Printf("Error reading: message: %s", err)
 | 
								log.Printf("Error reading from agent %s: %v", agentName, err)
 | 
				
			||||||
 | 
								break
 | 
				
			||||||
		}
 | 
							}
 | 
				
			||||||
 | 
							log.Printf("Message from agent %s: %s", agentName, message)
 | 
				
			||||||
 | 
					
 | 
				
			||||||
		log.Printf("Received message: %s", message)
 | 
							if err = c.WriteMessage(mt, message); err  != nil {
 | 
				
			||||||
		if err = c.WriteMessage(mt, message); err !=nil {
 | 
								log.Printf("Error writing to agent %s: %v", agentName, err)
 | 
				
			||||||
			log.Printf("Error writing the message: %s", err)
 | 
					 | 
				
			||||||
			break
 | 
								break
 | 
				
			||||||
		}
 | 
							}
 | 
				
			||||||
	}
 | 
						}
 | 
				
			||||||
}
 | 
					}
 | 
				
			||||||
 | 
					
 | 
				
			||||||
var agentSockets = make(map[string]*websocket.Conn)
 | 
					type Message struct {
 | 
				
			||||||
var agentSocketsMutex sync.Mutex
 | 
						Type 	string `json:"type"`
 | 
				
			||||||
 | 
						Payload string `json:"payload"`
 | 
				
			||||||
 | 
					}
 | 
				
			||||||
 | 
					
 | 
				
			||||||
var executeCommand http.HandlerFunc = func(w http.ResponseWriter, r *http.Request){
 | 
					var executeCommand http.HandlerFunc = func(w http.ResponseWriter, r *http.Request){
 | 
				
			||||||
	err := r.ParseForm()
 | 
						err := r.ParseForm()
 | 
				
			||||||
	if err != nil {
 | 
						if err != nil {
 | 
				
			||||||
		http.Error(w, "Invalid form data", http.StatusBadRequest)
 | 
							http.Error(w, "Invalid form data", http.StatusBadRequest)
 | 
				
			||||||
 | 
							return
 | 
				
			||||||
	}
 | 
						}
 | 
				
			||||||
 | 
					
 | 
				
			||||||
	agentName := r.FormValue("agentName")
 | 
						agentName := r.FormValue("agentName")
 | 
				
			||||||
| 
						 | 
					@ -55,7 +137,14 @@ var executeCommand http.HandlerFunc = func(w http.ResponseWriter, r *http.Reques
 | 
				
			||||||
		return
 | 
							return
 | 
				
			||||||
	}
 | 
						}
 | 
				
			||||||
 | 
					
 | 
				
			||||||
	err = conn.WriteMessage(websocket.TextMessage, []byte(command))
 | 
						message := Message {
 | 
				
			||||||
 | 
							Type: "command",
 | 
				
			||||||
 | 
							Payload: command,
 | 
				
			||||||
 | 
						}
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
						messageBytes, _ := json.Marshal(message)
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
						err = conn.WriteMessage(websocket.TextMessage, messageBytes)
 | 
				
			||||||
	if err != nil {
 | 
						if err != nil {
 | 
				
			||||||
		http.Error(w, "Failed to send command to the agent", http.StatusInternalServerError)
 | 
							http.Error(w, "Failed to send command to the agent", http.StatusInternalServerError)
 | 
				
			||||||
		return
 | 
							return
 | 
				
			||||||
| 
						 | 
					@ -69,11 +158,17 @@ func Server() (*http.Server) {
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
					
 | 
				
			||||||
	webSocketHandler := webSocketHandler {
 | 
						webSocketHandler := webSocketHandler {
 | 
				
			||||||
		upgrader: websocket.Upgrader{},
 | 
							upgrader: websocket.Upgrader{
 | 
				
			||||||
 | 
								CheckOrigin: func(r *http.Request) bool {
 | 
				
			||||||
 | 
									return true
 | 
				
			||||||
 | 
								},
 | 
				
			||||||
 | 
							},
 | 
				
			||||||
	}
 | 
						}
 | 
				
			||||||
 | 
					
 | 
				
			||||||
	webSocketMux := http.NewServeMux()
 | 
						webSocketMux := http.NewServeMux()
 | 
				
			||||||
	webSocketMux.Handle("/data", webSocketHandler)
 | 
						webSocketMux.Handle("/data", webSocketHandler)
 | 
				
			||||||
	webSocketMux.Handle("/executeCommand", executeCommand)
 | 
						webSocketMux.Handle("/executeCommand", executeCommand)
 | 
				
			||||||
 | 
						webSocketMux.Handle("/agentNames", getAgentNames)
 | 
				
			||||||
	websocketServer := &http.Server{
 | 
						websocketServer := &http.Server{
 | 
				
			||||||
		Addr: ":5555",
 | 
							Addr: ":5555",
 | 
				
			||||||
		Handler: webSocketMux,
 | 
							Handler: webSocketMux,
 | 
				
			||||||
| 
						 | 
					
 | 
				
			||||||
| 
						 | 
					@ -32,7 +32,7 @@
 | 
				
			||||||
            })
 | 
					            })
 | 
				
			||||||
            .catch(error => console.error('Error fetching agent names:', error));
 | 
					            .catch(error => console.error('Error fetching agent names:', error));
 | 
				
			||||||
    });
 | 
					    });
 | 
				
			||||||
</script>
 | 
					    </script>
 | 
				
			||||||
</head>
 | 
					</head>
 | 
				
			||||||
<body>
 | 
					<body>
 | 
				
			||||||
    <div class="container">
 | 
					    <div class="container">
 | 
				
			||||||
| 
						 | 
					
 | 
				
			||||||
		Loading…
	
		Reference in New Issue