icynet-auth-server/src/modules/features/login/login.controller.ts

import {
  Body,
  Controller,
  Get,
  Post,
  Query,
  Render,
  Req,
  Res,
  Session,
} from '@nestjs/common';
import { Request, Response } from 'express';
import { SessionData } from 'express-session';
import { User } from 'src/modules/objects/user/user.entity';
import { UserService } from 'src/modules/objects/user/user.service';
import { FormUtilityService } from 'src/modules/utility/services/form-utility.service';
import { TokenService } from 'src/modules/utility/services/token.service';

@Controller('/login')
export class LoginController {
  constructor(
    private readonly userService: UserService,
    private readonly formUtil: FormUtilityService,
    private readonly token: TokenService,
  ) {}

  @Get()
  @Render('login')
  public loginView(
    @Session() session: SessionData,
    @Req() req: Request,
  ): Record<string, any> {
    return this.formUtil.populateTemplate(req, session);
  }

  @Post()
  public async loginRequest(
    @Req() req: Request,
    @Res() res: Response,
    @Body() body: { username: string; password: string },
    @Query() query: { redirectTo?: string },
  ) {
    const { username, password } = body;
    const user = await this.userService.getByUsername(username);

    // User exists and password matches
    if (
      !user ||
      !user.activated ||
      !(await this.userService.comparePasswords(user.password, password))
    ) {
      req.flash('form', { username });
      req.flash('message', {
        error: true,
        text: 'Invalid username or password',
      });

      res.redirect(req.originalUrl);
      return;
    }

    if (this.userService.userHasTOTP(user)) {
      const challenge = { type: 'totp', user: user.uuid };
      req.session.challenge = await this.token.encryptChallenge(challenge);
      res.redirect(
        '/login/verify' +
          (query.redirectTo ? '?redirectTo=' + query.redirectTo : ''),
      );
    }

    req.session.user = user;
    res.redirect(query.redirectTo ? decodeURIComponent(query.redirectTo) : '/');
  }

  @Get('verify')
  public verifyUserTokenView(
    @Session() session: SessionData,
    @Query() query: { redirectTo?: string },
    @Req() req: Request,
    @Res() res: Response,
  ) {
    if (!session.challenge) {
      req.flash('message', {
        error: true,
        text: 'An unexpected error occured, please log in again.',
      });

      res.redirect(
        '/login' + (query.redirectTo ? '?redirectTo=' + query.redirectTo : ''),
      );
      return;
    }

    res.render('totp-verify', this.formUtil.populateTemplate(req, session));
  }

  @Post('verify')
  public async verifyUserToken(
    @Session() session: SessionData,
    @Query() query: { redirectTo?: string },
    @Body() body: { totp: string },
    @Req() req: Request,
    @Res() res: Response,
  ) {
    let user: User;

    try {
      if (!session.challenge) {
        throw new Error('No challenge');
      }

      const challenge = await this.token.decryptChallenge(session.challenge);
      if (!challenge || challenge.type !== 'totp' || !challenge.user) {
        throw new Error('Bad challenge');
      }

      user = await this.userService.getByUUID(challenge.user);
      if (!user) {
        throw new Error('Bad challenge');
      }
    } catch (e: any) {
      req.flash('message', {
        error: true,
        text: 'An unexpected error occured, please log in again.',
      });

      res.redirect(
        '/login' + (query.redirectTo ? '?redirectTo=' + query.redirectTo : ''),
      );
      return;
    }

    try {
      const totp = await this.userService.getUserTOTP(user);

      if (!this.userService.validateTOTP(totp.token, body.totp)) {
        throw new Error('Invalid code!');
      }
    } catch (e: any) {
      req.flash('message', {
        error: true,
        text: e.message,
      });
      res.redirect(req.originalUrl);
      return;
    }

    session.challenge = null;
    res.redirect(query.redirectTo ? decodeURIComponent(query.redirectTo) : '/');
  }
}