@oussema_mili/test-pkg-123 1.1.32

package/websocket-server.js

@@ -0,0 +1,415 @@
+import { WebSocketServer } from "ws";
+import { v4 as uuidv4 } from "uuid";
+import crypto from "crypto";
+import fs from "fs";
+import path from "path";
+import chalk from "chalk";
+import os from "os";
+import { loadConfig } from "./store/configStore.js";
+
+// Import handlers
+import containerHandlers from "./docker-actions/containers.js";
+import imageHandlers from "./docker-actions/images.js";
+import volumeHandlers from "./docker-actions/volumes.js";
+import appHandlers from "./docker-actions/apps.js";
+import registryHandlers from "./docker-actions/registry.js";
+import logHandlers from "./docker-actions/logs.js";
+import metricsHandlers from "./docker-actions/metrics.js";
+import terminalHandlers from "./docker-actions/terminal.js";
+import generalHandlers from "./docker-actions/general.js";
+import setupTaskHandlers from "./docker-actions/setup-tasks.js";
+import agentSessionStore from "./store/agentSessionStore.js";
+
+// Store active streams and sessions
+const activeStreams = new Map();
+const terminalSessions = new Map();
+
+// Load configuration
+const config = loadConfig();
+const AGENT_ROOT_DIR = path.join(os.homedir(), config.agentRootDir);
+const WS_TOKEN_FILE = path.join(AGENT_ROOT_DIR, "ws-token");
+const AUTH_TIMEOUT_MS = 30000; // 30 seconds to authenticate
+
+// Heartbeat settings
+const HEARTBEAT_INTERVAL_MS = 30000; // 30 seconds
+const HEARTBEAT_TIMEOUT_MS = 35000; // 35 seconds to respond
+
+/**
+ * Generate a secure random WebSocket authentication token
+ * @returns {string} 64-character hex token
+ */
+function generateWsToken() {
+  return crypto.randomBytes(32).toString("hex");
+}
+
+/**
+ * Save WebSocket token to file for local-env app to read
+ * @param {string} token - The WebSocket authentication token
+ */
+function saveWsToken(token) {
+  // Ensure directory exists
+  if (!fs.existsSync(AGENT_ROOT_DIR)) {
+    fs.mkdirSync(AGENT_ROOT_DIR, { recursive: true, mode: 0o700 });
+  }
+
+  // Write token to file with restricted permissions (owner only)
+  fs.writeFileSync(WS_TOKEN_FILE, token, { mode: 0o600 });
+}
+
+/**
+ * Read WebSocket token from file
+ * @returns {string|null} The WebSocket token or null if not found
+ */
+function readWsToken() {
+  try {
+    if (fs.existsSync(WS_TOKEN_FILE)) {
+      return fs.readFileSync(WS_TOKEN_FILE, "utf-8").trim();
+    }
+  } catch (error) {
+    console.error("Error reading WS token:", error.message);
+  }
+  return null;
+}
+
+function setupWebSocketServer(
+  server,
+  clients,
+  agentId,
+  sessionInfo = {},
+  existingToken = null
+) {
+  // Store agent session info for use in handlers
+  agentSessionStore.setAgentSessionInfo({
+    agentId: agentId,
+    userEntityRef: sessionInfo.userEntityRef || null,
+    sessionExpiresAt: sessionInfo.expiresAt || null,
+  });
+
+  // Use existing token or generate a new one
+  const wsToken = existingToken || generateWsToken();
+  if (!existingToken) {
+    // Only save if we generated a new token
+    saveWsToken(wsToken);
+  }
+
+  const wss = new WebSocketServer({ server });
+
+  // Set up cleanup timer for inactive terminal sessions
+  setInterval(() => {
+    terminalHandlers.cleanupInactiveTerminalSessions(terminalSessions);
+  }, 15 * 60 * 1000); // Run every 15 minutes
+
+  // Setup heartbeat interval to check client health
+  const heartbeatInterval = setInterval(() => {
+    wss.clients.forEach((ws) => {
+      if (ws.isAlive === false) {
+        // Client didn't respond to ping, terminate
+        console.log(chalk.yellow(`Terminating unresponsive client`));
+        return ws.terminate();
+      }
+
+      ws.isAlive = false;
+      ws.ping();
+    });
+  }, HEARTBEAT_INTERVAL_MS);
+
+  // Clean up heartbeat interval when server closes
+  wss.on("close", () => {
+    clearInterval(heartbeatInterval);
+  });
+
+  wss.on("connection", (ws) => {
+    // Generate a unique client ID for each WebSocket connection
+    const clientId = uuidv4();
+    const connectionTime = Date.now();
+    let authenticated = false;
+
+    // Initialize heartbeat tracking
+    ws.isAlive = true;
+
+    // Handle pong responses (heartbeat acknowledgment)
+    ws.on("pong", () => {
+      ws.isAlive = true;
+    });
+
+    // Set authentication timeout
+    const authTimeout = setTimeout(() => {
+      if (!authenticated) {
+        ws.send(
+          JSON.stringify({
+            type: "auth_error",
+            error: "Authentication timeout. Connection closed.",
+          })
+        );
+        ws.close(4001, "Authentication timeout");
+      }
+    }, AUTH_TIMEOUT_MS);
+
+    // Send authentication challenge
+    ws.send(
+      JSON.stringify({
+        type: "auth_required",
+        message:
+          "Authentication required. Please send auth message with token.",
+        clientId,
+      })
+    );
+
+    // Handle messages from client
+    ws.on("message", async (message) => {
+      let data;
+      try {
+        data = JSON.parse(message);
+      } catch (parseError) {
+        console.error("Error parsing message: ", parseError);
+        ws.send(
+          JSON.stringify({
+            type: "error",
+            error: "Failed to parse message: " + parseError.message,
+          })
+        );
+        return;
+      }
+
+      // Handle authentication
+      if (!authenticated) {
+        if (data.type === "auth" && data.token) {
+          if (data.token === wsToken) {
+            // Authentication successful
+            authenticated = true;
+            clearTimeout(authTimeout);
+
+            // Store client info
+            clients.set(clientId, {
+              ws,
+              connectionTime,
+              lastActivity: Date.now(),
+              isActive: true,
+              authenticated: true,
+            });
+
+            console.log(`✅ Client authenticated: ${clientId}`);
+
+            // Send success response
+            ws.send(
+              JSON.stringify({
+                type: "auth_success",
+                clientId,
+                agentId: agentId,
+                message: "Successfully authenticated with Fenwave Agent",
+              })
+            );
+          } else {
+            // Invalid token
+            console.log(
+              chalk.red(
+                `❌ Invalid authentication token from client: ${clientId}`
+              )
+            );
+            ws.send(
+              JSON.stringify({
+                type: "auth_error",
+                error: "Invalid authentication token",
+              })
+            );
+            ws.close(4003, "Invalid authentication token");
+          }
+        } else {
+          // Not an auth message
+          ws.send(
+            JSON.stringify({
+              type: "auth_error",
+              error: "Authentication required before sending messages",
+            })
+          );
+        }
+        return;
+      }
+
+      // Update last activity time for authenticated clients
+      const client = clients.get(clientId);
+      if (client) {
+        client.lastActivity = Date.now();
+      }
+
+      // Only log received messages in debug mode
+      if (process.env.DEBUG === 'true' || process.env.FW_VERBOSE === 'true') {
+        console.log(`📨 Received message from client ${clientId}:`, data.action);
+      }
+
+      // Route messages to appropriate handlers
+      await routeMessage(ws, clientId, data);
+    });
+
+    // Handle client disconnect
+    ws.on("close", () => {
+      console.log(`🔴 Client disconnected: ${clientId}`);
+      clients.delete(clientId);
+
+      // Clean up any active streams for this client
+      for (const [streamId, stream] of activeStreams.entries()) {
+        if (streamId.startsWith(clientId)) {
+          if (stream.destroy) stream.destroy();
+          activeStreams.delete(streamId);
+        }
+      }
+    });
+  });
+
+  return wss;
+}
+
+async function routeMessage(ws, clientId, data) {
+  const { action, payload = {} } = data;
+
+  // Container actions
+  if (
+    action === "fetchContainers" ||
+    action === "startContainer" ||
+    action === "stopContainer" ||
+    action === "restartContainer" ||
+    action === "deleteContainer" ||
+    action === "createContainer" ||
+    action === "inspectContainer"
+  ) {
+    return await containerHandlers.handleContainerAction(ws, action, payload);
+  }
+
+  // Image actions
+  if (
+    action === "fetchImages" ||
+    action === "pullImage" ||
+    action === "pushImage" ||
+    action === "deleteImage" ||
+    action === "tagImage" ||
+    action === "inspectImage"
+  ) {
+    return await imageHandlers.handleImageAction(ws, action, payload);
+  }
+
+  // Volume actions
+  if (
+    action === "fetchVolumes" ||
+    action === "createVolume" ||
+    action === "deleteVolume" ||
+    action === "inspectVolume" ||
+    action === "getContainersUsingVolume" ||
+    action === "fetchImages" ||
+    action === "pullImage" ||
+    action === "backupVolume" ||
+    action === "listVolumeBackups" ||
+    action === "fetchImages" ||
+    action === "pullImage" ||
+    action === "restoreVolumeFromBackup" ||
+    action === "deleteVolumeBackup" ||
+    action === "getBackupDownloadUrl"
+  ) {
+    return await volumeHandlers.handleVolumeAction(ws, action, payload);
+  }
+
+  // App actions
+  if (
+    action === "fetchApps" ||
+    action === "fetchAppVersions" ||
+    action === "startApp" ||
+    action === "stopApp" ||
+    action === "restartApp" ||
+    action === "deleteApp" ||
+    action === "deleteAppVersions" ||
+    action === "createApp" ||
+    action === "validateCompose" ||
+    action === "syncApp" ||
+    action === "changeVersion" ||
+    action === "cleanApp" ||
+    action === "fetchAppLogs" ||
+    action === "fetchLogContent"
+  ) {
+    return await appHandlers.handleAppAction(ws, action, payload);
+  }
+
+  // Registry actions
+  if (
+    action === "fetchRegistries" ||
+    action === "connectRegistry" ||
+    action === "disconnectRegistry" ||
+    action === "renameRegistry" ||
+    action === "fetchRegistryImages" ||
+    action === "setActiveRegistry"
+  ) {
+    return await registryHandlers.handleRegistryAction(ws, action, payload);
+  }
+
+  // Log actions
+  if (
+    action === "fetchContainerLogs" ||
+    action === "streamContainerLogs" ||
+    action === "stopStreamLogs"
+  ) {
+    return await logHandlers.handleLogAction(
+      ws,
+      clientId,
+      action,
+      payload,
+      activeStreams
+    );
+  }
+
+  // Metrics actions
+  if (
+    action === "fetchSystemMetrics" ||
+    action === "fetchContainerMetrics" ||
+    action === "systemMetrics"
+  ) {
+    return await metricsHandlers.handleMetricsAction(ws, action, payload);
+  }
+
+  // Terminal actions
+  if (
+    action === "openTerminalSession" ||
+    action === "closeTerminalSession" ||
+    action === "sendTerminalInput" ||
+    action === "resizeTerminal"
+  ) {
+    return await terminalHandlers.handleTerminalAction(
+      ws,
+      clientId,
+      action,
+      payload,
+      terminalSessions
+    );
+  }
+
+  // General actions
+  if (["getAgentInfo", "getDockerInfo", "execInContainer"].includes(action)) {
+    return await generalHandlers.handleGeneralAction(
+      ws,
+      clientId,
+      action,
+      payload
+    );
+  }
+
+  // Setup task actions
+  if (
+    action === "getSetupTasks" ||
+    action === "getSetupStatus" ||
+    action === "verifySetupTask" ||
+    action === "executeSetupTask" ||
+    action === "saveSetupProgress" ||
+    action === "clearSetupProgress" ||
+    action === "browseDirectory"
+  ) {
+    return await setupTaskHandlers.handleSetupAction(ws, action, payload);
+  }
+
+  // Unknown action
+  ws.send(
+    JSON.stringify({
+      type: "error",
+      requestId: payload.requestId,
+      error: `Unknown action: ${action}`,
+    })
+  );
+}
+
+export { setupWebSocketServer, generateWsToken, saveWsToken, readWsToken };