4bnode/add-websocket.js

A professional tool to generate a Node.js app by 4Brains Technologies

#!/usr/bin/env node

import fs from "fs";
import path from "path";
import readline from "readline";
import { execSync } from "child_process";
import chalk from "chalk";

const rl = readline.createInterface({
  input: process.stdin,
  output: process.stdout,
});

function updateEnvFile(filePath, wsPort) {
  if (!wsPort) return; // Don't update .env if using the same port as the server

  const wsPortLine = `WS_PORT=${wsPort}\n`;

  if (fs.existsSync(filePath)) {
    let envContent = fs.readFileSync(filePath, "utf8");
    if (envContent.includes("WS_PORT=")) {
      envContent = envContent.replace(/WS_PORT=.*/, wsPortLine.trim());
    } else {
      envContent = envContent.trim() + "\n" + wsPortLine;
    }
    fs.writeFileSync(filePath, envContent, "utf8");
  } else {
    fs.writeFileSync(filePath, wsPortLine, "utf8");
  }
}

function addWebSocket(samePort, wsPort = null) {
  const projectRoot = process.cwd();
  const indexPath = path.join(projectRoot, "index.js");

  let indexContent = fs.readFileSync(indexPath, "utf8");

  const wsImport = `import { WebSocketServer } from 'ws';\nimport http from 'http';\n`;

  // Define WebSocket setup based on `samePort`
  const websocketSetup = samePort
    ? `
const server = http.createServer(app);
const wsServer = new WebSocketServer({ server });
console.log(\`WebSocket server is running on the same port as the server (\${port})\`);
`
    : `
const wsPort = process.env.WS_PORT || ${wsPort};
const wsServer = new WebSocketServer({ port: wsPort });
console.log(\`WebSocket server is running on WebSocket port \${wsPort}\`);
`;

  const wsLogic = `
wsServer.on('connection', (ws) => {
  console.log('A client connected');

  ws.on('message', (message) => {
    try {
      const parsedMessage = JSON.parse(message);
      console.log('Received:', parsedMessage);
      
      // Broadcast to all clients except sender
      wsServer.clients.forEach((client) => {
        if (client !== ws && client.readyState === ws.OPEN) {
          client.send(JSON.stringify(parsedMessage));
        }
      });
    } catch (error) {
      console.error('Invalid message format:', error);
    }
  });

  ws.on('close', () => {
    console.log('Client disconnected');
  });
});

app.use((req, res, next) => {
  req.wss = wsServer;
  next();
});
`;

  // Select correct `listen()` method
  const listenSetup = samePort
    ? `
server.listen(port, () => {
  const localIp = getLocalNetworkIp();
  console.log(\`App is running on http://localhost:\${port}\`);
  if (localIp) {
    console.log(
      \`App is also accessible on your local network at http://\${localIp}:\${port}\`
    );
  }
});
`
    : `
app.listen(port, () => {
  const localIp = getLocalNetworkIp();
  console.log(\`App is running on http://localhost:\${port}\`);
  if (localIp) {
    console.log(
      \`App is also accessible on your local network at http://\${localIp}:\${port}\`
    );
  }
});
`;

  // Add imports if they don't already exist
  if (!indexContent.includes(wsImport)) {
    indexContent = wsImport + indexContent;
  }

  // Remove existing `app.listen()` or `server.listen()`
  indexContent = indexContent.replace(/server\.listen\(port.*\);/s, "");
  indexContent = indexContent.replace(/app\.listen\(port.*\);/s, "");

  // Ensure WebSocket setup is placed before routes
  if (!indexContent.includes("const wsServer = new WebSocketServer({")) {
    const routesStart = indexContent.indexOf("app.use(");
    if (routesStart !== -1) {
      indexContent =
        indexContent.slice(0, routesStart) +
        websocketSetup +
        wsLogic +
        indexContent.slice(routesStart);
    } else {
      indexContent += websocketSetup + wsLogic;
    }
  }

  // Append the correct `listen()` function
  if (
    !indexContent.includes("app.listen(port)") &&
    !indexContent.includes("server.listen(port)")
  ) {
    indexContent += listenSetup;
  }

  fs.writeFileSync(indexPath, indexContent, "utf8");

  // Update `.env` files only if using a separate WebSocket port
  if (!samePort) {
    updateEnvFile(path.join(projectRoot, ".env.development"), wsPort);
    updateEnvFile(path.join(projectRoot, ".env.production"), wsPort);
  }

  // Install ws library
  console.log(chalk.green("Installing ws..."));
  execSync("npm install ws", { stdio: "inherit" });

  console.log(chalk.green("WebSocket integration added successfully."));
}

// Get the server port from environment variables, default to 3000 if undefined
const serverPort = process.env.PORT || 3000;

function askWebSocketPort() {
  rl.question(
    chalk.yellow("Enter a different WebSocket port (default: 3001): "),
    (wsPortInput) => {
      let wsPort = parseInt(wsPortInput.trim()) || 3001;

      if (wsPort === serverPort) {
        console.log(
          chalk.red("WebSocket port cannot be the same as the server port.")
        );
        askWebSocketPort(); // Re-prompt if user enters the same port
      } else {
        addWebSocket(false, wsPort); // samePort = false, pass wsPort
        rl.close();
      }
    }
  );
}

function askUseSamePort() {
  rl.question(
    chalk.yellow(
      "Do you want to use the server port for WebSocket? (Yes/No): "
    ),
    (useSamePort) => {
      const answer = useSamePort.trim().toLowerCase();

      if (answer === "y" || answer === "yes") {
        addWebSocket(true); // samePort = true, no wsPort needed
        rl.close();
      } else if (answer === "n" || answer === "no") {
        askWebSocketPort(); // Ask for a different WebSocket port
      } else {
        console.log(chalk.red("Invalid input. Please enter Yes or No."));
        askUseSamePort(); // Re-prompt if invalid input
      }
    }
  );
}

// Start asking user input
askUseSamePort();