84 lines
1.8 KiB
Go
84 lines
1.8 KiB
Go
|
package websocketserver
|
||
|
|
|
||
|
|
import (
|
||
|
|
"net/http"
|
||
|
|
"log"
|
||
|
|
"sync"
|
||
|
|
"github.com/gorilla/websocket"
|
||
|
|
|
||
|
|
)
|
||
|
|
|
||
|
|
type webSocketHandler struct {
|
||
|
|
upgrader websocket.Upgrader
|
||
|
|
}
|
||
|
|
|
||
|
|
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()
|
||
|
|
|
||
|
|
for {
|
||
|
|
mt, message, err := c.ReadMessage()
|
||
|
|
if err != nil {
|
||
|
|
log.Printf("Error reading: message: %s", err)
|
||
|
|
}
|
||
|
|
|
||
|
|
log.Printf("Received message: %s", message)
|
||
|
|
if err = c.WriteMessage(mt, message); err !=nil {
|
||
|
|
log.Printf("Error writing the message: %s", err)
|
||
|
|
break
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
var agentSockets = make(map[string]*websocket.Conn)
|
||
|
|
var agentSocketsMutex sync.Mutex
|
||
|
|
|
||
|
|
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)
|
||
|
|
}
|
||
|
|
|
||
|
|
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
|
||
|
|
}
|
||
|
|
|
||
|
|
err = conn.WriteMessage(websocket.TextMessage, []byte(command))
|
||
|
|
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{},
|
||
|
|
}
|
||
|
|
webSocketMux := http.NewServeMux()
|
||
|
|
webSocketMux.Handle("/data", webSocketHandler)
|
||
|
|
webSocketMux.Handle("/executeCommand", executeCommand)
|
||
|
|
websocketServer := &http.Server{
|
||
|
|
Addr: ":5555",
|
||
|
|
Handler: webSocketMux,
|
||
|
|
}
|
||
|
|
|
||
|
|
return websocketServer
|
||
|
|
}
|