server.js

// server.js
const express = require('express');
const http = require('http');
const bcrypt = require('bcrypt');
const session = require('express-session');
const fs = require('fs');
const path = require('path');
const { v4: uuidv4 } = require('uuid');
const { WebSocketServer } = require('ws'); // Use 'ws' library for WebSockets

const app = express();
const server = http.createServer(app);
const wss = new WebSocketServer({ noServer: true }); // Create WebSocket server without a port

// rate limiting: track user -> timestamps of sent messages
const messageRateLimit = {};
const MESSAGE_LIMIT = 5;
const TIME_WINDOW_MS = 1000;
const MAX_MESSAGE_LENGTH = 10000;

const PORT = 3000;
app.use((req, res, next) => {
  const origin = req.headers.origin;
  if (origin && (origin.startsWith('https://') || origin === 'null')) {
    res.setHeader('Access-Control-Allow-Origin', origin);
  }
  res.setHeader('Access-Control-Allow-Credentials', 'true');
  res.setHeader('Access-Control-Allow-Headers', 'Content-Type');
  res.setHeader('Access-Control-Allow-Methods', 'GET, POST, OPTIONS');
  if (req.method === 'OPTIONS') {
    return res.sendStatus(200);
  }
  next();
});
app.use((req, res, next) => {
  res.setHeader('X-Frame-Options', 'ALLOWALL');
  next();
});

const USERS_DB = path.join(__dirname, 'users.json');
const CHATS_DB = path.join(__dirname, 'chats.json');
const MESSAGES_DB = path.join(__dirname, 'messages.json');

app.use(express.urlencoded({ extended: true }));
app.use(express.json());

const sessionMiddleware = session({
  genid: () => uuidv4(),
  secret: 'super-secret-key',
  resave: false,
  saveUninitialized: false,
  cookie: {
    httpOnly: true,
    // Setting secure to false and sameSite to 'none' is the most permissive config.
    // This is required for local file access to work with secure remote servers.
    secure: false, 
    sameSite: 'none',
    maxAge: 1000 * 60 * 60 * 2
  }
});

app.use(sessionMiddleware);

function loadJSON(file) {
  return fs.existsSync(file) ? JSON.parse(fs.readFileSync(file, 'utf-8')) : [];
}

function saveJSON(file, data) {
  fs.writeFileSync(file, JSON.stringify(data, null, 2));
}

function generateSuffix() {
  return Math.floor(1000 + Math.random() * 9000);
}

function normalizeUsername(name) {
  return name.trim().toLowerCase();
}

function isAuthenticated(req, res, next) {
  if (req.session && req.session.username) return next();
  res.status(401).send('Unauthorized');
}

// Map to store connected clients by username
const clients = new Map();

// HTTP server listens for upgrade events to handle WebSocket connections
server.on('upgrade', (request, socket, head) => {
  sessionMiddleware(request, {}, () => {
    if (!request.session || !request.session.username) {
      socket.write('HTTP/1.1 401 Unauthorized\r\n\r\n');
      socket.destroy();
      return;
    }
    wss.handleUpgrade(request, socket, head, ws => {
      wss.emit('connection', ws, request);
    });
  });
});

wss.on('connection', (ws, request) => {
  const username = request.session.username;
  ws.username = username; // Store username on the WebSocket object
  clients.set(username, ws);

  console.log(`User ${username} connected via WebSocket.`);

  // Send a message to the new client to signal a successful connection
  ws.send(JSON.stringify({ type: 'connected', message: 'Connected to WebSocket server.' }));

  ws.on('message', async message => {
    try {
      const data = JSON.parse(message);
      const { type } = data;

      switch (type) {
        case 'refreshChats':
          // Re-send the chatUpdated event to the client to trigger a refresh
          ws.send(JSON.stringify({ type: 'chatsUpdated' }));
          break;
        case 'sendMessage':
          // Re-use the existing sendMessage logic from the HTTP POST route
          // This simulates the client-side logic from the old code
          // and adapts it for the WebSocket message format
          const { chatId, encryptedContent, iv } = data;

          if (!encryptedContent || !iv || encryptedContent.length > MAX_MESSAGE_LENGTH) {
            return ws.send(JSON.stringify({ type: 'error', message: 'Encrypted message content or IV missing/too long.' }));
          }

          const now = Date.now();
          if (!messageRateLimit[username]) {
            messageRateLimit[username] = [];
          }
          messageRateLimit[username] = messageRateLimit[username].filter(ts => now - ts < TIME_WINDOW_MS);

          if (messageRateLimit[username].length >= MESSAGE_LIMIT) {
            return ws.send(JSON.stringify({ type: 'error', message: 'Rate limit exceeded' }));
          }

          messageRateLimit[username].push(now);

          const chats = loadJSON(CHATS_DB);
          const chat = chats.find(c => c.chatId === chatId);
          if (!chat || !chat.members.includes(username)) {
            return ws.send(JSON.stringify({ type: 'error', message: 'Not allowed to send message to this chat.' }));
          }

          const messages = loadJSON(MESSAGES_DB);
          messages.push({
            chatId,
            from: username,
            iv: iv,
            content: encryptedContent
          });
          saveJSON(MESSAGES_DB, messages);

          // Notify all chat members using the client map
          chat.members.forEach(member => {
            const memberWs = clients.get(member);
            if (memberWs) {
              memberWs.send(JSON.stringify({
                type: 'newMessage',
                chatId,
                from: username,
                iv,
                content: encryptedContent
              }));
            }
          });
          break;
        default:
          console.warn(`Unknown message type: ${type}`);
      }
    } catch (e) {
      console.error('Failed to parse WebSocket message:', e);
      ws.send(JSON.stringify({ type: 'error', message: 'Invalid message format.' }));
    }
  });

  ws.on('close', () => {
    console.log(`User ${username} disconnected.`);
    clients.delete(username);
  });

  ws.on('error', error => {
    console.error(`WebSocket error for user ${username}:`, error);
  });
});

// All API routes remain unchanged as they are standard HTTP POST/GET requests.
app.post('/signup', async (req, res) => {
  const { username, password, publicKey } = req.body;
  if (!username || !password || username.length > 200 || password.length > 200) {
    return res.status(400).send('Username and password must be less than 200 characters');
  }
  if (!publicKey) {
    return res.status(400).send('Public key is required for secure signup.');
  }

  const users = loadJSON(USERS_DB);
  const baseName = normalizeUsername(username);

  if (users.find(u => u.baseName === baseName)) {
    return res.status(400).send('Username taken');
  }

  const suffix = generateSuffix();
  const hashedPassword = await bcrypt.hash(password, 10);
  users.push({ baseName, suffix, password: hashedPassword, publicKey });
  saveJSON(USERS_DB, users);
  res.status(200).send('Signup successful.');
});

app.post('/logout', (req, res) => {
  req.session.destroy(() => {
    res.clearCookie('connect.sid');
    res.sendStatus(200);
  });
});

app.post('/login', async (req, res) => {
  const { username, password } = req.body;
  const users = loadJSON(USERS_DB);
  const baseName = normalizeUsername(username);
  const user = users.find(u => u.baseName === baseName);
  if (!user) return res.status(401).send('User not found');
  const match = await bcrypt.compare(password, user.password);
  if (!match) return res.status(401).send('Invalid password');

  req.session.username = user.baseName;
  res.status(200).send('Login successful');
});

app.get('/api/users', isAuthenticated, (req, res) => {
  const users = loadJSON(USERS_DB);
  const filteredUsers = users.filter(u => u.baseName !== req.session.username);
  res.json({
    currentUser: req.session.username,
    users: filteredUsers.map(u => ({ baseName: u.baseName, suffix: u.suffix, display: `${u.baseName}(${u.suffix})` }))
  });
});

app.get('/api/user-public-key/:username', isAuthenticated, (req, res) => {
  const { username } = req.params;
  const users = loadJSON(USERS_DB);
  const user = users.find(u => normalizeUsername(u.baseName) === normalizeUsername(username));
  if (user && user.publicKey) {
    res.json({ publicKey: user.publicKey });
  } else {
    res.status(404).send('User or public key not found');
  }
});

app.post('/api/create-chat', isAuthenticated, (req, res) => {
  const { users, name, encryptedKeys } = req.body;
  const creator = normalizeUsername(req.session.username);
  
  const chatMembers = [...new Set([creator, ...users.map(normalizeUsername)])];

  if (!encryptedKeys || !Array.isArray(encryptedKeys) || encryptedKeys.length === 0) {
    return res.status(400).send('Encrypted keys for chat members are required.');
  }

  const receivedMembers = new Set(encryptedKeys.map(k => k.memberUsername));
  const missingKeys = chatMembers.filter(member => !receivedMembers.has(member));
  if (missingKeys.length > 0) {
    return res.status(400).send(`Missing encrypted keys for members: ${missingKeys.join(', ')}`);
  }

  const chats = loadJSON(CHATS_DB);
  const chatId = uuidv4();
  
  const encryptedKeysMap = encryptedKeys.reduce((acc, current) => {
    acc[current.memberUsername] = current.encryptedAesKeyBase64;
    return acc;
  }, {});

  chats.push({ chatId, name, members: chatMembers, encryptedKeys: encryptedKeysMap });
  saveJSON(CHATS_DB, chats);
  
  // Notify all chat members about the new chat via WebSocket
  chatMembers.forEach(member => {
    const memberWs = clients.get(member);
    if (memberWs) {
      memberWs.send(JSON.stringify({ type: 'chatCreated', chatId, name }));
    }
  });

  res.status(200).send('Chat created');
});

app.get('/api/chats', isAuthenticated, (req, res) => {
  const chats = loadJSON(CHATS_DB);
  const userChats = chats.filter(chat => chat.members.includes(req.session.username));
  res.json(userChats.map(chat => ({ 
    chatId: chat.chatId, 
    name: chat.name || 'Unnamed Chat',
    encryptedAesKeyBase64: chat.encryptedKeys[req.session.username] 
  })));
});

app.get('/api/messages/:chatId', isAuthenticated, (req, res) => {
  const { chatId } = req.params;
  const messages = loadJSON(MESSAGES_DB).filter(m => m.chatId === chatId);
  const chats = loadJSON(CHATS_DB);
  const chat = chats.find(c => c.chatId === chatId);

  if (!chat || !chat.members.includes(req.session.username)) {
    return res.status(403).send('Unauthorized to view messages in this chat.');
  }

  const encryptedAesKeyBase64 = chat.encryptedKeys[req.session.username];
  if (!encryptedAesKeyBase64) {
    return res.status(500).send('Missing encrypted chat key for user.');
  }

  res.json({
    messages: messages.map(m => ({
      from: m.from,
      iv: m.iv,
      content: m.content
    })),
    encryptedAesKeyBase64: encryptedAesKeyBase64
  });
});

// Since client will open a local file, there is no need to serve index.html from here.
// The client will need to connect to the server's domain directly.

server.listen(PORT, () => console.log(`running on http://localhost:${PORT}`));