websocketManager.js

Go to the documentation of this file.

/**
 * @file WebSocket Manager - Centralized Real-Time Communication Hub
 * @module services/websocketManager
 * @description
 * Enterprise WebSocket infrastructure managing bidirectional real-time communication
 * between RMM agents, dashboard clients, and remote access sessions. Provides event-driven
 * architecture for metrics streaming, command execution, file transfers, and remote desktop.
 *
 * **Architecture Overview:**
 * ```
 * ┌─────────────────┐      ┌──────────────────────┐      ┌─────────────────┐
 * │  RMM Agents     │◄────►│  WebSocketManager    │◄────►│  Dashboard UI   │
 * │  (/ws)          │      │  (Event Hub)         │      │  (/dashboard-ws)│
 * └─────────────────┘      └──────────────────────┘      └─────────────────┘
 *                                    │
 *                          ┌─────────┴─────────┐
 *                          │                   │
 *                    ┌─────▼──────┐     ┌──────▼──────┐
 *                    │ MeshCentral│     │    RDS      │
 *                    │  Proxy     │     │  Sessions   │
 *                    └────────────┘     └─────────────┘
 * ```
 *
 * **Connection Types:**
 * 
 * 1. **Agent Connections** (`/ws?agent_uuid=...`)
 *    - Persistent WebSocket from RMM agent to server
 *    - Real-time metrics streaming (CPU, memory, disk, network)
 *    - Command execution and script running
 *    - File transfer coordination
 *    - Persistent tunnel management for RDP/VNC
 * 
 * 2. **Dashboard Connections** (`/dashboard-ws?connectionId=...`)
 *    - Dashboard client subscriptions to agent streams
 *    - Live metrics, logs, process lists
 *    - Script execution monitoring
 *    - Real-time alerts and notifications
 * 
 * 3. **RDS Sessions** (`/api/rds/:sessionId`)
 *    - Remote Desktop Service broker
 *    - Dual WebSocket coordination (viewer ↔ agent)
 *    - Mouse/keyboard input relay
 *    - Screen update streaming
 * 
 * 4. **MeshCentral Proxy** (`/api/meshcentral/terminal`, `/api/meshcentral/files`)
 *    - WebSocket-to-WebSocket proxying for MeshCentral integration
 *    - Terminal sessions and file browser
 *
 * **Message Protocol:**
 * 
 * All WebSocket messages use JSON format:
 * ```json
 * {
 *   "type": "metrics|script_output|command|subscribe|...",
 *   "data": { ... },
 *   "requestId": "optional-correlation-id"
 * }
 * ```
 *
 * **Agent Message Types:**
 * - `connected` - Agent initial handshake
 * - `metrics` - Performance metrics (CPU, RAM, disk, network)
 * - `processes` - Running process snapshot
 * - `services` - Windows services status
 * - `script_output` - Script execution output chunk
 * - `script_complete` - Script execution finished
 * - `tunnel_ready` - Persistent tunnel established (RDP/VNC)
 * - `response` - Response to request (with requestId correlation)
 *
 * **Dashboard Message Types:**
 * - `subscribe` - Subscribe to agent's real-time data
 * - `unsubscribe` - Unsubscribe from agent
 * - `execute_script` - Run script on agent
 * - `request_processes` - Query running processes
 * - `request_files` - Query filesystem
 *
 * **State Management:**
 * - `agentConnections`: Map<agent_uuid, WebSocket> - Active agent connections
 * - `dashboardConnections`: Map<connection_id, {ws, subscriptions}> - Dashboard subscriptions
 * - `subscriptions`: Map<agent_uuid, Set<connection_id>> - Agent → Subscribers index
 * - `activeExecutions`: Map<execution_id, {...}> - Running script tracking
 * - `rdsSessions`: Map<sessionId, {...}> - Remote desktop sessions
 * - `persistentTunnels`: Map<agent_uuid, {...}> - RDP/VNC tunnel state
 *
 * **Key Features:**
 * - **Automatic reconnection**: Agents reconnect and resume subscriptions
 * - **Request/response**: Promise-based request-reply pattern with timeouts
 * - **Broadcast**: Efficient fan-out to multiple dashboard subscribers
 * - **Session cleanup**: Automatic cleanup on disconnect
 * - **Port allocation**: Dynamic port assignment for persistent tunnels (45000+)
 * - **Metrics persistence**: Network I/O saved to database for history
 * @class WebSocketManager
 * @augments EventEmitter
 * @example
 * // Initialize WebSocket manager
 * const WebSocketManager = require('./services/websocketManager');
 * const wsManager = new WebSocketManager();
 * wsManager.init(httpServer);
 * @example
 * // Send command to agent
 * wsManager.sendToAgent(agent_uuid, {
 *   type: 'execute_script',
 *   data: { script_id: 123, code: 'Get-Service' }
 * });
 * @example
 * // Request data from agent with timeout
 * const response = await wsManager.requestFromAgent(agent_uuid, {
 *   type: 'get_processes'
 * }, 5000);
 * @example
 * // Broadcast metrics to all subscribers
 * wsManager.broadcastToSubscribers(agent_uuid, {
 *   type: 'metrics',
 *   data: { cpu: 45.2, memory: 78.5, disk: 65.0 }
 * });
 * @see {@link module:routes/agent} for agent registration endpoints
 * @see {@link module:services/metricsAggregator} for metrics buffering
 * @see {@link module:services/guacamoleService} for RDS session management
 * @requires ws - WebSocket server implementation
 * @requires events - EventEmitter for pub/sub
 * @requires db - PostgreSQL connection pool from services/db
 */
 
/**
 * WebSocket Manager - Centralized real-time communication
 * 
 * Handles:
 * - Agent connections and message routing
 * - Dashboard subscriptions to agent data streams
 * - Real-time metrics, processes, logs, events
 * - Script execution with output streaming
 * - RDS (Remote Desktop) sessions
 */
 
const { WebSocketServer } = require('ws');
const url = require('url');
const EventEmitter = require('events');
const pool = require('./db');
 
/**
 * WebSocket Manager class.
 * Handles real-time bidirectional communication for agents, dashboards, and RDS sessions.
 */
class WebSocketManager extends EventEmitter {
  /**
   * Initialize WebSocketManager.
   * Sets up connection maps, subscriptions, and execution tracking.
   */
  constructor() {
    super();
    
    // Agent connections: agent_uuid → ws
    this.agentConnections = new Map();
    
    // Dashboard connections: connection_id → { ws, subscriptions: Set<agent_uuid> }
    this.dashboardConnections = new Map();
    
    // Active subscriptions: agent_uuid → Set<connection_id>
    this.subscriptions = new Map();
    
    // Script executions: execution_id → { agent_uuid, dashboard_connections: Set<connection_id> }
    this.activeExecutions = new Map();
    
    // RDS sessions: sessionId → { agentId, viewer, agent, control }
    this.rdsSessions = new Map();
    
    // Persistent tunnels: agent_uuid → { protocol, localPort, assignedPort, sessionId, ready }
    this.persistentTunnels = new Map();
    
    // Port assignments: track which ports are in use
    this.assignedPorts = new Set();
    this.nextAvailablePort = 45000; // Start assigning ports from 45000
    
    // Pending requests: requestId → { resolve, reject, timeout }
    this.pendingRequests = new Map();
  }
 
  /**
   * Initialize WebSocket server and attach to HTTP server
   * @description Sets up WebSocket server to handle multiple connection types:
   * - `/ws` - Agent connections
   * - `/dashboard-ws` - Dashboard subscriptions
   * - `/api/rds/:sessionId` - Remote Desktop sessions
   * - `/api/meshcentral/terminal` - MeshCentral terminal proxy
   * - `/api/meshcentral/files` - MeshCentral file browser proxy
   * @param {object} server - HTTP/HTTPS server instance to attach WebSocket upgrade handler
   * @returns {WebSocketManager} Returns this instance for method chaining
   * @fires connection WebSocket connection established event
   */
  init(server) {
    const wss = new WebSocketServer({ noServer: true });
 
    // Handle upgrade requests
    server.on('upgrade', (req, socket, head) => {
      const pathname = req.url.split('?')[0];
 
      if (
        pathname === '/ws' ||
        pathname === '/dashboard-ws' ||
        pathname.startsWith('/api/rds/') ||
        pathname === '/api/meshcentral/terminal' ||
        pathname === '/api/meshcentral/files'
      ) {
        wss.handleUpgrade(req, socket, head, (ws) => {
          wss.emit('connection', ws, req);
        });
      } else {
        socket.destroy();
      }
    });
 
    // Handle connections
    wss.on('connection', (ws, req) => {
      const parsed = url.parse(req.url, true);
      const path = parsed.pathname;
 
      if (path === '/ws') {
        this.handleAgentConnection(ws, parsed.query);
      } else if (path === '/dashboard-ws') {
        this.handleDashboardConnection(ws, parsed.query);
      } else if (path.startsWith('/api/rds/')) {
        this.handleRDSConnection(ws, req);
      } else if (path === '/api/meshcentral/terminal') {
        this.handleMeshCentralTerminal(ws, parsed.query);
      } else if (path === '/api/meshcentral/files') {
        this.handleMeshCentralFiles(ws, parsed.query);
      } else {
        ws.close();
      }
    });
 
    console.log('✅ WebSocket Manager initialized');
    return this;
  }
 
  /**
   * Handle new agent WebSocket connection
   * @description Registers agent connection, stores in connection map, and sets up message
   * handlers for metrics, script output, and commands. Updates agent last_seen timestamp.
   * @param {WebSocket} ws - WebSocket connection from agent
   * @param {object} query - URL query parameters from connection request
   * @param {string} query.agent_uuid - Unique agent identifier (required)
   * @returns {void}
   * @fires agent_disconnected When agent connection closes
   */
  handleAgentConnection(ws, query) {
    const agent_uuid = query.agent_uuid;
    if (!agent_uuid) {
      console.log('[WS] Agent connection rejected: missing agent_uuid');
      return ws.close();
    }
 
    // Store connection
    this.agentConnections.set(agent_uuid, ws);
    console.log(`[WS] Agent connected: ${agent_uuid} (${this.agentConnections.size} total)`);
 
    // Update database last_seen
    this.updateLastSeen(agent_uuid);
 
    // Handle messages from agent
    ws.on('message', (raw) => {
      try {
        const msg = JSON.parse(raw.toString());
        this.handleAgentMessage(agent_uuid, msg, ws);
      } catch (err) {
        console.error(`[WS] Invalid message from agent ${agent_uuid}:`, err.message);
      }
    });
 
    // Handle disconnect
    ws.on('close', () => {
      this.agentConnections.delete(agent_uuid);
      console.log(`[WS] Agent disconnected: ${agent_uuid}`);
      this.updateLastSeen(agent_uuid);
      this.cleanupAgentSessions(agent_uuid);
      
      // Notify subscribed dashboards
      this.broadcastToSubscribers(agent_uuid, {
        type: 'agent_disconnected',
        agent_uuid,
        timestamp: Date.now()
      });
    });
 
    // Send welcome message
    ws.send(JSON.stringify({
      type: 'welcome',
      agent_uuid,
      timestamp: Date.now()
    }));
  }
 
  /**
   * Handle new dashboard WebSocket connection
   * @description Creates dashboard connection with unique ID, manages subscriptions to agents,
   * and handles subscribe/unsubscribe messages.
   * @param {WebSocket} ws - WebSocket connection from dashboard client
   * @returns {void}
   */
  handleDashboardConnection(ws) {
    const connection_id = `dashboard-${Date.now()}-${Math.random().toString(36).slice(2)}`;
    
    this.dashboardConnections.set(connection_id, {
      ws,
      subscriptions: new Set(),
      authenticated: false,
      user_id: null
    });
 
    console.log(`[WS] Dashboard connected: ${connection_id}`);
 
    // Handle messages from dashboard
    ws.on('message', (raw) => {
      try {
        const msg = JSON.parse(raw.toString());
        this.handleDashboardMessage(connection_id, msg, ws);
      } catch (err) {
        console.error(`[WS] Invalid message from dashboard ${connection_id}:`, err.message);
      }
    });
 
    // Handle disconnect
    ws.on('close', () => {
      const conn = this.dashboardConnections.get(connection_id);
      if (conn) {
        // Unsubscribe from all agents
        for (const agent_uuid of conn.subscriptions) {
          const subs = this.subscriptions.get(agent_uuid);
          if (subs) {
            subs.delete(connection_id);
            if (subs.size === 0) {
              this.subscriptions.delete(agent_uuid);
            }
          }
        }
      }
      this.dashboardConnections.delete(connection_id);
      console.log(`[WS] Dashboard disconnected: ${connection_id}`);
    });
 
    // Send welcome
    ws.send(JSON.stringify({
      type: 'welcome',
      connection_id,
      timestamp: Date.now()
    }));
  }
 
  /**
   * Process and route messages received from agent
   * @description Handles all message types from agents including metrics, script output,
   * tunnel status, and responses to requests. Broadcasts relevant messages to subscribed dashboards.
   * @param {string} agent_uuid - Unique agent identifier
   * @param {object} msg - Parsed message object from agent
   * @param {string} msg.type - Message type (metrics|script_output|tunnel_ready|response|...)
   * @param {object} msg.data - Message payload data
   * @param {string} [msg.requestId] - Optional correlation ID for request/response pattern
   * @param {WebSocket} ws - Agent's WebSocket connection
   * @returns {Promise<void>}
   */
  async handleAgentMessage(agent_uuid, msg, ws) {
    console.log(`[WS] Agent ${agent_uuid} → ${msg.type}`);
 
    switch (msg.type) {
      case 'agent_hello':
        // Update agent version in database if provided
        if (msg.version) {
          try {
            await pool.query(
              `UPDATE agents SET version = $1, last_seen = NOW() WHERE agent_uuid = $2`,
              [msg.version, agent_uuid]
            );
            console.log(`[WS] Updated agent ${agent_uuid} to version ${msg.version}`);
          } catch (err) {
            console.error('[WS] Failed to update agent version:', err);
          }
        }
        
        // Send plugin enabled/disabled status back to agent
        try {
          const pluginsResult = await pool.query(
            `SELECT name, enabled FROM plugins ORDER BY name ASC`
          );
          const pluginStatus = {};
          pluginsResult.rows.forEach(row => {
            pluginStatus[row.name] = row.enabled;
          });
          
          // Send plugin status to agent
          ws.send(JSON.stringify({
            type: 'plugin_status_update',
            plugins: pluginStatus,
            timestamp: Date.now()
          }));
          
          console.log(`[WS] Sent plugin status to agent ${agent_uuid}: ${Object.keys(pluginStatus).length} plugins`);
        } catch (err) {
          console.error('[WS] Failed to fetch plugin status:', err);
        }
        
        this.emit('agent_hello', { agent_uuid, ...msg });
        break;
 
      case 'metrics_update':
        // Save network info (IP) to database for offline fallback
        this.saveNetworkInfo(agent_uuid, msg.data);
        
        // Broadcast to subscribed dashboards
        this.broadcastToSubscribers(agent_uuid, {
          type: 'metrics_update',
          agent_uuid,
          data: msg.data,
          timestamp: msg.timestamp || Date.now()
        });
        break;
 
      case 'processes_snapshot':
        this.broadcastToSubscribers(agent_uuid, {
          type: 'processes_snapshot',
          agent_uuid,
          processes: msg.processes,
          timestamp: msg.timestamp || Date.now()
        });
        break;
 
      case 'services_snapshot':
        this.broadcastToSubscribers(agent_uuid, {
          type: 'services_snapshot',
          agent_uuid,
          services: msg.services,
          timestamp: msg.timestamp || Date.now()
        });
        break;
 
      case 'script_output':
        // Forward to dashboards watching this execution
        this.broadcastScriptOutput(msg.execution_id, {
          type: 'script_output',
          execution_id: msg.execution_id,
          line: msg.line,
          timestamp: msg.timestamp || Date.now()
        });
        break;
 
      case 'script_complete':
        this.broadcastScriptOutput(msg.execution_id, {
          type: 'script_complete',
          execution_id: msg.execution_id,
          exit_code: msg.exit_code,
          timestamp: msg.timestamp || Date.now()
        });
        this.activeExecutions.delete(msg.execution_id);
        break;
 
      case 'agent_log':
        this.broadcastToSubscribers(agent_uuid, {
          type: 'agent_log',
          agent_uuid,
          level: msg.level,
          message: msg.message,
          timestamp: msg.timestamp || Date.now()
        });
        break;
 
      case 'event_log':
        this.broadcastToSubscribers(agent_uuid, {
          type: 'event_log',
          agent_uuid,
          event: msg.event,
          timestamp: msg.timestamp || Date.now()
        });
        break;
 
      case 'response':
        // Handle response to dashboard request
        if (msg.request_id) {
          const pending = this.pendingRequests.get(msg.request_id);
          if (pending) {
            clearTimeout(pending.timeout);
            pending.resolve(msg.data);
            this.pendingRequests.delete(msg.request_id);
          }
        }
        break;
 
      case 'file-list-response':
        // Handle file list response (existing RDS functionality)
        if (msg.requestId) {
          const pending = this.pendingRequests.get(msg.requestId);
          if (pending) {
            clearTimeout(pending.timeout);
            pending.resolve(msg.files || []);
            this.pendingRequests.delete(msg.requestId);
          }
        }
        break;
      
      case 'file_download_chunk':
      case 'file_download_complete':
      case 'file_download_error':
      case 'file_upload_success':
      case 'file_upload_error': {
        // Forward to file transfer service
        const fileTransferService = require('./fileTransferService');
        if (msg.type === 'file_download_chunk') {
          fileTransferService.handleDownloadChunk(msg);
        } else if (msg.type === 'file_download_complete') {
          fileTransferService.handleDownloadComplete(msg);
        }
        break;
      }
      
      case 'connection_approval_response': {
        // Forward to tunnel service via event emitter
        this.emit('connection_approval_response', {
          agent_uuid,
          requestId: msg.requestId,
          approved: msg.approved,
          timestamp: msg.timestamp
        });
        console.log(`[WS] Forwarded approval response from agent ${agent_uuid}: ${msg.approved ? 'APPROVED' : 'DENIED'}`);
        break;
      }
      
      case 'rds_data': {
        // Handle RDS data from agent (new single-WebSocket architecture)
        // Forward to viewer WebSocket
        const sessionId = msg.sessionId;
        const session = this.rdsSessions.get(sessionId);
        if (session && session.viewer && session.viewer.readyState === 1) {
          // Decode base64 and send raw binary to viewer
          const buffer = Buffer.from(msg.data, 'base64');
          session.viewer.send(buffer);
        } else {
          console.warn(`[RDS] Received data for session ${sessionId} but no viewer connected`);
        }
        break;
      }
      
      case 'rds_tunnel_closed': {
        // Handle tunnel closure notification from agent
        const closedSessionId = msg.sessionId;
        const closedSession = this.rdsSessions.get(closedSessionId);
        if (closedSession) {
          // Only clean up if NOT a persistent tunnel
          if (!closedSession.persistent) {
            // Notify viewer
            if (closedSession.viewer && closedSession.viewer.readyState === 1) {
              closedSession.viewer.close();
            }
            // Notify control
            if (closedSession.control && closedSession.control.readyState === 1) {
              closedSession.control.close();
            }
            this.rdsSessions.delete(closedSessionId);
            console.log(`[RDS] Session ${closedSessionId} closed by agent`);
          } else {
            console.log(`[RDS] Ignored tunnel close for persistent session ${closedSessionId}`);
          }
        }
        break;
      }
      
      case 'rds_started': {
        // Handle RDS tunnel started confirmation
        console.log(`[RDS] Agent confirmed tunnel started: ${msg.sessionId}, proto: ${msg.proto}, port: ${msg.port}`);
        const startedSession = this.rdsSessions.get(msg.sessionId);
        if (startedSession) {
          startedSession.proto = msg.proto;
          startedSession.port = msg.port;
          startedSession.active = true;
        }
        break;
      }
      
      case 'rds_failed': {
        // Handle RDS tunnel failure
        console.error(`[RDS] Agent reported tunnel failure: ${msg.sessionId}, reason: ${msg.reason}`);
        const failedSession = this.rdsSessions.get(msg.sessionId);
        if (failedSession) {
          if (failedSession.viewer) failedSession.viewer.close();
          if (failedSession.control) failedSession.control.close();
          this.rdsSessions.delete(msg.sessionId);
        }
        break;
      }
 
      case 'persistent_tunnel_request':
        // Handle agent requesting persistent tunnel assignment
        this.handlePersistentTunnelRequest(agent_uuid, msg, ws);
        break;
 
      case 'persistent_tunnel_ready':
        // Handle agent confirming persistent tunnel is ready
        this.handlePersistentTunnelReady(agent_uuid, msg);
        break;
 
      default:
        console.log(`[WS] Unknown message type from agent: ${msg.type}`);
    }
  }
 
  /**
   * Process and route messages received from dashboard clients
   * @description Handles dashboard commands including authenticate, subscribe/unsubscribe,
   * request processes/services, execute scripts, and control services on agents.
   * @param {string} connection_id - Unique dashboard connection identifier
   * @param {object} msg - Parsed message object from dashboard
   * @param {string} msg.type - Message type (authenticate|subscribe|execute_script|...)
   * @param {object} [msg.data] - Message payload data
   * @param {WebSocket} ws - Dashboard's WebSocket connection
   * @returns {void}
   */
  handleDashboardMessage(connection_id, msg, ws) {
    console.log(`[WS] Dashboard ${connection_id} → ${msg.type}`);
 
    const conn = this.dashboardConnections.get(connection_id);
    if (!conn) return;
 
    switch (msg.type) {
      case 'authenticate':
        // Validate JWT token
        try {
          const jwt = require('jsonwebtoken');
          const JWT_SECRET = process.env.JWT_SECRET;
          
          if (!msg.token) {
            ws.send(JSON.stringify({
              type: 'auth_error',
              error: 'No token provided',
              timestamp: Date.now()
            }));
            return;
          }
          
          const decoded = jwt.verify(msg.token, JWT_SECRET);
          conn.authenticated = true;
          conn.user_id = decoded.user_id || decoded.userId;
          conn.tenant_id = decoded.tenant_id || decoded.tenantId;
          
          ws.send(JSON.stringify({
            type: 'authenticated',
            connection_id,
            user_id: conn.user_id,
            tenant_id: conn.tenant_id,
            timestamp: Date.now()
          }));
          
          console.log(`[WS] Dashboard ${connection_id} authenticated as user ${conn.user_id}`);
        } catch (err) {
          console.error(`[WS] Dashboard ${connection_id} authentication failed:`, err.message);
          ws.send(JSON.stringify({
            type: 'auth_error',
            error: 'Invalid token',
            timestamp: Date.now()
          }));
        }
        break;
 
      case 'subscribe':
        // Subscribe to agent updates
        if (!Array.isArray(msg.agent_uuids)) return;
        
        for (const agent_uuid of msg.agent_uuids) {
          conn.subscriptions.add(agent_uuid);
          
          if (!this.subscriptions.has(agent_uuid)) {
            this.subscriptions.set(agent_uuid, new Set());
          }
          this.subscriptions.get(agent_uuid).add(connection_id);
          
          console.log(`[WS] Dashboard ${connection_id} subscribed to agent ${agent_uuid}`);
        }
 
        ws.send(JSON.stringify({
          type: 'subscribed',
          agent_uuids: msg.agent_uuids,
          timestamp: Date.now()
        }));
        break;
 
      case 'unsubscribe':
        if (!Array.isArray(msg.agent_uuids)) return;
        
        for (const agent_uuid of msg.agent_uuids) {
          conn.subscriptions.delete(agent_uuid);
          
          const subs = this.subscriptions.get(agent_uuid);
          if (subs) {
            subs.delete(connection_id);
            if (subs.size === 0) {
              this.subscriptions.delete(agent_uuid);
            }
          }
        }
        break;
 
      case 'request_processes':
        this.sendToAgent(msg.agent_uuid, {
          type: 'get_processes',
          request_id: msg.request_id
        });
        break;
 
      case 'request_services':
        this.sendToAgent(msg.agent_uuid, {
          type: 'get_services',
          request_id: msg.request_id
        });
        break;
 
      case 'execute_script':
        this.executeScript(msg.agent_uuid, msg.script_id, msg.code, connection_id);
        break;
 
      case 'kill_process':
        this.sendToAgent(msg.agent_uuid, {
          type: 'kill_process',
          pid: msg.pid,
          request_id: msg.request_id
        });
        break;
 
      case 'control_service':
        this.sendToAgent(msg.agent_uuid, {
          type: 'control_service',
          service_name: msg.service_name,
          action: msg.action, // start, stop, restart
          request_id: msg.request_id
        });
        break;
 
      default:
        console.log(`[WS] Unknown message type from dashboard: ${msg.type}`);
    }
  }
 
  /**
   * Send message to specific agent
   * @description Sends JSON message to agent if connected. Logs approval requests and
   * handles connection failures gracefully.
   * @param {string} agent_uuid - Target agent's unique identifier
   * @param {object} message - Message object to send
   * @param {string} message.type - Message type
   * @param {object} [message.data] - Message payload
   * @param {string} [message.request_id] - Optional correlation ID for responses
   * @returns {boolean} True if message sent successfully, false if agent not connected
   */
  sendToAgent(agent_uuid, message) {
    const ws = this.agentConnections.get(agent_uuid);
    if (!ws || ws.readyState !== 1) {
      console.log(`[WS] Agent ${agent_uuid} not connected`);
      return false;
    }
 
    try {
      // Log approval requests specifically
      if (message.type === 'connection_approval_request') {
        console.log(`[WS] 📤 Sending approval request to agent ${agent_uuid}:`, JSON.stringify(message));
      }
      ws.send(JSON.stringify(message));
      console.log(`[WS] ✅ Message sent to agent ${agent_uuid}, type: ${message.type}`);
      return true;
    } catch (err) {
      console.error(`[WS] Failed to send to agent ${agent_uuid}:`, err.message);
      return false;
    }
  }
 
  /**
   * Broadcast message to all dashboards subscribed to agent
   * @description Efficiently fan-out message to multiple dashboard subscribers.
   * Automatically handles connection cleanup for closed connections.
   * @param {string} agent_uuid - Agent whose subscribers should receive message
   * @param {object} message - Message object to broadcast
   * @param {string} message.type - Message type (metrics|agent_disconnected|...)
   * @param {object} [message.data] - Message payload
   * @returns {void}
   */
  broadcastToSubscribers(agent_uuid, message) {
    const subs = this.subscriptions.get(agent_uuid);
    if (!subs || subs.size === 0) return;
 
    const payload = JSON.stringify(message);
    let sent = 0;
 
    for (const connection_id of subs) {
      const conn = this.dashboardConnections.get(connection_id);
      if (conn && conn.ws.readyState === 1) {
        try {
          conn.ws.send(payload);
          sent++;
        } catch (err) {
          console.error(`[WS] Failed to send to dashboard ${connection_id}:`, err.message);
        }
      }
    }
 
    if (sent > 0) {
      console.log(`[WS] Broadcasted ${message.type} to ${sent} dashboards`);
    }
  }
 
  /**
   * Execute script on agent and stream output to dashboard
   * @description Creates execution session, sends script to agent, and registers
   * dashboard for output streaming. Supports PowerShell, Bash, or Python scripts.
   * @param {string} agent_uuid - Target agent's unique identifier
   * @param {number} script_id - Script ID from database
   * @param {string} code - Script code to execute
   * @param {string} dashboard_connection_id - Dashboard requesting execution
   * @returns {string} execution_id - Unique execution identifier for tracking
   */
  executeScript(agent_uuid, script_id, code, dashboard_connection_id) {
    const execution_id = `exec-${Date.now()}-${Math.random().toString(36).slice(2)}`;
    
    this.activeExecutions.set(execution_id, {
      agent_uuid,
      dashboard_connections: new Set([dashboard_connection_id])
    });
 
    this.sendToAgent(agent_uuid, {
      type: 'execute_script',
      execution_id,
      script_id,
      code
    });
 
    return execution_id;
  }
 
  /**
   * Broadcast script output to watching dashboards
   * @description Sends script output chunk or completion message to all dashboards
   * watching this execution.
   * @param {string} execution_id - Script execution identifier
   * @param {object} message - Output message to broadcast
   * @param {string} message.type - Message type (script_output|script_complete)
   * @param {object} message.data - Output data (stdout, stderr, exit code)
   * @returns {void}
   */
  broadcastScriptOutput(execution_id, message) {
    const execution = this.activeExecutions.get(execution_id);
    if (!execution) return;
 
    const payload = JSON.stringify(message);
 
    for (const connection_id of execution.dashboard_connections) {
      const conn = this.dashboardConnections.get(connection_id);
      if (conn && conn.ws.readyState === 1) {
        try {
          conn.ws.send(payload);
        } catch (err) {
          console.error(`[WS] Failed to send script output to ${connection_id}:`, err.message);
        }
      }
    }
  }
 
  /**
   * Send request to agent and wait for response
   * @description Implements request-reply pattern with timeout. Correlates responses
   * using request_id. Used for synchronous queries like get_processes, get_services.
   * @param {string} agent_uuid - Target agent's unique identifier
   * @param {object} message - Request message object
   * @param {string} message.type - Request type (get_processes|get_services|...)
   * @param {number} [timeout] - Timeout in milliseconds
   * @returns {Promise<object>} Resolves with agent's response data, rejects on timeout
   * @throws {Error} Throws 'Request timeout' if agent doesn't respond within timeout
   * @throws {Error} Throws 'Agent not connected' if agent is offline
   */
  requestFromAgent(agent_uuid, message, timeout = 10000) {
    return new Promise((resolve, reject) => {
      const request_id = `req-${Date.now()}-${Math.random().toString(36).slice(2)}`;
      
      const timer = setTimeout(() => {
        this.pendingRequests.delete(request_id);
        reject(new Error('Request timeout'));
      }, timeout);
 
      this.pendingRequests.set(request_id, { resolve, reject, timeout: timer });
 
      const success = this.sendToAgent(agent_uuid, {
        ...message,
        request_id
      });
 
      if (!success) {
        clearTimeout(timer);
        this.pendingRequests.delete(request_id);
        reject(new Error('Agent not connected'));
      }
    });
  }
 
  /**
   * Request file list from agent.
   * Sends file list request to agent and returns promise with results.
   * @param {string} agentId - Agent UUID
   * @param {string} path - File system path to list
   * @returns {Promise<Array>} Resolves with file list array
   * @throws {Error} Agent not connected or timeout
   */
  requestAgentFileList(agentId, path) {
    return new Promise((resolve, reject) => {
      const ws = this.agentConnections.get(agentId);
      if (!ws || ws.readyState !== 1) {
        return reject(new Error('Agent not connected'));
      }
 
      const requestId = `filelist-${Date.now()}-${Math.random().toString(36).slice(2)}`;
      
      const timer = setTimeout(() => {
        this.pendingRequests.delete(requestId);
        reject(new Error('Timeout waiting for agent file list'));
      }, 10000);
 
      this.pendingRequests.set(requestId, { resolve, reject, timeout: timer });
 
      // Send file list request
      ws.send(JSON.stringify({ type: 'file-list', requestId, path }));
    });
  }
 
  /**
   * Update agent last_seen in database.
   * @param {string} agent_uuid - Agent UUID to update
   * @returns {Promise<void>}
   */
  async updateLastSeen(agent_uuid) {
    try {
      const pool = require('./db');
      await pool.query(
        'UPDATE agents SET last_seen = NOW() WHERE agent_uuid = $1',
        [agent_uuid]
      );
    } catch (err) {
      console.error(`[WS] Failed to update last_seen for ${agent_uuid}:`, err.message);
    }
  }
 
  /**
   * Save network and disk info from WebSocket metrics to database.
   * This ensures offline fallback data is current.
   * @param {string} agent_uuid - Agent UUID
   * @param {object} metrics - Metrics object with network and disk data
   * @returns {Promise<void>}
   */
  async saveNetworkInfo(agent_uuid, metrics) {
    try {
      const pool = require('./db');
      
      // Extract IP address from metrics
      if (metrics.network && metrics.network.ipv4) {
        await pool.query(
          'UPDATE agents SET ip_address = $1, updated_at = NOW() WHERE agent_uuid = $2',
          [metrics.network.ipv4, agent_uuid]
        );
      }
      
      // Update system_info with current drives data
      if (metrics.disk && Array.isArray(metrics.disk)) {
        // Get existing system_info and merge drives data
        const result = await pool.query(
          'SELECT system_info FROM agents WHERE agent_uuid = $1',
          [agent_uuid]
        );
        
        if (result.rows.length > 0) {
          const systemInfo = result.rows[0].system_info || {};
          
          // Convert WebSocket disk format to drives format
          const drives = metrics.disk.map(d => ({
            Drive: d.mount || d.filesystem || 'Unknown',
            Root: d.mount || d.filesystem || 'Unknown',
            Used: Math.round((d.used / (1024 * 1024 * 1024)) * 100) / 100, // Convert to GB
            Free: Math.round(((d.total - d.used) / (1024 * 1024 * 1024)) * 100) / 100,
            Total: Math.round((d.total / (1024 * 1024 * 1024)) * 100) / 100,
            UsedPercent: Math.round(d.used_percent * 10) / 10
          }));
          
          systemInfo.drives = drives;
          
          await pool.query(
            'UPDATE agents SET system_info = $1, updated_at = NOW() WHERE agent_uuid = $2',
            [JSON.stringify(systemInfo), agent_uuid]
          );
        }
      }
    } catch (err) {
      console.error(`[WS] Failed to save network info for ${agent_uuid}:`, err.message);
    }
  }
 
  /**
   * Cleanup sessions when agent disconnects.
   * Removes RDS sessions and active executions for disconnected agent.
   * @param {string} agent_uuid - Agent UUID to clean up
   * @returns {void}
   */
  cleanupAgentSessions(agent_uuid) {
    // Clean up RDS sessions
    for (const [sessionId, session] of this.rdsSessions.entries()) {
      if (session.agentId === agent_uuid) {
        if (session.viewer) session.viewer.close();
        if (session.control) session.control.close();
        this.rdsSessions.delete(sessionId);
        console.log(`[RDS] Cleaned session ${sessionId} for agent ${agent_uuid}`);
      }
    }
 
    // Clean up active executions
    for (const [execution_id, execution] of this.activeExecutions.entries()) {
      if (execution.agent_uuid === agent_uuid) {
        this.broadcastScriptOutput(execution_id, {
          type: 'script_error',
          execution_id,
          error: 'Agent disconnected',
          timestamp: Date.now()
        });
        this.activeExecutions.delete(execution_id);
      }
    }
  }
 
  /**
   * Handle RDS connections (existing functionality).
   * Routes viewer and control WebSocket connections for remote desktop sessions.
   * @param {WebSocket} ws - WebSocket connection
   * @param {object} req - HTTP request object with URL
   * @returns {void}
   */
  handleRDSConnection(ws, req) {
    const parsed = url.parse(req.url, true);
    const path = parsed.pathname;
 
    if (path.startsWith('/api/rds/view/')) {
      const sessionId = parsed.query.session;
      if (!sessionId) return ws.close();
 
      const session = this.rdsSessions.get(sessionId);
      if (session) session.viewer = ws;
 
      console.log(`[RDS] Viewer connected: ${sessionId}`);
 
      // Handle messages from viewer (keyboard/mouse) - forward to agent via main WebSocket
      ws.on('message', (data) => {
        if (session && session.agentId) {
          const agentWs = this.agentConnections.get(session.agentId);
          if (agentWs && agentWs.readyState === 1) {
            // Send viewer data to agent through main WebSocket
            agentWs.send(JSON.stringify({
              type: 'rds_data',
              sessionId: sessionId,
              data: data.toString('base64') // Encode binary data as base64
            }));
          }
        }
      });
 
      ws.on('close', () => {
        if (session) {
          session.viewer = null;
          // Notify agent to close tunnel
          const agentWs = this.agentConnections.get(session.agentId);
          if (agentWs && agentWs.readyState === 1) {
            agentWs.send(JSON.stringify({
              type: 'rds_stop',
              sessionId: sessionId
            }));
          }
        }
        console.log(`[RDS] Viewer disconnected: ${sessionId}`);
      });
 
    } else if (path.startsWith('/api/rds/control/')) {
      const sessionId = parsed.query.session;
      if (!sessionId) return ws.close();
 
      const session = this.rdsSessions.get(sessionId);
      if (!session) return ws.close();
 
      session.control = ws;
      console.log(`[RDS] Control connected: ${sessionId}`);
 
      // Forward control messages to agent
      ws.on('message', (raw) => {
        try {
          const msg = JSON.parse(raw.toString());
          const agentWs = this.agentConnections.get(session.agentId);
          if (agentWs && agentWs.readyState === 1) {
            agentWs.send(JSON.stringify({ ...msg, sessionId }));
          }
        } catch (err) {
          console.error('[RDS] Error forwarding control:', err);
        }
      });
 
      ws.on('close', () => {
        if (session) session.control = null;
        console.log(`[RDS] Control disconnected: ${sessionId}`);
      });
    }
  }
 
  /**
   * Handle persistent tunnel request from agent.
   * Assigns port and creates persistent tunnel session.
   * @param {string} agent_uuid - Agent UUID
   * @param {object} msg - Message with protocol and port
   * @param {WebSocket} ws - Agent WebSocket connection
   * @returns {Promise<void>}
   */
  async handlePersistentTunnelRequest(agent_uuid, msg, ws) {
    const { protocol, port: localPort } = msg;
    console.log(`[Persistent Tunnel] Agent ${agent_uuid} requesting ${protocol}:${localPort} tunnel`);
 
    // Clean up any existing tunnel for this agent
    const existing = this.persistentTunnels.get(agent_uuid);
    if (existing) {
      this.assignedPorts.delete(existing.assignedPort);
      this.rdsSessions.delete(existing.sessionId);
      console.log(`[Persistent Tunnel] Cleaned up existing tunnel for agent ${agent_uuid}`);
    }
 
    // Assign a new port
    const assignedPort = this.getNextAvailablePort();
    const sessionId = require('crypto').randomUUID();
 
    // Store tunnel info
    this.persistentTunnels.set(agent_uuid, {
      protocol,
      localPort,
      assignedPort,
      sessionId,
      ready: false
    });
 
    // Register in RDS sessions for data routing
    this.rdsSessions.set(sessionId, {
      agentId: agent_uuid,
      viewer: null,
      agent: null,
      control: null,
      persistent: true
    });
 
    console.log(`[Persistent Tunnel] Assigned ${agent_uuid} → localhost:${assignedPort} (${protocol}:${localPort})`);
 
    // Send assignment to agent
    ws.send(JSON.stringify({
      type: 'persistent_tunnel_assigned',
      sessionId,
      assignedPort,
      protocol,
      localPort,
      timestamp: Date.now()
    }));
  }
 
  /**
   * Handle persistent tunnel ready confirmation.
   * Marks tunnel as ready after agent confirms initialization.
   * @param {string} agent_uuid - Agent UUID
   * @param {object} msg - Message with sessionId
   * @returns {void}
   */
  handlePersistentTunnelReady(agent_uuid, msg) {
    const { sessionId } = msg;
    const tunnel = this.persistentTunnels.get(agent_uuid);
    
    if (tunnel && tunnel.sessionId === sessionId) {
      tunnel.ready = true;
      console.log(`[Persistent Tunnel] ✅ Agent ${agent_uuid} tunnel ready: localhost:${tunnel.assignedPort} → ${tunnel.protocol}:${tunnel.localPort}`);
    }
  }
 
  /**
   * Get next available port for persistent tunnels.
   * Increments port counter and tracks assigned ports.
   * @returns {number} Available port number
   */
  getNextAvailablePort() {
    while (this.assignedPorts.has(this.nextAvailablePort)) {
      this.nextAvailablePort++;
    }
    const port = this.nextAvailablePort++;
    this.assignedPorts.add(port);
    return port;
  }
 
  /**
   * Get persistent tunnel info for agent.
   * @param {string} agent_uuid - Agent UUID
   * @returns {object|undefined} Tunnel info object or undefined if not found
   */
  getPersistentTunnel(agent_uuid) {
    return this.persistentTunnels.get(agent_uuid);
  }
 
  /**
   * Get agent connection status.
   * @param {string} agent_uuid - Agent UUID
   * @returns {boolean} True if agent is connected and ready
   */
  isAgentConnected(agent_uuid) {
    const ws = this.agentConnections.get(agent_uuid);
    return ws && ws.readyState === 1;
  }
 
  /**
   * Get stats.
   * Returns current connection and session counts.
   * @returns {object} Stats object with connection counts
   */
  getStats() {
    return {
      agents: this.agentConnections.size,
      dashboards: this.dashboardConnections.size,
      subscriptions: this.subscriptions.size,
      executions: this.activeExecutions.size,
      rds_sessions: this.rdsSessions.size,
      persistent_tunnels: this.persistentTunnels.size
    };
  }
 
  /**
   * Handle MeshCentral terminal WebSocket connection.
   * Proxies terminal session between dashboard and MeshCentral.
   * @param {WebSocket} ws - WebSocket connection from dashboard
   * @param {object} query - Query parameters with sessionToken and agentId
   * @returns {Promise<void>}
   */
  async handleMeshCentralTerminal(ws, query) {
    const { sessionToken, agentId } = query;
    
    console.log(`[MeshCentral Terminal] New connection: agentId=${agentId}, sessionToken=${sessionToken?.substring(0, 8)}...`);
 
    if (!sessionToken || !agentId) {
      console.error('[MeshCentral Terminal] Missing sessionToken or agentId');
      ws.close(1008, 'Missing required parameters');
      return;
    }
 
    try {
      // Validate session and get device info
      const sessionResult = await pool.query(
        `SELECT s.*, a.meshcentral_nodeid, a.tenant_id 
         FROM meshcentral_sessions s
         JOIN agents a ON a.agent_id = s.agent_id
         WHERE s.session_token = $1 AND s.agent_id = $2 AND s.expires_at > NOW()`,
        [sessionToken, agentId]
      );
 
      if (sessionResult.rows.length === 0) {
        console.error('[MeshCentral Terminal] Invalid or expired session');
        ws.close(1008, 'Invalid or expired session');
        return;
      }
 
      const session = sessionResult.rows[0];
      const nodeId = session.meshcentral_nodeid;
      const tenantId = session.tenant_id;
 
      console.log(`[MeshCentral Terminal] Session validated: tenantId=${tenantId}, nodeId=${nodeId}`);
 
      // Get tenant's MeshCentral API connection
      const { getTenantMeshAPI } = require('../routes/meshcentral');
      const api = await getTenantMeshAPI(tenantId);
      if (!api || !api.ws) {
        console.error('[MeshCentral Terminal] Failed to connect to MeshCentral');
        ws.close(1011, 'MeshCentral connection failed');
        return;
      }
 
      console.log(`[MeshCentral Terminal] Connected to MeshCentral for tenant ${tenantId}`);
 
      // Forward messages from MeshCentral to client
      const meshMessageHandler = (data) => {
        try {
          const message = JSON.parse(data.toString());
          if (message.type === 'console') {
            console.log(`[MeshCentral Terminal] -> Client: console message`);
            ws.send(JSON.stringify(message));
          }
        } catch (err) {
          console.error('[MeshCentral Terminal] Error parsing MeshCentral message:', err);
        }
      };
      api.ws.on('message', meshMessageHandler);
 
      // Forward messages from client to MeshCentral
      ws.on('message', (data) => {
        try {
          const clientMessage = JSON.parse(data.toString());
          console.log(`[MeshCentral Terminal] <- Client:`, clientMessage.type, clientMessage.action);
 
          const meshMessage = {
            action: 'msg',
            type: 'console',
            nodeid: nodeId,
            msg: JSON.stringify(clientMessage)
          };
 
          api.ws.send(JSON.stringify(meshMessage));
        } catch (err) {
          console.error('[MeshCentral Terminal] Error forwarding client message:', err);
        }
      });
 
      // Handle disconnection
      ws.on('close', () => {
        console.log(`[MeshCentral Terminal] Client disconnected: agentId=${agentId}`);
        api.ws.off('message', meshMessageHandler);
      });
 
      // Send initial connection success
      ws.send(JSON.stringify({
        type: 'connected',
        nodeId: nodeId
      }));
 
    } catch (err) {
      console.error('[MeshCentral Terminal] Connection error:', err);
      ws.close(1011, 'Internal server error');
    }
  }
 
  /**
   * Handle MeshCentral file browser WebSocket connection.
   * Proxies file browser session between dashboard and MeshCentral.
   * @param {WebSocket} ws - WebSocket connection from dashboard
   * @param {object} query - Query parameters with sessionToken and agentId
   * @returns {Promise<void>}
   */
  async handleMeshCentralFiles(ws, query) {
    const { sessionToken, agentId } = query;
    
    console.log(`[MeshCentral Files] New connection: agentId=${agentId}, sessionToken=${sessionToken?.substring(0, 8)}...`);
 
    if (!sessionToken || !agentId) {
      console.error('[MeshCentral Files] Missing sessionToken or agentId');
      ws.close(1008, 'Missing required parameters');
      return;
    }
 
    try {
      // Validate session and get device info
      const sessionResult = await pool.query(
        `SELECT s.*, a.meshcentral_nodeid, a.tenant_id 
         FROM meshcentral_sessions s
         JOIN agents a ON a.agent_id = s.agent_id
         WHERE s.session_token = $1 AND s.agent_id = $2 AND s.expires_at > NOW()`,
        [sessionToken, agentId]
      );
 
      if (sessionResult.rows.length === 0) {
        console.error('[MeshCentral Files] Invalid or expired session');
        ws.close(1008, 'Invalid or expired session');
        return;
      }
 
      const session = sessionResult.rows[0];
      const nodeId = session.meshcentral_nodeid;
      const tenantId = session.tenant_id;
 
      console.log(`[MeshCentral Files] Session validated: tenantId=${tenantId}, nodeId=${nodeId}`);
 
      // Get tenant's MeshCentral API connection
      const { getTenantMeshAPI } = require('../routes/meshcentral');
      const api = await getTenantMeshAPI(tenantId);
      if (!api || !api.ws) {
        console.error('[MeshCentral Files] Failed to connect to MeshCentral');
        ws.close(1011, 'MeshCentral connection failed');
        return;
      }
 
      console.log(`[MeshCentral Files] Connected to MeshCentral for tenant ${tenantId}`);
 
      // Forward messages from MeshCentral to client
      const meshMessageHandler = (data) => {
        try {
          const message = JSON.parse(data.toString());
          
          // Only forward file-related messages
          if (message.type === 'notify' || 
              message.type === 'download' ||
              message.type === 'uploaddone' ||
              message.action === 'ls' ||
              message.action === 'download' ||
              message.action === 'upload') {
            
            console.log(`[MeshCentral Files] -> Client: ${message.action || message.type}`);
            ws.send(JSON.stringify(message));
          }
        } catch (err) {
          console.error('[MeshCentral Files] Error parsing MeshCentral message:', err);
        }
      };
      api.ws.on('message', meshMessageHandler);
 
      // Forward messages from client to MeshCentral
      ws.on('message', (data) => {
        try {
          const clientMessage = JSON.parse(data.toString());
          console.log(`[MeshCentral Files] <- Client: ${clientMessage.action}`, {
            path: clientMessage.path,
            reqid: clientMessage.reqid
          });
 
          const meshMessage = {
            action: 'msg',
            type: 'files',
            nodeid: nodeId,
            ...clientMessage
          };
 
          api.ws.send(JSON.stringify(meshMessage));
        } catch (err) {
          console.error('[MeshCentral Files] Error forwarding client message:', err);
          ws.send(JSON.stringify({
            type: 'error',
            message: 'Failed to process request'
          }));
        }
      });
 
      // Handle disconnection
      ws.on('close', () => {
        console.log(`[MeshCentral Files] Client disconnected: agentId=${agentId}`);
        api.ws.off('message', meshMessageHandler);
      });
 
      // Send initial connection success
      ws.send(JSON.stringify({
        type: 'connected',
        nodeId: nodeId
      }));
 
    } catch (err) {
      console.error('[MeshCentral Files] Connection error:', err);
      ws.close(1011, 'Internal server error');
    }
  }
}
 
// Create singleton instance
const wsManager = new WebSocketManager();
 
module.exports = wsManager;