template-saas-backend-v2/src/modules/auth/services/mfa.service.ts

import {
  Injectable,
  BadRequestException,
  UnauthorizedException,
  NotFoundException,
} from '@nestjs/common';
import { InjectRepository } from '@nestjs/typeorm';
import { Repository } from 'typeorm';
import { ConfigService } from '@nestjs/config';
import * as speakeasy from 'speakeasy';
import * as QRCode from 'qrcode';
import * as bcrypt from 'bcrypt';
import * as crypto from 'crypto';
import { User } from '../entities';
import {
  SetupMfaResponseDto,
  VerifyMfaSetupDto,
  DisableMfaDto,
  MfaStatusDto,
  BackupCodesResponseDto,
} from '../dto/mfa.dto';

@Injectable()
export class MfaService {
  private readonly appName: string;

  constructor(
    @InjectRepository(User)
    private readonly userRepository: Repository<User>,
    private readonly configService: ConfigService,
  ) {
    this.appName = this.configService.get<string>('app.name') || 'Template SaaS';
  }

  /**
   * Initialize MFA setup - generate secret and QR code
   */
  async setupMfa(userId: string): Promise<SetupMfaResponseDto> {
    const user = await this.userRepository.findOne({ where: { id: userId } });

    if (!user) {
      throw new NotFoundException('User not found');
    }

    if (user.mfa_enabled) {
      throw new BadRequestException('MFA is already enabled');
    }

    // Generate TOTP secret
    const secret = speakeasy.generateSecret({
      name: `${this.appName} (${user.email})`,
      length: 20,
    });

    // Generate QR code
    const qrCodeDataUrl = await QRCode.toDataURL(secret.otpauth_url!);

    // Generate backup codes
    const backupCodes = this.generateBackupCodes();

    return {
      secret: secret.base32,
      qrCodeDataUrl,
      backupCodes,
    };
  }

  /**
   * Verify MFA setup and enable
   */
  async verifyMfaSetup(
    userId: string,
    dto: VerifyMfaSetupDto,
  ): Promise<{ success: boolean; message: string }> {
    const user = await this.userRepository.findOne({ where: { id: userId } });

    if (!user) {
      throw new NotFoundException('User not found');
    }

    if (user.mfa_enabled) {
      throw new BadRequestException('MFA is already enabled');
    }

    // Verify the TOTP code
    const isValid = speakeasy.totp.verify({
      secret: dto.secret,
      encoding: 'base32',
      token: dto.code,
      window: 1, // Allow 1 step (30 seconds) tolerance
    });

    if (!isValid) {
      throw new BadRequestException('Invalid verification code');
    }

    // Generate and hash backup codes
    const backupCodes = this.generateBackupCodes();
    const hashedBackupCodes = await Promise.all(
      backupCodes.map((code) => bcrypt.hash(code, 10)),
    );

    // Enable MFA and store secret
    await this.userRepository.update(
      { id: userId },
      {
        mfa_enabled: true,
        mfa_secret: this.encryptSecret(dto.secret),
        mfa_backup_codes: hashedBackupCodes,
        mfa_enabled_at: new Date(),
      },
    );

    return {
      success: true,
      message: 'MFA enabled successfully. Please save your backup codes.',
    };
  }

  /**
   * Verify TOTP code during login
   */
  async verifyMfaCode(
    userId: string,
    code: string,
    isBackupCode: boolean = false,
  ): Promise<boolean> {
    const user = await this.userRepository.findOne({ where: { id: userId } });

    if (!user) {
      throw new NotFoundException('User not found');
    }

    if (!user.mfa_enabled || !user.mfa_secret) {
      throw new BadRequestException('MFA is not enabled for this user');
    }

    if (isBackupCode) {
      return this.verifyBackupCode(user, code);
    }

    // Decrypt secret and verify TOTP
    const secret = this.decryptSecret(user.mfa_secret);

    const isValid = speakeasy.totp.verify({
      secret,
      encoding: 'base32',
      token: code,
      window: 1,
    });

    return isValid;
  }

  /**
   * Disable MFA for user
   */
  async disableMfa(
    userId: string,
    dto: DisableMfaDto,
  ): Promise<{ success: boolean; message: string }> {
    const user = await this.userRepository.findOne({ where: { id: userId } });

    if (!user) {
      throw new NotFoundException('User not found');
    }

    if (!user.mfa_enabled) {
      throw new BadRequestException('MFA is not enabled');
    }

    // Verify password
    const isPasswordValid = await bcrypt.compare(dto.password, user.password_hash);
    if (!isPasswordValid) {
      throw new UnauthorizedException('Invalid password');
    }

    // Verify MFA code
    const isMfaValid = await this.verifyMfaCode(userId, dto.code, dto.code.length > 6);
    if (!isMfaValid) {
      throw new BadRequestException('Invalid verification code');
    }

    // Disable MFA
    await this.userRepository.update(
      { id: userId },
      {
        mfa_enabled: false,
        mfa_secret: null,
        mfa_backup_codes: null,
        mfa_enabled_at: null,
      },
    );

    return {
      success: true,
      message: 'MFA disabled successfully',
    };
  }

  /**
   * Get MFA status for user
   */
  async getMfaStatus(userId: string): Promise<MfaStatusDto> {
    const user = await this.userRepository.findOne({ where: { id: userId } });

    if (!user) {
      throw new NotFoundException('User not found');
    }

    const backupCodesRemaining = user.mfa_backup_codes?.length || 0;

    return {
      enabled: user.mfa_enabled || false,
      enabledAt: user.mfa_enabled_at || undefined,
      backupCodesRemaining,
    };
  }

  /**
   * Regenerate backup codes
   */
  async regenerateBackupCodes(
    userId: string,
    password: string,
    code: string,
  ): Promise<BackupCodesResponseDto> {
    const user = await this.userRepository.findOne({ where: { id: userId } });

    if (!user) {
      throw new NotFoundException('User not found');
    }

    if (!user.mfa_enabled) {
      throw new BadRequestException('MFA is not enabled');
    }

    // Verify password
    const isPasswordValid = await bcrypt.compare(password, user.password_hash);
    if (!isPasswordValid) {
      throw new UnauthorizedException('Invalid password');
    }

    // Verify MFA code
    const isMfaValid = await this.verifyMfaCode(userId, code);
    if (!isMfaValid) {
      throw new BadRequestException('Invalid verification code');
    }

    // Generate new backup codes
    const backupCodes = this.generateBackupCodes();
    const hashedBackupCodes = await Promise.all(
      backupCodes.map((code) => bcrypt.hash(code, 10)),
    );

    // Update user
    await this.userRepository.update(
      { id: userId },
      { mfa_backup_codes: hashedBackupCodes },
    );

    return {
      backupCodes,
      message: 'New backup codes generated. Please save them securely.',
    };
  }

  // ==================== Private Methods ====================

  /**
   * Generate 10 random backup codes
   */
  private generateBackupCodes(): string[] {
    const codes: string[] = [];
    for (let i = 0; i < 10; i++) {
      const code = crypto.randomBytes(4).toString('hex').toUpperCase();
      // Format as XXXX-XXXX for readability
      codes.push(`${code.slice(0, 4)}-${code.slice(4)}`);
    }
    return codes;
  }

  /**
   * Verify a backup code
   */
  private async verifyBackupCode(user: User, code: string): Promise<boolean> {
    if (!user.mfa_backup_codes || user.mfa_backup_codes.length === 0) {
      return false;
    }

    // Normalize code (remove dashes, uppercase)
    const normalizedCode = code.replace(/-/g, '').toUpperCase();
    const formattedCode = `${normalizedCode.slice(0, 4)}-${normalizedCode.slice(4)}`;

    // Check each hashed backup code
    for (let i = 0; i < user.mfa_backup_codes.length; i++) {
      const isMatch = await bcrypt.compare(formattedCode, user.mfa_backup_codes[i]);
      if (isMatch) {
        // Remove the used backup code
        const updatedCodes = [...user.mfa_backup_codes];
        updatedCodes.splice(i, 1);
        await this.userRepository.update(
          { id: user.id },
          { mfa_backup_codes: updatedCodes },
        );
        return true;
      }
    }

    return false;
  }

  /**
   * Encrypt MFA secret for storage
   */
  private encryptSecret(secret: string): string {
    const encryptionKey = this.configService.get<string>('mfa.encryptionKey') ||
      this.configService.get<string>('jwt.secret') ||
      'default-encryption-key-change-me';

    // Use first 32 bytes of key for AES-256
    const key = crypto.createHash('sha256').update(encryptionKey).digest();
    const iv = crypto.randomBytes(16);
    const cipher = crypto.createCipheriv('aes-256-cbc', key, iv);

    let encrypted = cipher.update(secret, 'utf8', 'hex');
    encrypted += cipher.final('hex');

    // Return IV + encrypted data
    return iv.toString('hex') + ':' + encrypted;
  }

  /**
   * Decrypt MFA secret from storage
   */
  private decryptSecret(encryptedSecret: string): string {
    const encryptionKey = this.configService.get<string>('mfa.encryptionKey') ||
      this.configService.get<string>('jwt.secret') ||
      'default-encryption-key-change-me';

    const key = crypto.createHash('sha256').update(encryptionKey).digest();
    const [ivHex, encrypted] = encryptedSecret.split(':');
    const iv = Buffer.from(ivHex, 'hex');
    const decipher = crypto.createDecipheriv('aes-256-cbc', key, iv);

    let decrypted = decipher.update(encrypted, 'hex', 'utf8');
    decrypted += decipher.final('utf8');

    return decrypted;
  }
}