icynet-auth-server/src/modules/ssr-front-end/two-factor/two-factor.controller.ts

import { Body, Controller, Get, Post, Req, Res } from '@nestjs/common';
import { Request, Response } from 'express';
import { AuditAction } from 'src/modules/objects/audit/audit.enum';
import { AuditService } from 'src/modules/objects/audit/audit.service';
import { UserTOTPService } from 'src/modules/objects/user-token/user-totp-token.service';
import { UserService } from 'src/modules/objects/user/user.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';

interface ChallengeType {
  secret: string;
  type: 'totp';
  user: string;
}

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

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

    if (!twoFA) {
      const challengeString = req.query.challenge as string;
      if (challengeString) {
        const challenge = await this.token.decryptChallenge<ChallengeType>(
          challengeString,
        );
        if (challenge.type === 'totp' && challenge.user === req.user.uuid) {
          secret = challenge.secret;
        }
      }

      if (!secret) {
        secret = this.totp.createTOTPSecret();
        const challenge = { type: 'totp', secret, user: req.user.uuid };
        const encrypted = await this.token.encryptChallenge(challenge);
        const cleanURL = req.originalUrl.replace(/\?(.*)$/, '');
        res.redirect(`${cleanURL}?challenge=${encrypted}`);
        return;
      }

      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(
    @Body() body: { code: string },
    @Req() req: Request,
    @Res() res: Response,
  ) {
    let secret: string;
    try {
      const challengeString = req.query.challenge as string;
      if (!challengeString || !body.code) {
        throw new Error('Invalid request');
      }

      const challenge = await this.token.decryptChallenge<ChallengeType>(
        challengeString,
      );
      secret = challenge.secret;

      if (
        challenge.type !== 'totp' ||
        challenge.user !== req.user.uuid ||
        !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: unknown) {
      req.flash('message', {
        error: true,
        text: (e as Error).message,
      });
      res.redirect('/two-factor');
      return;
    }

    // TODO: show the recovery tokens to the user
    await this.totp.activateTOTP(req.user, secret);
    await this.audit.auditRequest(req, AuditAction.TOTP_ACTIVATE);
    req.flash('message', {
      error: false,
      text: 'Two-factor authenticator has been enabled successfully. Your account is now more secure!',
    });
    res.redirect('/');
  }

  @Get('disable')
  public async disableTwoFA(@Req() req: Request, @Res() res: Response) {
    const twoFA = await this.totp.getUserTOTP(req.user);
    if (!twoFA) {
      return res.redirect('/');
    }

    res.render('password', this.form.populateTemplate(req));
  }

  @Post('disable')
  public async disableTwoFAPost(
    @Req() req: Request,
    @Res() res: Response,
    @Body() body: { password: string },
  ) {
    const twoFA = await this.totp.getUserTOTP(req.user);
    if (!twoFA) {
      return res.redirect('/');
    }

    try {
      if (!body.password) {
        throw new Error('Please enter your password');
      }

      if (
        !(await this.user.comparePasswords(req.user.password, body.password))
      ) {
        throw new Error('The entered password is invalid.');
      }

      await this.totp.deactivateTOTP(twoFA);
      await this.audit.auditRequest(req, AuditAction.TOTP_DEACTIVATE);
    } catch (e: unknown) {
      req.flash('message', {
        error: true,
        text: (e as Error).message,
      });
      res.redirect('/account/two-factor/disable');
      return;
    }

    req.flash('message', {
      error: false,
      text: 'Two-factor authenticator has been disabled successfully. Your account is now less secure!!',
    });
    res.redirect('/');
  }
}