import { Body, Controller, Get, Post, Req, Res, Session } from '@nestjs/common'; import { Request, Response } from 'express'; import { SessionData } from 'express-session'; 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'; @Controller('/two-factor') export class TwoFactorController { constructor( private userService: UserService, private qr: QRCodeService, private token: TokenService, private form: FormUtilityService, ) {} @Get() public async twoFAStatus( @Session() session: SessionData, @Req() req: Request, @Res() res: Response, ) { const twoFA = await this.userService.getUserTOTP(session.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.userService.createTOTPSecret(); const challenge = { type: 'totp', secret }; session.challenge = await this.token.encryptChallenge(challenge); } const url = this.userService.getTOTPURL(secret, session.user.username); const qrcode = await this.qr.createQRDataURI(url); res.render('two-factor/activate', { ...this.form.populateTemplate(req, session), qrcode, }); return; } res.redirect('/'); } @Post() 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.userService.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; } await this.userService.activateTOTP(session.user, secret); session.challenge = null; res.redirect('/'); } }