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:
gurkenhabicht 2025-01-19 22:13:52 +01:00
parent 68f135fb7b
commit 3e4722e3ca
5 changed files with 334 additions and 216 deletions

6
.env
View File

@ -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"

View File

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

View File

@ -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{

View File

@ -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,

View File

@ -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">