From ea9742a1c64b826ad42e95554b09f147431d20ea Mon Sep 17 00:00:00 2001
From: Bernd Storath <999999bst@gmail.com>
Date: Tue, 1 Apr 2025 13:28:30 +0200
Subject: [PATCH] require 2fa on login

if enabled
---
 src/app/composables/useSubmit.ts              |  7 --
 src/app/pages/admin/interface.vue             |  2 -
 src/app/pages/login.vue                       | 51 ++++++++++++--
 src/app/pages/me.vue                          |  2 +-
 src/i18n/locales/en.json                      | 21 +++---
 src/server/api/me/totp.post.ts                |  8 ++-
 src/server/api/session.post.ts                | 56 ++++++++-------
 .../database/repositories/user/service.ts     | 70 +++++++++++++++++++
 .../database/repositories/user/types.ts       |  8 ++-
 src/server/utils/types.ts                     |  7 ++
 10 files changed, 172 insertions(+), 60 deletions(-)

diff --git a/src/app/composables/useSubmit.ts b/src/app/composables/useSubmit.ts
index 048aadfc..805c2b66 100644
--- a/src/app/composables/useSubmit.ts
+++ b/src/app/composables/useSubmit.ts
@@ -28,7 +28,6 @@ type SubmitOpts<
 > = {
   revert: RevertFn<R, T, O>;
   successMsg?: string;
-  errorMsg?: string;
   noSuccessToast?: boolean;
 };
 
@@ -38,7 +37,6 @@ export function useSubmit<
   T = unknown,
 >(url: R, options: O, opts: SubmitOpts<R, T, O>) {
   const toast = useToast();
-  const { t: $t } = useI18n();
 
   return async (data: unknown) => {
     try {
@@ -47,11 +45,6 @@ export function useSubmit<
         body: data,
       });
 
-      // eslint-disable-next-line @typescript-eslint/no-explicit-any
-      if (!(res as any).success) {
-        throw new Error(opts.errorMsg || $t('toast.errored'));
-      }
-
       if (!opts.noSuccessToast) {
         toast.showToast({
           type: 'success',
diff --git a/src/app/pages/admin/interface.vue b/src/app/pages/admin/interface.vue
index 18988943..55e7d283 100644
--- a/src/app/pages/admin/interface.vue
+++ b/src/app/pages/admin/interface.vue
@@ -86,7 +86,6 @@ const _changeCidr = useSubmit(
   {
     revert,
     successMsg: t('admin.interface.cidrSuccess'),
-    errorMsg: t('admin.interface.cidrError'),
   }
 );
 
@@ -102,7 +101,6 @@ const _restartInterface = useSubmit(
   {
     revert,
     successMsg: t('admin.interface.restartSuccess'),
-    errorMsg: t('admin.interface.restartError'),
   }
 );
 
diff --git a/src/app/pages/login.vue b/src/app/pages/login.vue
index 9c342f61..76d34ca7 100644
--- a/src/app/pages/login.vue
+++ b/src/app/pages/login.vue
@@ -30,6 +30,18 @@
         autocomplete="current-password"
       />
 
+      <BaseInput
+        v-if="totpRequired"
+        v-model="totp"
+        type="text"
+        name="totp"
+        :placeholder="$t('general.2faCode')"
+        autocomplete="one-time-code"
+        inputmode="numeric"
+        maxlength="6"
+        pattern="\d{6}"
+      />
+
       <label
         class="flex gap-2 whitespace-nowrap"
         :title="$t('login.rememberMeDesc')"
@@ -58,10 +70,15 @@
 const authStore = useAuthStore();
 authStore.update();
 
+const toast = useToast();
+const { t } = useI18n();
+
 const authenticating = ref(false);
 const remember = ref(false);
-const username = ref<null | string>(null);
-const password = ref<null | string>(null);
+const username = ref<string>('');
+const password = ref<string>('');
+const totpRequired = ref(false);
+const totp = ref<string>('');
 
 const _submit = useSubmit(
   '/api/session',
@@ -69,13 +86,32 @@ const _submit = useSubmit(
     method: 'post',
   },
   {
-    revert: async (success) => {
-      authenticating.value = false;
-      password.value = null;
-
+    revert: async (success, data) => {
       if (success) {
-        await navigateTo('/');
+        if (data?.status === 'success') {
+          await navigateTo('/');
+        } else if (data?.status === 'TOTP_REQUIRED') {
+          authenticating.value = false;
+          totpRequired.value = true;
+          toast.showToast({
+            title: t('general.2fa'),
+            message: t('login.2faRequired'),
+            type: 'error',
+          });
+          return;
+        } else if (data?.status === 'INVALID_TOTP_CODE') {
+          authenticating.value = false;
+          totp.value = '';
+          toast.showToast({
+            title: t('general.2fa'),
+            message: t('login.2faWrong'),
+            type: 'error',
+          });
+          return;
+        }
       }
+      authenticating.value = false;
+      password.value = '';
     },
     noSuccessToast: true,
   }
@@ -90,6 +126,7 @@ async function submit() {
     username: username.value,
     password: password.value,
     remember: remember.value,
+    totpCode: totpRequired.value ? totp.value : undefined,
   });
 }
 </script>
diff --git a/src/app/pages/me.vue b/src/app/pages/me.vue
index 57bc79b3..021853a7 100644
--- a/src/app/pages/me.vue
+++ b/src/app/pages/me.vue
@@ -79,7 +79,7 @@
                 <FormTextField
                   id="2facode"
                   v-model="code"
-                  :label="$t('me.2faCode')"
+                  :label="$t('general.2faCode')"
                 />
                 <FormActionField
                   :label="$t('me.enable2fa')"
diff --git a/src/i18n/locales/en.json b/src/i18n/locales/en.json
index 604d837d..060f4f91 100644
--- a/src/i18n/locales/en.json
+++ b/src/i18n/locales/en.json
@@ -17,8 +17,7 @@
     "currentPassword": "Current Password",
     "enable2fa": "Enable Two Factor Authentication",
     "enable2faDesc": "Scan the QR code with your authenticator app or enter the key manually.",
-    "2faKey": "2FA Key",
-    "2faCode": "2FA Code",
+    "2faKey": "TOTP Key",
     "2faCodeDesc": "Enter the code from your authenticator app.",
     "disable2fa": "Disable Two Factor Authentication",
     "disable2faDesc": "Enter your password to disable Two Factor Authentication."
@@ -41,7 +40,8 @@
     "no": "No",
     "confirmPassword": "Confirm Password",
     "loading": "Loading...",
-    "2fa": "Two Factor Authentication"
+    "2fa": "Two Factor Authentication",
+    "2faCode": "TOTP Code"
   },
   "setup": {
     "welcome": "Welcome to your first setup of wg-easy",
@@ -74,11 +74,9 @@
     "signIn": "Sign In",
     "rememberMe": "Remember me",
     "rememberMeDesc": "Stay logged after closing the browser",
-    "insecure": "You can't log in with an insecure connection. Use HTTPS."
-  },
-  "error": {
-    "clear": "Clear",
-    "login": "Log in error"
+    "insecure": "You can't log in with an insecure connection. Use HTTPS.",
+    "2faRequired": "Two Factor Authentication is required",
+    "2faWrong": "Two Factor Authentication is wrong"
   },
   "client": {
     "empty": "There are no clients yet.",
@@ -125,8 +123,7 @@
   "toast": {
     "success": "Success",
     "saved": "Saved",
-    "error": "Error",
-    "errored": "Failed to save"
+    "error": "Error"
   },
   "form": {
     "actions": "Actions",
@@ -163,7 +160,6 @@
     },
     "interface": {
       "cidrSuccess": "Changed CIDR",
-      "cidrError": "Failed to change CIDR",
       "device": "Device",
       "deviceDesc": "Ethernet device the wireguard traffic should be forwarded through",
       "mtuDesc": "MTU WireGuard will use",
@@ -172,8 +168,7 @@
       "restart": "Restart Interface",
       "restartDesc": "Restart the WireGuard interface",
       "restartWarn": "Are you sure to restart the interface? This will disconnect all clients.",
-      "restartSuccess": "Interface restarted",
-      "restartError": "Failed to restart interface"
+      "restartSuccess": "Interface restarted"
     },
     "introText": "Welcome to the admin panel.\n\nHere you can manage the general settings, the configuration, the interface settings and the hooks.\n\nStart by choosing one of the sections in the sidebar."
   },
diff --git a/src/server/api/me/totp.post.ts b/src/server/api/me/totp.post.ts
index 99d0fc6e..53c87f8a 100644
--- a/src/server/api/me/totp.post.ts
+++ b/src/server/api/me/totp.post.ts
@@ -9,8 +9,7 @@ type Response =
       uri: string;
     }
   | { success: boolean; type: 'created' }
-  | { success: boolean; type: 'deleted' }
-  | { success: boolean; type: 'error' };
+  | { success: boolean; type: 'deleted' };
 
 export default definePermissionEventHandler(
   'me',
@@ -58,6 +57,9 @@ export default definePermissionEventHandler(
         type: 'deleted',
       } as Response;
     }
-    return { success: false, type: 'error' } as Response;
+    throw createError({
+      statusCode: 400,
+      statusMessage: 'Invalid request',
+    });
   }
 );
diff --git a/src/server/api/session.post.ts b/src/server/api/session.post.ts
index 0dcc008a..a11f14f4 100644
--- a/src/server/api/session.post.ts
+++ b/src/server/api/session.post.ts
@@ -1,37 +1,41 @@
 import { UserLoginSchema } from '#db/repositories/user/types';
 
 export default defineEventHandler(async (event) => {
-  const { username, password, remember } = await readValidatedBody(
+  const { username, password, remember, totpCode } = await readValidatedBody(
     event,
     validateZod(UserLoginSchema, event)
   );
 
-  // TODO: timing can be used to enumerate usernames
-
-  const user = await Database.users.getByUsername(username);
-  if (!user) {
-    throw createError({
-      statusCode: 401,
-      statusMessage: 'Incorrect credentials',
-    });
-  }
-
-  if (!user.enabled) {
-    throw createError({
-      statusCode: 403,
-      statusMessage: 'User is disabled',
-    });
+  const result = await Database.users.login(username, password, totpCode);
+
+  // TODO: add localization support
+
+  if (!result.success) {
+    switch (result.error) {
+      case 'INCORRECT_CREDENTIALS':
+        throw createError({
+          statusCode: 401,
+          statusMessage: 'Invalid username or password',
+        });
+      case 'TOTP_REQUIRED':
+        return { status: 'TOTP_REQUIRED' };
+      case 'INVALID_TOTP_CODE':
+        return { status: 'INVALID_TOTP_CODE' };
+      case 'USER_DISABLED':
+        throw createError({
+          statusCode: 401,
+          statusMessage: 'User disabled',
+        });
+      case 'UNEXPECTED_ERROR':
+        throw createError({
+          statusCode: 500,
+          statusMessage: 'Unexpected error',
+        });
+    }
+    assertUnreachable(result.error);
   }
 
-  // todo: handle in service
-  const userHashPassword = user.password;
-  const passwordValid = await isPasswordValid(password, userHashPassword);
-  if (!passwordValid) {
-    throw createError({
-      statusCode: 401,
-      statusMessage: 'Incorrect credentials',
-    });
-  }
+  const user = result.user;
 
   const session = await useWGSession(event, remember);
 
@@ -43,5 +47,5 @@ export default defineEventHandler(async (event) => {
 
   SERVER_DEBUG(`New Session: ${data.id} for ${user.id} (${user.username})`);
 
-  return { success: true };
+  return { status: 'success' };
 });
diff --git a/src/server/database/repositories/user/service.ts b/src/server/database/repositories/user/service.ts
index 4c139ab4..9a073249 100644
--- a/src/server/database/repositories/user/service.ts
+++ b/src/server/database/repositories/user/service.ts
@@ -1,8 +1,24 @@
 import { eq, sql } from 'drizzle-orm';
 import { TOTP } from 'otpauth';
 import { user } from './schema';
+import type { UserType } from './types';
 import type { DBType } from '#db/sqlite';
 
+type LoginResult =
+  | {
+      success: true;
+      user: UserType;
+    }
+  | {
+      success: false;
+      error:
+        | 'INCORRECT_CREDENTIALS'
+        | 'TOTP_REQUIRED'
+        | 'USER_DISABLED'
+        | 'INVALID_TOTP_CODE'
+        | 'UNEXPECTED_ERROR';
+    };
+
 function createPreparedStatement(db: DBType) {
   return {
     findAll: db.query.user.findMany().prepare(),
@@ -120,6 +136,60 @@ export class UserService {
     return this.#statements.updateKey.execute({ id, key });
   }
 
+  login(username: string, password: string, code: string | undefined) {
+    return this.#db.transaction(async (tx): Promise<LoginResult> => {
+      const txUser = await tx.query.user
+        .findFirst({ where: eq(user.username, username) })
+        .execute();
+
+      if (!txUser) {
+        return { success: false, error: 'INCORRECT_CREDENTIALS' };
+      }
+
+      const passwordValid = await isPasswordValid(password, txUser.password);
+
+      if (!passwordValid) {
+        return { success: false, error: 'INCORRECT_CREDENTIALS' };
+      }
+
+      if (txUser.totpVerified) {
+        if (!code) {
+          return { success: false, error: 'TOTP_REQUIRED' };
+        } else {
+          if (!txUser.totpKey) {
+            return { success: false, error: 'UNEXPECTED_ERROR' };
+          }
+
+          const totp = new TOTP({
+            issuer: 'wg-easy',
+            label: txUser.username,
+            algorithm: 'SHA1',
+            digits: 6,
+            period: 30,
+            secret: txUser.totpKey,
+          });
+
+          console.log(txUser.totpKey);
+          console.log(code);
+
+          const valid = totp.validate({ token: code, window: 1 });
+
+          console.log(valid);
+
+          if (valid === null) {
+            return { success: false, error: 'INVALID_TOTP_CODE' };
+          }
+        }
+      }
+
+      if (!txUser.enabled) {
+        return { success: false, error: 'USER_DISABLED' };
+      }
+
+      return { success: true, user: txUser };
+    });
+  }
+
   verifyTotp(id: ID, code: string) {
     return this.#db.transaction(async (tx) => {
       const txUser = await tx.query.user
diff --git a/src/server/database/repositories/user/types.ts b/src/server/database/repositories/user/types.ts
index 0d0100d1..4743be73 100644
--- a/src/server/database/repositories/user/types.ts
+++ b/src/server/database/repositories/user/types.ts
@@ -16,10 +16,16 @@ const password = z
 
 const remember = z.boolean({ message: t('zod.user.remember') });
 
+const totpCode = z
+  .string({ message: t('zod.user.totpCode') })
+  .min(6, t('zod.user.totpCode'))
+  .pipe(safeStringRefine);
+
 export const UserLoginSchema = z.object({
   username: username,
   password: password,
   remember: remember,
+  totpCode: totpCode.optional(),
 });
 
 export const UserSetupSchema = z
@@ -65,7 +71,7 @@ export const UserUpdateTotpSchema = z.union([
   }),
   z.object({
     type: z.literal('create'),
-    code: z.string().min(6, t('zod.user.totpCode')),
+    code: totpCode,
   }),
   z.object({
     type: z.literal('delete'),
diff --git a/src/server/utils/types.ts b/src/server/utils/types.ts
index 898cf957..ad43fd1c 100644
--- a/src/server/utils/types.ts
+++ b/src/server/utils/types.ts
@@ -147,3 +147,10 @@ export function validateZod<T>(
     }
   };
 }
+
+/**
+ * exhaustive check
+ */
+export function assertUnreachable(_: never): never {
+  throw new Error("Didn't expect to get here");
+}