streamflow/backend/routes/twoFactor.js

const express = require('express');
const router = express.Router();
const speakeasy = require('speakeasy');
const QRCode = require('qrcode');
const crypto = require('crypto');
const { authenticate } = require('../middleware/auth');
const { modifyLimiter, authLimiter, readLimiter } = require('../middleware/rateLimiter');
const { db } = require('../database/db');
const logger = require('../utils/logger');
const { promisify } = require('util');

const dbRun = promisify(db.run.bind(db));
const dbGet = promisify(db.get.bind(db));
const dbAll = promisify(db.all.bind(db));

// Generate 2FA secret and QR code
router.post('/setup', authenticate, modifyLimiter, async (req, res) => {
  try {
    const userId = req.user.userId;

    // Check if user already has 2FA enabled
    const user = await dbGet('SELECT two_factor_enabled FROM users WHERE id = ?', [userId]);
    
    if (user.two_factor_enabled) {
      return res.status(400).json({ error: '2FA is already enabled' });
    }

    // Generate secret with StreamFlow branding
    const secret = speakeasy.generateSecret({
      name: `StreamFlow IPTV:${req.user.username || 'User'}`,
      issuer: 'StreamFlow'
    });

    // Generate QR code with custom options for better readability
    const qrCodeDataURL = await QRCode.toDataURL(secret.otpauth_url, {
      width: 300,
      margin: 2,
      color: {
        dark: '#000000',
        light: '#FFFFFF'
      },
      errorCorrectionLevel: 'H'
    });

    // Store secret temporarily (not enabled yet)
    await dbRun(
      'UPDATE users SET two_factor_secret = ? WHERE id = ?',
      [secret.base32, userId]
    );

    res.json({
      secret: secret.base32,
      qrCode: qrCodeDataURL,
      manualEntryKey: secret.base32
    });
  } catch (error) {
    logger.error('2FA setup error:', error);
    res.status(500).json({ error: 'Failed to setup 2FA' });
  }
});

// Verify and enable 2FA
router.post('/enable', authenticate, authLimiter, async (req, res) => {
  try {
    const { token } = req.body;
    const userId = req.user.userId;

    if (!token) {
      return res.status(400).json({ error: 'Verification token required' });
    }

    // Get user's secret
    const user = await dbGet(
      'SELECT two_factor_secret FROM users WHERE id = ?',
      [userId]
    );

    if (!user.two_factor_secret) {
      return res.status(400).json({ error: '2FA not set up. Call /setup first' });
    }

    // Verify token
    const verified = speakeasy.totp.verify({
      secret: user.two_factor_secret,
      encoding: 'base32',
      token: token,
      window: 2 // Allow 2 time steps before/after
    });

    if (!verified) {
      return res.status(400).json({ error: 'Invalid verification code' });
    }

    // Generate backup codes
    const backupCodes = [];
    for (let i = 0; i < 10; i++) {
      const code = crypto.randomBytes(4).toString('hex').toUpperCase();
      backupCodes.push(code);
      
      await dbRun(
        'INSERT INTO two_factor_backup_codes (user_id, code) VALUES (?, ?)',
        [userId, code]
      );
    }

    // Enable 2FA
    await dbRun(
      'UPDATE users SET two_factor_enabled = 1 WHERE id = ?',
      [userId]
    );

    logger.info(`2FA enabled for user ${userId}`);

    res.json({
      success: true,
      message: '2FA enabled successfully',
      backupCodes: backupCodes
    });
  } catch (error) {
    logger.error('2FA enable error:', error);
    res.status(500).json({ error: 'Failed to enable 2FA' });
  }
});

// Disable 2FA
router.post('/disable', authenticate, authLimiter, async (req, res) => {
  try {
    const { password, token } = req.body;
    const userId = req.user.userId;

    if (!password) {
      return res.status(400).json({ error: 'Password required to disable 2FA' });
    }

    // Verify password
    const bcrypt = require('bcrypt');
    const user = await dbGet('SELECT password, two_factor_secret FROM users WHERE id = ?', [userId]);
    const validPassword = await bcrypt.compare(password, user.password);

    if (!validPassword) {
      return res.status(401).json({ error: 'Invalid password' });
    }

    // Verify 2FA token
    const verified = speakeasy.totp.verify({
      secret: user.two_factor_secret,
      encoding: 'base32',
      token: token,
      window: 2
    });

    if (!verified) {
      return res.status(400).json({ error: 'Invalid 2FA code' });
    }

    // Disable 2FA and remove secret
    await dbRun(
      'UPDATE users SET two_factor_enabled = 0, two_factor_secret = NULL WHERE id = ?',
      [userId]
    );

    // Delete all backup codes
    await dbRun('DELETE FROM two_factor_backup_codes WHERE user_id = ?', [userId]);

    logger.info(`2FA disabled for user ${userId}`);

    res.json({ success: true, message: '2FA disabled successfully' });
  } catch (error) {
    logger.error('2FA disable error:', error);
    res.status(500).json({ error: 'Failed to disable 2FA' });
  }
});

// Verify 2FA token (for login)
router.post('/verify', authLimiter, async (req, res) => {
  try {
    const { userId, token } = req.body;

    if (!userId || !token) {
      return res.status(400).json({ error: 'User ID and token required' });
    }

    // Get user's secret
    const user = await dbGet(
      'SELECT two_factor_secret, two_factor_enabled FROM users WHERE id = ?',
      [userId]
    );

    if (!user || !user.two_factor_enabled) {
      return res.status(400).json({ error: '2FA not enabled for this user' });
    }

    // Check if it's a backup code
    const backupCode = await dbGet(
      'SELECT id FROM two_factor_backup_codes WHERE user_id = ? AND code = ? AND used = 0',
      [userId, token.toUpperCase()]
    );

    if (backupCode) {
      // Mark backup code as used
      await dbRun(
        'UPDATE two_factor_backup_codes SET used = 1, used_at = CURRENT_TIMESTAMP WHERE id = ?',
        [backupCode.id]
      );

      logger.info(`Backup code used for user ${userId}`);
      return res.json({ valid: true, method: 'backup_code' });
    }

    // Verify TOTP token
    const verified = speakeasy.totp.verify({
      secret: user.two_factor_secret,
      encoding: 'base32',
      token: token,
      window: 2
    });

    if (verified) {
      return res.json({ valid: true, method: 'totp' });
    } else {
      return res.status(400).json({ error: 'Invalid 2FA code' });
    }
  } catch (error) {
    logger.error('2FA verify error:', error);
    res.status(500).json({ error: 'Failed to verify 2FA code' });
  }
});

// Get backup codes
router.get('/backup-codes', authenticate, readLimiter, async (req, res) => {
  try {
    const userId = req.user.userId;

    const codes = await dbAll(
      'SELECT code, used, used_at, created_at FROM two_factor_backup_codes WHERE user_id = ? ORDER BY created_at DESC',
      [userId]
    );

    res.json(codes);
  } catch (error) {
    logger.error('Get backup codes error:', error);
    res.status(500).json({ error: 'Failed to retrieve backup codes' });
  }
});

// Regenerate backup codes
router.post('/backup-codes/regenerate', authenticate, modifyLimiter, async (req, res) => {
  try {
    const { password, token } = req.body;
    const userId = req.user.userId;

    if (!password || !token) {
      return res.status(400).json({ error: 'Password and 2FA token required' });
    }

    // Verify password
    const bcrypt = require('bcrypt');
    const user = await dbGet('SELECT password, two_factor_secret FROM users WHERE id = ?', [userId]);
    const validPassword = await bcrypt.compare(password, user.password);

    if (!validPassword) {
      return res.status(401).json({ error: 'Invalid password' });
    }

    // Verify 2FA token
    const verified = speakeasy.totp.verify({
      secret: user.two_factor_secret,
      encoding: 'base32',
      token: token,
      window: 2
    });

    if (!verified) {
      return res.status(400).json({ error: 'Invalid 2FA code' });
    }

    // Delete old backup codes
    await dbRun('DELETE FROM two_factor_backup_codes WHERE user_id = ?', [userId]);

    // Generate new backup codes
    const backupCodes = [];
    for (let i = 0; i < 10; i++) {
      const code = crypto.randomBytes(4).toString('hex').toUpperCase();
      backupCodes.push(code);
      
      await dbRun(
        'INSERT INTO two_factor_backup_codes (user_id, code) VALUES (?, ?)',
        [userId, code]
      );
    }

    logger.info(`Backup codes regenerated for user ${userId}`);

    res.json({
      success: true,
      message: 'Backup codes regenerated',
      backupCodes: backupCodes
    });
  } catch (error) {
    logger.error('Backup codes regenerate error:', error);
    res.status(500).json({ error: 'Failed to regenerate backup codes' });
  }
});

// Check 2FA status
router.get('/status', authenticate, readLimiter, async (req, res) => {
  try {
    const userId = req.user.userId;

    const user = await dbGet(
      'SELECT two_factor_enabled FROM users WHERE id = ?',
      [userId]
    );

    const backupCodesCount = await dbGet(
      'SELECT COUNT(*) as total, SUM(CASE WHEN used = 0 THEN 1 ELSE 0 END) as unused FROM two_factor_backup_codes WHERE user_id = ?',
      [userId]
    );

    res.json({
      enabled: !!user.two_factor_enabled,
      backupCodesTotal: backupCodesCount.total || 0,
      backupCodesUnused: backupCodesCount.unused || 0
    });
  } catch (error) {
    logger.error('2FA status error:', error);
    res.status(500).json({ error: 'Failed to get 2FA status' });
  }
});

module.exports = router;