HOA_Financial_Platform/backend/src/modules/auth/mfa.service.ts

import { Injectable, Logger, BadRequestException, UnauthorizedException } from '@nestjs/common';
import { DataSource } from 'typeorm';
import * as bcrypt from 'bcryptjs';
import { generateSecret, generateURI, verifySync } from 'otplib';
import * as QRCode from 'qrcode';
import { randomBytes } from 'crypto';

@Injectable()
export class MfaService {
  private readonly logger = new Logger(MfaService.name);

  constructor(private dataSource: DataSource) {}

  /**
   * Generate MFA setup data (secret + QR code) for a user.
   */
  async generateSetup(userId: string): Promise<{ secret: string; qrDataUrl: string; otpauthUrl: string }> {
    const userRows = await this.dataSource.query(
      `SELECT email, mfa_enabled FROM shared.users WHERE id = $1`,
      [userId],
    );
    if (userRows.length === 0) throw new BadRequestException('User not found');

    const secret = generateSecret();
    const otpauthUrl = generateURI({ secret, issuer: 'HOA LedgerIQ', label: userRows[0].email });
    const qrDataUrl = await QRCode.toDataURL(otpauthUrl);

    // Store the secret temporarily (not verified yet)
    await this.dataSource.query(
      `UPDATE shared.users SET mfa_secret = $1, updated_at = NOW() WHERE id = $2`,
      [secret, userId],
    );

    return { secret, qrDataUrl, otpauthUrl };
  }

  /**
   * Enable MFA after verifying the initial TOTP code.
   * Returns recovery codes.
   */
  async enableMfa(userId: string, token: string): Promise<{ recoveryCodes: string[] }> {
    const userRows = await this.dataSource.query(
      `SELECT mfa_secret, mfa_enabled FROM shared.users WHERE id = $1`,
      [userId],
    );
    if (userRows.length === 0) throw new BadRequestException('User not found');
    if (!userRows[0].mfa_secret) throw new BadRequestException('MFA setup not initiated');
    if (userRows[0].mfa_enabled) throw new BadRequestException('MFA is already enabled');

    // Verify the token
    const result = verifySync({ token, secret: userRows[0].mfa_secret });
    if (!result.valid) throw new BadRequestException('Invalid verification code');

    // Generate recovery codes
    const recoveryCodes = Array.from({ length: 10 }, () =>
      randomBytes(4).toString('hex').toUpperCase(),
    );

    // Hash recovery codes for storage
    const hashedCodes = await Promise.all(
      recoveryCodes.map((code) => bcrypt.hash(code, 10)),
    );

    // Enable MFA
    await this.dataSource.query(
      `UPDATE shared.users SET
        mfa_enabled = true,
        totp_verified_at = NOW(),
        recovery_codes = $1,
        updated_at = NOW()
       WHERE id = $2`,
      [JSON.stringify(hashedCodes), userId],
    );

    this.logger.log(`MFA enabled for user ${userId}`);
    return { recoveryCodes };
  }

  /**
   * Verify a TOTP code during login.
   */
  async verifyMfa(userId: string, token: string): Promise<boolean> {
    const userRows = await this.dataSource.query(
      `SELECT mfa_secret, mfa_enabled FROM shared.users WHERE id = $1`,
      [userId],
    );
    if (userRows.length === 0 || !userRows[0].mfa_enabled) return false;

    const result = verifySync({ token, secret: userRows[0].mfa_secret });
    return result.valid;
  }

  /**
   * Verify a recovery code (consumes it on success).
   */
  async verifyRecoveryCode(userId: string, code: string): Promise<boolean> {
    const userRows = await this.dataSource.query(
      `SELECT recovery_codes FROM shared.users WHERE id = $1`,
      [userId],
    );
    if (userRows.length === 0 || !userRows[0].recovery_codes) return false;

    const hashedCodes: string[] = JSON.parse(userRows[0].recovery_codes);

    for (let i = 0; i < hashedCodes.length; i++) {
      const match = await bcrypt.compare(code.toUpperCase(), hashedCodes[i]);
      if (match) {
        // Remove the used code
        hashedCodes.splice(i, 1);
        await this.dataSource.query(
          `UPDATE shared.users SET recovery_codes = $1, updated_at = NOW() WHERE id = $2`,
          [JSON.stringify(hashedCodes), userId],
        );
        this.logger.log(`Recovery code used for user ${userId}`);
        return true;
      }
    }

    return false;
  }

  /**
   * Disable MFA (requires password verification done by caller).
   */
  async disableMfa(userId: string): Promise<void> {
    await this.dataSource.query(
      `UPDATE shared.users SET
        mfa_enabled = false,
        mfa_secret = NULL,
        totp_verified_at = NULL,
        recovery_codes = NULL,
        updated_at = NOW()
       WHERE id = $1`,
      [userId],
    );
    this.logger.log(`MFA disabled for user ${userId}`);
  }

  /**
   * Get MFA status for a user.
   */
  async getStatus(userId: string): Promise<{ enabled: boolean; hasRecoveryCodes: boolean }> {
    const rows = await this.dataSource.query(
      `SELECT mfa_enabled, recovery_codes FROM shared.users WHERE id = $1`,
      [userId],
    );
    if (rows.length === 0) return { enabled: false, hasRecoveryCodes: false };

    return {
      enabled: rows[0].mfa_enabled || false,
      hasRecoveryCodes: !!rows[0].recovery_codes && JSON.parse(rows[0].recovery_codes || '[]').length > 0,
    };
  }
}