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('/'); } }