108 lines
2.6 KiB
Go
108 lines
2.6 KiB
Go
package main
|
|
|
|
import (
|
|
"encoding/json"
|
|
"log"
|
|
"net/url"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/gorilla/websocket"
|
|
)
|
|
|
|
type WSMessage struct {
|
|
Type string `json:"type"`
|
|
Payload map[string]interface{} `json:"payload"`
|
|
}
|
|
|
|
func connectWebSocket(cfg *Config) {
|
|
for {
|
|
// Ensure protocol is ws/wss
|
|
serverURL := strings.Replace(cfg.ServerURL, "http://", "ws://", 1)
|
|
serverURL = strings.Replace(serverURL, "https://", "wss://", 1)
|
|
|
|
u, err := url.Parse(serverURL)
|
|
if err != nil {
|
|
log.Fatalf("Invalid ServerURL: %v", err)
|
|
}
|
|
u.Path = "/api/agent/ws"
|
|
u.RawQuery = "token=" + cfg.AuthToken
|
|
|
|
log.Printf("Connecting to %s", u.String())
|
|
|
|
c, _, err := websocket.DefaultDialer.Dial(u.String(), nil)
|
|
if err != nil {
|
|
log.Printf("Dial error: %v. Retrying in 5 seconds...", err)
|
|
time.Sleep(5 * time.Second)
|
|
continue
|
|
}
|
|
|
|
log.Println("Successfully connected to SSO Manager.")
|
|
|
|
// Start telemetry if enabled
|
|
var telemetryTicker *time.Ticker
|
|
var telemetryDone chan bool
|
|
if cfg.Capabilities.Telemetry {
|
|
telemetryTicker, telemetryDone = startTelemetry(c, cfg)
|
|
}
|
|
|
|
// Read loop
|
|
for {
|
|
_, message, err := c.ReadMessage()
|
|
if err != nil {
|
|
log.Println("WebSocket read error:", err)
|
|
break // break read loop, reconnect
|
|
}
|
|
|
|
var msg WSMessage
|
|
if err := json.Unmarshal(message, &msg); err != nil {
|
|
log.Printf("Error unmarshaling message: %v", err)
|
|
continue
|
|
}
|
|
|
|
handleCommand(cfg, msg, c)
|
|
}
|
|
|
|
// Cleanup on disconnect
|
|
c.Close()
|
|
if telemetryTicker != nil {
|
|
telemetryTicker.Stop()
|
|
telemetryDone <- true
|
|
}
|
|
|
|
log.Println("WebSocket disconnected. Reconnecting in 5 seconds...")
|
|
time.Sleep(5 * time.Second)
|
|
}
|
|
}
|
|
|
|
func handleCommand(cfg *Config, msg WSMessage, c *websocket.Conn) {
|
|
log.Printf("Received command: %s", msg.Type)
|
|
|
|
switch msg.Type {
|
|
case "config":
|
|
log.Printf("Received config payload: %v", msg.Payload)
|
|
case "reboot":
|
|
if !cfg.Capabilities.Reboot {
|
|
log.Println("Reboot rejected: capability disabled in agent.yml")
|
|
return
|
|
}
|
|
log.Println("Reboot capability enabled. (Simulation: rebooting system...)")
|
|
case "service_restart":
|
|
serviceName, ok := msg.Payload["service"].(string)
|
|
if !ok || !cfg.Capabilities.CanManageService(serviceName) {
|
|
log.Printf("Service restart rejected for '%s': not in allowed service list", serviceName)
|
|
return
|
|
}
|
|
log.Printf("Restarting service %s...", serviceName)
|
|
default:
|
|
log.Printf("Unknown command type: %s", msg.Type)
|
|
}
|
|
}
|
|
|
|
func startTelemetry(c *websocket.Conn, cfg *Config) (*time.Ticker, chan bool) {
|
|
ticker := time.Ticker{C: nil} // placeholder logic
|
|
done := make(chan bool)
|
|
log.Println("Telemetry loop started (placeholder).")
|
|
return &ticker, done
|
|
}
|