icynet-auth-server/src/modules/features/two-factor/two-factor.controller.ts

import { Body, Controller, Get, Post, Req, Res, Session } from '@nestjs/common';
import { Request, Response } from 'express';
import { SessionData } from 'express-session';
import { UserTOTPService } from 'src/modules/objects/user-token/user-totp-token.service';
import { FormUtilityService } from 'src/modules/utility/services/form-utility.service';
import { QRCodeService } from 'src/modules/utility/services/qr-code.service';
import { TokenService } from 'src/modules/utility/services/token.service';

@Controller('/account/two-factor')
export class TwoFactorController {
  constructor(
    private totp: UserTOTPService,
    private qr: QRCodeService,
    private token: TokenService,
    private form: FormUtilityService,
  ) {}

  @Get('activate')
  public async twoFAStatus(
    @Session() session: SessionData,
    @Req() req: Request,
    @Res() res: Response,
  ) {
    const twoFA = await this.totp.getUserTOTP(req.user);
    let secret: string;

    if (!twoFA) {
      if (session.challenge) {
        const challenge = await this.token.decryptChallenge(session.challenge);
        if (challenge.type === 'totp') {
          secret = challenge.secret;
        }
      }

      if (!secret) {
        secret = this.totp.createTOTPSecret();
        const challenge = { type: 'totp', secret };
        session.challenge = await this.token.encryptChallenge(challenge);
      }

      const url = this.totp.getTOTPURL(secret, req.user.username);
      const qrcode = await this.qr.createQRDataURI(url);

      res.render('two-factor/activate', {
        ...this.form.populateTemplate(req),
        qrcode,
      });

      return;
    }

    res.redirect('/');
  }

  @Post('activate')
  public async twoFAActivate(
    @Session() session: SessionData,
    @Body() body: { code: string },
    @Req() req: Request,
    @Res() res: Response,
  ) {
    let secret: string;
    try {
      if (!session.challenge || !body.code) {
        throw new Error('Invalid request');
      }

      const challenge = await this.token.decryptChallenge(session.challenge);
      secret = challenge.secret;

      if (challenge.type !== 'totp' || !secret) {
        throw new Error('Invalid request');
      }

      const verify = this.totp.validateTOTP(secret, body.code);
      if (!verify) {
        throw new Error('Invalid code! Try again.');
      }
    } catch (e: any) {
      req.flash('message', {
        error: true,
        text: e.message,
      });
      res.redirect('/two-factor');
      return;
    }

    // TODO: show the recovery tokens to the user
    await this.totp.activateTOTP(req.user, secret);
    session.challenge = null;
    res.redirect('/');
  }
}