Feat: 2fa (#1783)

* preplan otp, better qrcode library

* add 2fa as feature

* add totp generation

* working totp lifecycle

* don't allow disabled user to log in

not a security issue as permission handler would fail anyway

* require 2fa on login

if enabled

* update packages

* fix typo

* remove console.logs

src/server/database/repositories/user/schema.ts

@@ -10,6 +10,8 @@ export const user = sqliteTable('users_table', {
   email: text(),
   name: text().notNull(),
   role: int().$type<Role>().notNull(),
+  totpKey: text('totp_key'),
+  totpVerified: int('totp_verified', { mode: 'boolean' }).notNull(),
   enabled: int({ mode: 'boolean' }).notNull(),
   createdAt: text('created_at')
     .notNull()

src/server/database/repositories/user/service.ts

@@ -1,7 +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(),
@@ -21,6 +38,14 @@ function createPreparedStatement(db: DBType) {
       })
       .where(eq(user.id, sql.placeholder('id')))
       .prepare(),
+    updateKey: db
+      .update(user)
+      .set({
+        totpKey: sql.placeholder('key') as never as string,
+        totpVerified: false,
+      })
+      .where(eq(user.id, sql.placeholder('id')))
+      .prepare(),
   };
 }
 
@@ -67,6 +92,7 @@ export class UserService {
         email: null,
         name: 'Administrator',
         role: userCount === 0 ? roles.ADMIN : roles.CLIENT,
+        totpVerified: false,
         enabled: true,
       });
     });
@@ -105,4 +131,121 @@ export class UserService {
         .execute();
     });
   }
+
+  updateTotpKey(id: ID, key: string | null) {
+    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,
+          });
+
+          const valid = totp.validate({ token: code, window: 1 });
+
+          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
+        .findFirst({ where: eq(user.id, id) })
+        .execute();
+
+      if (!txUser) {
+        throw new Error('User not found');
+      }
+
+      if (!txUser.totpKey) {
+        throw new Error('TOTP key is not set');
+      }
+
+      const totp = new TOTP({
+        issuer: 'wg-easy',
+        label: txUser.username,
+        algorithm: 'SHA1',
+        digits: 6,
+        period: 30,
+        secret: txUser.totpKey,
+      });
+
+      const valid = totp.validate({ token: code, window: 1 });
+
+      if (valid === null) {
+        throw new Error('Invalid TOTP code');
+      }
+
+      await tx
+        .update(user)
+        .set({ totpVerified: true })
+        .where(eq(user.id, id))
+        .execute();
+    });
+  }
+
+  deleteTotpKey(id: ID, currentPassword: string) {
+    return this.#db.transaction(async (tx) => {
+      const txUser = await tx.query.user
+        .findFirst({ where: eq(user.id, id) })
+        .execute();
+
+      if (!txUser) {
+        throw new Error('User not found');
+      }
+
+      const passwordValid = await isPasswordValid(
+        currentPassword,
+        txUser.password
+      );
+
+      if (!passwordValid) {
+        throw new Error('Invalid password');
+      }
+
+      await tx
+        .update(user)
+        .set({ totpKey: null, totpVerified: false })
+        .where(eq(user.id, id))
+        .execute();
+    });
+  }
 }

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
@@ -58,3 +64,17 @@ export const UserUpdatePasswordSchema = z
   .refine((val) => val.newPassword === val.confirmPassword, {
     message: t('zod.user.passwordMatch'),
   });
+
+export const UserUpdateTotpSchema = z.union([
+  z.object({
+    type: z.literal('setup'),
+  }),
+  z.object({
+    type: z.literal('create'),
+    code: totpCode,
+  }),
+  z.object({
+    type: z.literal('delete'),
+    currentPassword: password,
+  }),
+]);