377 lines
9.0 KiB
Go
377 lines
9.0 KiB
Go
package ws
|
|
|
|
import (
|
|
"encoding/json"
|
|
"log"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
type AgentConnection struct {
|
|
AgentID string
|
|
Connection *Connection
|
|
LastSeen time.Time
|
|
}
|
|
|
|
type Hub struct {
|
|
agents map[string]*AgentConnection
|
|
agentsMux sync.RWMutex
|
|
|
|
register chan *AgentConnection
|
|
unregister chan *AgentConnection
|
|
broadcast chan []byte
|
|
|
|
tunnelManager *TunnelManager
|
|
|
|
// Response-Waiters: Command-ID -> Response-Channel
|
|
responseWaiters map[string]chan map[string]interface{}
|
|
waitersMux sync.RWMutex
|
|
|
|
// Event-Callback fuer DB-Events bei Connect/Disconnect
|
|
eventCallback func(agentID, eventType, message string)
|
|
|
|
// PTY-Sessions: session_id -> channel fuer Ausgabe an Browser
|
|
ptySessions map[string]chan []byte
|
|
ptySessionsMux sync.RWMutex
|
|
}
|
|
|
|
func NewHub() *Hub {
|
|
hub := &Hub{
|
|
agents: make(map[string]*AgentConnection),
|
|
register: make(chan *AgentConnection),
|
|
unregister: make(chan *AgentConnection),
|
|
broadcast: make(chan []byte, 256),
|
|
responseWaiters: make(map[string]chan map[string]interface{}),
|
|
}
|
|
|
|
hub.tunnelManager = NewTunnelManager(hub)
|
|
hub.ptySessions = make(map[string]chan []byte)
|
|
return hub
|
|
}
|
|
|
|
// SetEventCallback setzt die Callback-Funktion fuer Connect/Disconnect Events
|
|
func (h *Hub) SetEventCallback(cb func(agentID, eventType, message string)) {
|
|
h.eventCallback = cb
|
|
}
|
|
|
|
func (h *Hub) Run() {
|
|
log.Println("WebSocket Hub startet...")
|
|
|
|
// Cleanup-Timer für inaktive Verbindungen
|
|
cleanupTicker := time.NewTicker(5 * time.Minute)
|
|
defer cleanupTicker.Stop()
|
|
|
|
for {
|
|
select {
|
|
case agent := <-h.register:
|
|
h.registerAgent(agent)
|
|
|
|
case agent := <-h.unregister:
|
|
h.unregisterAgent(agent)
|
|
|
|
case message := <-h.broadcast:
|
|
h.broadcastMessage(message)
|
|
|
|
case <-cleanupTicker.C:
|
|
h.cleanupInactiveAgents()
|
|
}
|
|
}
|
|
}
|
|
|
|
func (h *Hub) registerAgent(agent *AgentConnection) {
|
|
h.agentsMux.Lock()
|
|
defer h.agentsMux.Unlock()
|
|
|
|
// Bestehende Verbindung schließen falls vorhanden
|
|
if existing, ok := h.agents[agent.AgentID]; ok {
|
|
log.Printf("Agent %s: Ersetzt bestehende Verbindung", agent.AgentID)
|
|
existing.Connection.Close()
|
|
}
|
|
|
|
h.agents[agent.AgentID] = agent
|
|
agent.LastSeen = time.Now()
|
|
|
|
log.Printf("Agent %s: WebSocket-Verbindung registriert", agent.AgentID)
|
|
|
|
// Event schreiben
|
|
if h.eventCallback != nil {
|
|
go h.eventCallback(agent.AgentID, "connected", "WebSocket-Verbindung hergestellt")
|
|
}
|
|
}
|
|
|
|
func (h *Hub) unregisterAgent(agent *AgentConnection) {
|
|
h.agentsMux.Lock()
|
|
defer h.agentsMux.Unlock()
|
|
|
|
if existing, ok := h.agents[agent.AgentID]; ok && existing.Connection == agent.Connection {
|
|
delete(h.agents, agent.AgentID)
|
|
existing.Connection.Close()
|
|
log.Printf("Agent %s: WebSocket-Verbindung entfernt", agent.AgentID)
|
|
|
|
// Event schreiben
|
|
if h.eventCallback != nil {
|
|
go h.eventCallback(agent.AgentID, "disconnected", "WebSocket-Verbindung getrennt")
|
|
}
|
|
|
|
// Alle Tunnel für diesen Agent schließen
|
|
h.tunnelManager.CloseAgentTunnels(agent.AgentID)
|
|
}
|
|
}
|
|
|
|
func (h *Hub) broadcastMessage(message []byte) {
|
|
h.agentsMux.RLock()
|
|
defer h.agentsMux.RUnlock()
|
|
|
|
for agentID, agent := range h.agents {
|
|
select {
|
|
case agent.Connection.Send <- message:
|
|
default:
|
|
// Agent kann keine Nachrichten mehr empfangen
|
|
log.Printf("Agent %s: Send-Channel blockiert, trenne Verbindung", agentID)
|
|
close(agent.Connection.Send)
|
|
delete(h.agents, agentID)
|
|
}
|
|
}
|
|
}
|
|
|
|
func (h *Hub) SendToAgent(agentID string, message []byte) error {
|
|
h.agentsMux.RLock()
|
|
agent, exists := h.agents[agentID]
|
|
h.agentsMux.RUnlock()
|
|
|
|
if !exists {
|
|
return ErrAgentNotConnected
|
|
}
|
|
|
|
agent.LastSeen = time.Now()
|
|
|
|
select {
|
|
case agent.Connection.Send <- message:
|
|
return nil
|
|
case <-time.After(5 * time.Second):
|
|
return ErrAgentSendBlocked
|
|
}
|
|
}
|
|
|
|
func (h *Hub) SendBinaryToAgent(agentID string, data []byte) error {
|
|
h.agentsMux.RLock()
|
|
agent, exists := h.agents[agentID]
|
|
h.agentsMux.RUnlock()
|
|
|
|
if !exists {
|
|
return ErrAgentNotConnected
|
|
}
|
|
|
|
agent.LastSeen = time.Now()
|
|
|
|
select {
|
|
case agent.Connection.BinarySend <- data:
|
|
return nil
|
|
case <-time.After(5 * time.Second):
|
|
return ErrAgentSendBlocked
|
|
}
|
|
}
|
|
|
|
func (h *Hub) IsAgentConnected(agentID string) bool {
|
|
h.agentsMux.RLock()
|
|
defer h.agentsMux.RUnlock()
|
|
|
|
_, exists := h.agents[agentID]
|
|
return exists
|
|
}
|
|
|
|
// WaitForAgentOnline polls until the agent reconnects or timeout is reached.
|
|
// Returns true if agent came online within the timeout, false otherwise.
|
|
func (h *Hub) WaitForAgentOnline(agentID string, timeout time.Duration) bool {
|
|
deadline := time.Now().Add(timeout)
|
|
for time.Now().Before(deadline) {
|
|
if h.IsAgentConnected(agentID) {
|
|
return true
|
|
}
|
|
time.Sleep(10 * time.Second)
|
|
}
|
|
return false
|
|
}
|
|
|
|
func (h *Hub) GetConnectedAgents() []string {
|
|
h.agentsMux.RLock()
|
|
defer h.agentsMux.RUnlock()
|
|
|
|
agents := make([]string, 0, len(h.agents))
|
|
for agentID := range h.agents {
|
|
agents = append(agents, agentID)
|
|
}
|
|
|
|
return agents
|
|
}
|
|
|
|
func (h *Hub) GetAgentStats() map[string]interface{} {
|
|
h.agentsMux.RLock()
|
|
defer h.agentsMux.RUnlock()
|
|
|
|
stats := map[string]interface{}{
|
|
"connected_agents": len(h.agents),
|
|
"active_tunnels": h.tunnelManager.GetActiveTunnelCount(),
|
|
}
|
|
|
|
return stats
|
|
}
|
|
|
|
func (h *Hub) cleanupInactiveAgents() {
|
|
h.agentsMux.Lock()
|
|
defer h.agentsMux.Unlock()
|
|
|
|
timeout := 10 * time.Minute
|
|
now := time.Now()
|
|
|
|
for agentID, agent := range h.agents {
|
|
if now.Sub(agent.LastSeen) > timeout {
|
|
log.Printf("Agent %s: Timeout, trenne Verbindung", agentID)
|
|
delete(h.agents, agentID)
|
|
agent.Connection.Close()
|
|
h.tunnelManager.CloseAgentTunnels(agentID)
|
|
}
|
|
}
|
|
}
|
|
|
|
// RegisterResponseWaiter registriert einen Channel fuer eine Command-Response
|
|
func (h *Hub) RegisterResponseWaiter(cmdID string) chan map[string]interface{} {
|
|
ch := make(chan map[string]interface{}, 1)
|
|
h.waitersMux.Lock()
|
|
h.responseWaiters[cmdID] = ch
|
|
h.waitersMux.Unlock()
|
|
return ch
|
|
}
|
|
|
|
// UnregisterResponseWaiter entfernt einen Response-Waiter
|
|
func (h *Hub) UnregisterResponseWaiter(cmdID string) {
|
|
h.waitersMux.Lock()
|
|
delete(h.responseWaiters, cmdID)
|
|
h.waitersMux.Unlock()
|
|
}
|
|
|
|
// GetTunnelManager gibt den TunnelManager zurück (für API-Handler)
|
|
func (h *Hub) GetTunnelManager() *TunnelManager {
|
|
return h.tunnelManager
|
|
}
|
|
|
|
// RegisterAgent registriert eine neue Agent-Verbindung
|
|
func (h *Hub) RegisterAgent(agentID string, conn *Connection) {
|
|
agent := &AgentConnection{
|
|
AgentID: agentID,
|
|
Connection: conn,
|
|
LastSeen: time.Now(),
|
|
}
|
|
|
|
h.register <- agent
|
|
}
|
|
|
|
// UnregisterAgent entfernt eine Agent-Verbindung
|
|
func (h *Hub) UnregisterAgent(agentID string, conn *Connection) {
|
|
agent := &AgentConnection{
|
|
AgentID: agentID,
|
|
Connection: conn,
|
|
}
|
|
|
|
h.unregister <- agent
|
|
}
|
|
|
|
// ProcessAgentMessage verarbeitet eingehende Nachrichten von Agents
|
|
func (h *Hub) ProcessAgentMessage(agentID string, messageType int, data []byte) error {
|
|
h.agentsMux.RLock()
|
|
agent, exists := h.agents[agentID]
|
|
h.agentsMux.RUnlock()
|
|
|
|
if exists {
|
|
agent.LastSeen = time.Now()
|
|
}
|
|
|
|
// Binary Messages: PTY oder Tunnel
|
|
if messageType == 2 { // websocket.BinaryMessage
|
|
// PTY-Session erkennen: session_id beginnt mit "pty-"
|
|
if len(data) > 1 {
|
|
idLen := int(data[0])
|
|
if 1+idLen <= len(data) {
|
|
sessionID := string(data[1 : 1+idLen])
|
|
payload := data[1+idLen:]
|
|
if len(sessionID) >= 4 && sessionID[:4] == "pty-" {
|
|
h.ptySessionsMux.RLock()
|
|
ch, ok := h.ptySessions[sessionID]
|
|
h.ptySessionsMux.RUnlock()
|
|
if ok {
|
|
select {
|
|
case ch <- payload:
|
|
default:
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
}
|
|
}
|
|
return h.tunnelManager.ProcessBinaryMessage(agentID, data)
|
|
}
|
|
|
|
// Text Messages sind JSON-Responses/Events vom Agent
|
|
var msg struct {
|
|
Type string `json:"type"`
|
|
ID string `json:"id"`
|
|
Status string `json:"status"`
|
|
Data map[string]interface{} `json:"data"`
|
|
}
|
|
if err := json.Unmarshal(data, &msg); err != nil {
|
|
return nil // Ignorieren wenn nicht parsebar
|
|
}
|
|
|
|
if msg.Type == "response" {
|
|
// Tunnel-Connect Responses verarbeiten
|
|
if msg.Status == "ok" && msg.Data != nil {
|
|
if sessionID, ok := msg.Data["session_id"].(string); ok {
|
|
h.tunnelManager.ConfirmSession(sessionID)
|
|
}
|
|
}
|
|
|
|
// Response an wartende API-Handler weiterleiten
|
|
if msg.ID != "" {
|
|
h.waitersMux.RLock()
|
|
ch, exists := h.responseWaiters[msg.ID]
|
|
h.waitersMux.RUnlock()
|
|
|
|
if exists {
|
|
resp := map[string]interface{}{
|
|
"status": msg.Status,
|
|
"data": msg.Data,
|
|
}
|
|
if msg.Status == "error" {
|
|
var errMsg struct {
|
|
Error string `json:"error"`
|
|
}
|
|
json.Unmarshal(data, &errMsg)
|
|
resp["error"] = errMsg.Error
|
|
}
|
|
|
|
select {
|
|
case ch <- resp:
|
|
default:
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
return nil
|
|
}
|
|
// RegisterPTYSession registriert einen Output-Channel fuer eine PTY-Session
|
|
func (h *Hub) RegisterPTYSession(sessionID string) chan []byte {
|
|
ch := make(chan []byte, 256)
|
|
h.ptySessionsMux.Lock()
|
|
h.ptySessions[sessionID] = ch
|
|
h.ptySessionsMux.Unlock()
|
|
return ch
|
|
}
|
|
|
|
// UnregisterPTYSession entfernt eine PTY-Session
|
|
func (h *Hub) UnregisterPTYSession(sessionID string) {
|
|
h.ptySessionsMux.Lock()
|
|
delete(h.ptySessions, sessionID)
|
|
h.ptySessionsMux.Unlock()
|
|
}
|