Add MFA to login screen

seanmorley15 · Dec 12, 2024 · 673a56c · 673a56c
1 parent 54d7a1a
commit 673a56c
Show file tree

Hide file tree

Showing 3 changed files with 86 additions and 44 deletions.
diff --git a/backend/server/main/settings.py b/backend/server/main/settings.py
@@ -47,7 +47,7 @@
     "allauth_ui",
     'allauth',
     'allauth.account',
-    # 'allauth.mfa',
+    'allauth.mfa',
     'allauth.headless',
     'allauth.socialaccount',
     "widget_tweaks",

diff --git a/frontend/src/routes/login/+page.server.ts b/frontend/src/routes/login/+page.server.ts
@@ -1,6 +1,6 @@
-import { fail, redirect } from '@sveltejs/kit';
+import { fail, redirect, type RequestEvent } from '@sveltejs/kit';
 
-import type { Actions, PageServerLoad } from './$types';
+import type { Actions, PageServerLoad, RouteParams } from './$types';
 import { getRandomBackground, getRandomQuote } from '$lib';
 import { fetchCSRFToken } from '$lib/index.server';
 const PUBLIC_SERVER_URL = process.env['PUBLIC_SERVER_URL'];
@@ -25,66 +25,99 @@ export const actions: Actions = {
 	default: async (event) => {
 		const formData = await event.request.formData();
 		const formUsername = formData.get('username');
-
-		let username = formUsername?.toString().toLocaleLowerCase();
-
+		const username = formUsername?.toString().toLowerCase();
 		const password = formData.get('password');
+		const totp = formData.get('totp');
 
 		const serverEndpoint = PUBLIC_SERVER_URL || 'http://localhost:8000';
-
 		const csrfToken = await fetchCSRFToken();
 
+		// Initial login attempt
 		const loginFetch = await event.fetch(`${serverEndpoint}/_allauth/browser/v1/auth/login`, {
 			method: 'POST',
 			headers: {
 				'X-CSRFToken': csrfToken,
 				'Content-Type': 'application/json',
 				Cookie: `csrftoken=${csrfToken}`
 			},
-			body: JSON.stringify({
-				username,
-				password
-			}),
+			body: JSON.stringify({ username, password }),
 			credentials: 'include'
 		});
 
-		const loginResponse = await loginFetch.json();
-		if (!loginFetch.ok) {
-			// get the value of the first key in the object
-			const firstKey = Object.keys(loginResponse)[0] || 'error';
-			const error = loginResponse[firstKey][0] || 'Invalid username or password';
-			return fail(400, {
-				message: error
-			});
-		} else {
-			const setCookieHeader = loginFetch.headers.get('Set-Cookie');
-
-			console.log('setCookieHeader:', setCookieHeader);
-
-			if (setCookieHeader) {
-				// Regular expression to match sessionid cookie and its expiry
-				const sessionIdRegex = /sessionid=([^;]+).*?expires=([^;]+)/;
-				const match = setCookieHeader.match(sessionIdRegex);
-
-				if (match) {
-					const sessionId = match[1];
-					const expiryString = match[2];
-					const expiryDate = new Date(expiryString);
+		if (loginFetch.status === 200) {
+			// Login successful without MFA
+			handleSuccessfulLogin(event, loginFetch);
+			return redirect(302, '/');
+		} else if (loginFetch.status === 401) {
+			// MFA required
+			if (!totp) {
+				return fail(401, {
+					message: 'Multi-factor authentication required',
+					mfa_required: true
+				});
+			} else {
+				// Attempt MFA authentication
+				const sessionId = extractSessionId(loginFetch.headers.get('Set-Cookie'));
+				const mfaLoginFetch = await event.fetch(
+					`${serverEndpoint}/_allauth/browser/v1/auth/2fa/authenticate`,
+					{
+						method: 'POST',
+						headers: {
+							'X-CSRFToken': csrfToken,
+							'Content-Type': 'application/json',
+							Cookie: `csrftoken=${csrfToken}; sessionid=${sessionId}`
+						},
+						body: JSON.stringify({ code: totp }),
+						credentials: 'include'
+					}
+				);
 
-					console.log('Session ID:', sessionId);
-					console.log('Expiry Date:', expiryDate);
-
-					// Set the sessionid cookie
-					event.cookies.set('sessionid', sessionId, {
-						path: '/',
-						httpOnly: true,
-						sameSite: 'lax',
-						secure: true,
-						expires: expiryDate
+				if (mfaLoginFetch.ok) {
+					// MFA successful
+					handleSuccessfulLogin(event, mfaLoginFetch);
+					return redirect(302, '/');
+				} else {
+					// MFA failed
+					const mfaLoginResponse = await mfaLoginFetch.json();
+					return fail(401, {
+						message: mfaLoginResponse.error || 'Invalid MFA code',
+						mfa_required: true
 					});
 				}
 			}
-			redirect(302, '/');
+		} else {
+			// Login failed
+			const loginResponse = await loginFetch.json();
+			const firstKey = Object.keys(loginResponse)[0] || 'error';
+			const error = loginResponse[firstKey][0] || 'Invalid username or password';
+			return fail(400, { message: error });
 		}
 	}
 };
+
+function handleSuccessfulLogin(event: RequestEvent<RouteParams, '/login'>, response: Response) {
+	const setCookieHeader = response.headers.get('Set-Cookie');
+	if (setCookieHeader) {
+		const sessionIdRegex = /sessionid=([^;]+).*?expires=([^;]+)/;
+		const match = setCookieHeader.match(sessionIdRegex);
+		if (match) {
+			const [, sessionId, expiryString] = match;
+			event.cookies.set('sessionid', sessionId, {
+				path: '/',
+				httpOnly: true,
+				sameSite: 'lax',
+				secure: true,
+				expires: new Date(expiryString)
+			});
+		}
+	}
+}
+
+function extractSessionId(setCookieHeader: string | null) {
+	if (setCookieHeader) {
+		const sessionIdRegex = /sessionid=([^;]+)/;
+		const match = setCookieHeader.match(sessionIdRegex);
+		return match ? match[1] : '';
+	}
+	return '';
+}
diff --git a/frontend/src/routes/login/+page.svelte b/frontend/src/routes/login/+page.svelte
@@ -51,6 +51,15 @@
 						id="password"
 						class="block input input-bordered w-full max-w-xs"
 					/><br />
+					{#if $page.form?.mfa_required}
+						<label for="password">TOTP</label>
+						<input
+							type="password"
+							name="totp"
+							id="totp"
+							class="block input input-bordered w-full max-w-xs"
+						/><br />
+					{/if}
 					<button class="py-2 px-4 btn btn-primary mr-2">{$t('auth.login')}</button>
 
 					<div class="flex justify-between mt-4">