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

import {
  Controller,
  Post,
  Get,
  Body,
  UseGuards,
  Request,
  Res,
  BadRequestException,
  UnauthorizedException,
} from '@nestjs/common';
import { ApiTags, ApiOperation, ApiBearerAuth } from '@nestjs/swagger';
import { Throttle } from '@nestjs/throttler';
import { JwtService } from '@nestjs/jwt';
import { Response } from 'express';
import { MfaService } from './mfa.service';
import { AuthService } from './auth.service';
import { JwtAuthGuard } from './guards/jwt-auth.guard';
import { AllowViewer } from '../../common/decorators/allow-viewer.decorator';

const COOKIE_NAME = 'ledgeriq_rt';
const isProduction = process.env.NODE_ENV === 'production';

@ApiTags('auth')
@Controller('auth/mfa')
export class MfaController {
  constructor(
    private mfaService: MfaService,
    private authService: AuthService,
    private jwtService: JwtService,
  ) {}

  @Post('setup')
  @ApiOperation({ summary: 'Generate MFA setup (QR code + secret)' })
  @ApiBearerAuth()
  @UseGuards(JwtAuthGuard)
  async setup(@Request() req: any) {
    return this.mfaService.generateSetup(req.user.sub);
  }

  @Post('enable')
  @ApiOperation({ summary: 'Enable MFA after verifying TOTP code' })
  @ApiBearerAuth()
  @UseGuards(JwtAuthGuard)
  async enable(@Request() req: any, @Body() body: { token: string }) {
    if (!body.token) throw new BadRequestException('TOTP code required');
    return this.mfaService.enableMfa(req.user.sub, body.token);
  }

  @Post('verify')
  @ApiOperation({ summary: 'Verify MFA during login flow' })
  @Throttle({ default: { limit: 5, ttl: 60000 } })
  async verify(
    @Body() body: { mfaToken: string; token: string; useRecovery?: boolean },
    @Res({ passthrough: true }) res: Response,
  ) {
    if (!body.mfaToken || !body.token) {
      throw new BadRequestException('mfaToken and verification code required');
    }

    // Decode the MFA challenge token
    let payload: any;
    try {
      payload = this.jwtService.verify(body.mfaToken);
      if (payload.type !== 'mfa_challenge') throw new Error('Wrong token type');
    } catch {
      throw new UnauthorizedException('Invalid or expired MFA challenge');
    }

    const userId = payload.sub;
    let verified = false;

    if (body.useRecovery) {
      verified = await this.mfaService.verifyRecoveryCode(userId, body.token);
    } else {
      verified = await this.mfaService.verifyMfa(userId, body.token);
    }

    if (!verified) {
      throw new UnauthorizedException('Invalid verification code');
    }

    // MFA passed — issue full session
    const result = await this.authService.completeMfaLogin(userId);
    if (result.refreshToken) {
      res.cookie(COOKIE_NAME, result.refreshToken, {
        httpOnly: true,
        secure: isProduction,
        sameSite: 'strict',
        path: '/api/auth',
        maxAge: 30 * 24 * 60 * 60 * 1000,
      });
    }
    const { refreshToken: _rt, ...response } = result;
    return response;
  }

  @Post('disable')
  @ApiOperation({ summary: 'Disable MFA (requires password)' })
  @ApiBearerAuth()
  @UseGuards(JwtAuthGuard)
  async disable(@Request() req: any, @Body() body: { password: string }) {
    if (!body.password) throw new BadRequestException('Password required to disable MFA');

    // Verify password first
    const user = await this.authService.validateUser(req.user.email, body.password);
    if (!user) throw new UnauthorizedException('Invalid password');

    await this.mfaService.disableMfa(req.user.sub);
    return { success: true };
  }

  @Get('status')
  @ApiOperation({ summary: 'Get MFA status' })
  @ApiBearerAuth()
  @UseGuards(JwtAuthGuard)
  @AllowViewer()
  async status(@Request() req: any) {
    return this.mfaService.getStatus(req.user.sub);
  }
}