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