Add 4 files, update 14 files

src/app/api/auth/route.ts

@@ -14,6 +14,7 @@ import {
 } from '@/lib/auth';
 import { withRateLimit, checkRateLimit, getRateLimitKey, getRateLimitHeaders, createRateLimitResponse } from '@/lib/rate-limit';
 import { createPasswordResetToken, hashResetToken, passwordResetExpiry } from '@/lib/password-reset';
+import { createRecoveryCodes, createTotpSecret, hashRecoveryCode, totpUri, verifyTotp } from '@/lib/two-factor';
 
 const loginSchema = z.object({
   email: z.string().email(),
@@ -37,11 +38,56 @@ const resetPasswordSchema = z.object({
   password: z.string().min(12),
 });
 
+const twoFactorTokenSchema = z.object({
+  token: z.string().min(6).max(20),
+});
+
 export const POST = withRateLimit(async function POST(request: Request) {
   try {
     const body = await request.json();
     const { action } = body;
 
+    if (action === 'setup-two-factor') {
+      const session = await getSession();
+      if (!session?.user) return NextResponse.json({ error: 'Not authenticated' }, { status: 401 });
+      if (process.env.AUTH_TOTP_ENABLED !== 'true') return NextResponse.json({ error: 'Two-factor authentication is disabled' }, { status: 403 });
+      const secret = createTotpSecret();
+      return NextResponse.json({ secret, otpauthUrl: totpUri(session.user.email, secret) });
+    }
+
+    if (action === 'enable-two-factor') {
+      const session = await getSession();
+      const parsed = twoFactorTokenSchema.extend({ secret: z.string().min(10) }).safeParse(body);
+      if (!session?.user) return NextResponse.json({ error: 'Not authenticated' }, { status: 401 });
+      if (process.env.AUTH_TOTP_ENABLED !== 'true') return NextResponse.json({ error: 'Two-factor authentication is disabled' }, { status: 403 });
+      if (!parsed.success || !verifyTotp(parsed.data.secret, parsed.data.token)) {
+        return NextResponse.json({ error: 'Invalid authenticator code' }, { status: 400 });
+      }
+      const recoveryCodes = createRecoveryCodes();
+      await db.$transaction([
+        db.user.update({ where: { id: session.userId }, data: { twoFactorSecret: parsed.data.secret } }),
+        db.twoFactorRecoveryCode.deleteMany({ where: { userId: session.userId } }),
+        db.twoFactorRecoveryCode.createMany({
+          data: recoveryCodes.map((code) => ({ userId: session.userId, codeHash: hashRecoveryCode(code) })),
+        }),
+      ]);
+      return NextResponse.json({ recoveryCodes });
+    }
+
+    if (action === 'disable-two-factor') {
+      const session = await getSession();
+      const parsed = twoFactorTokenSchema.safeParse(body);
+      if (!session?.user) return NextResponse.json({ error: 'Not authenticated' }, { status: 401 });
+      if (!parsed.success || !session.user.twoFactorSecret || !verifyTotp(session.user.twoFactorSecret, parsed.data.token)) {
+        return NextResponse.json({ error: 'Invalid authenticator code' }, { status: 400 });
+      }
+      await db.$transaction([
+        db.user.update({ where: { id: session.userId }, data: { twoFactorSecret: null } }),
+        db.twoFactorRecoveryCode.deleteMany({ where: { userId: session.userId } }),
+      ]);
+      return NextResponse.json({ message: 'Two-factor authentication disabled' });
+    }
+
     if (action === 'request-password-reset') {
       const parsed = requestPasswordResetSchema.safeParse(body);
       if (!parsed.success) return NextResponse.json({ error: 'Validation failed' }, { status: 400 });
@@ -129,6 +175,24 @@ export const POST = withRateLimit(async function POST(request: Request) {
         });
       }
 
+      if (user.twoFactorSecret) {
+        const parsedTwoFactor = twoFactorTokenSchema.safeParse(body);
+        if (!parsedTwoFactor.success) {
+          return NextResponse.json({ error: 'Two-factor code required', twoFactorRequired: true }, { status: 401 });
+        }
+        const providedCode = parsedTwoFactor.data.token;
+        const validTotp = verifyTotp(user.twoFactorSecret, providedCode);
+        const recoveryCode = validTotp ? null : await db.twoFactorRecoveryCode.findFirst({
+          where: { userId: user.id, codeHash: hashRecoveryCode(providedCode), usedAt: null },
+        });
+        if (!validTotp && !recoveryCode) {
+          return NextResponse.json({ error: 'Invalid two-factor code', twoFactorRequired: true }, { status: 401 });
+        }
+        if (recoveryCode) {
+          await db.twoFactorRecoveryCode.update({ where: { id: recoveryCode.id }, data: { usedAt: new Date() } });
+        }
+      }
+
       const secure = isSecureRequest(request);
       await createSession(user.id, secure);
 

src/lib/auth.ts

@@ -76,6 +76,7 @@ const sessionUserSelect = Prisma.validator<Prisma.UserSelect>()({
   schoolId: true,
   locale: true,
   isDemo: true,
+  twoFactorSecret: true,
   createdAt: true,
   updatedAt: true,
 });

src/lib/two-factor.test.ts

@@ -0,0 +1,21 @@
+import { describe, expect, test } from 'bun:test';
+import { createRecoveryCodes, createTotpSecret, hashRecoveryCode, totpUri, verifyTotp } from './two-factor';
+
+describe('two factor authentication', () => {
+  test('creates TOTP enrollment data', () => {
+    const secret = createTotpSecret();
+    expect(secret.length).toBeGreaterThan(10);
+    expect(totpUri('teacher@school.example', secret)).toStartWith('otpauth://totp/');
+  });
+
+  test('creates unique hashed recovery codes', () => {
+    const codes = createRecoveryCodes();
+    expect(codes).toHaveLength(10);
+    expect(new Set(codes).size).toBe(10);
+    expect(hashRecoveryCode(codes[0])).not.toBe(codes[0]);
+  });
+
+  test('rejects malformed TOTP values', () => {
+    expect(verifyTotp(createTotpSecret(), 'invalid')).toBeFalse();
+  });
+});

src/lib/two-factor.ts

@@ -0,0 +1,26 @@
+import { createHash, randomBytes } from 'node:crypto';
+import { TOTP } from 'otpauth';
+
+const RECOVERY_CODE_COUNT = 10;
+
+export function createTotpSecret(): string {
+  return new TOTP({ issuer: 'SchoolOS', label: 'SchoolOS account' }).secret.base32;
+}
+
+export function totpUri(email: string, secret: string): string {
+  return new TOTP({ issuer: 'SchoolOS', label: email, secret }).toString();
+}
+
+export function verifyTotp(secret: string, token: string): boolean {
+  const normalizedToken = token.replace(/\s/g, '');
+  if (!/^\d{6}$/.test(normalizedToken)) return false;
+  return new TOTP({ issuer: 'SchoolOS', label: 'SchoolOS account', secret }).validate({ token: normalizedToken, window: 1 }) !== null;
+}
+
+export function hashRecoveryCode(code: string): string {
+  return createHash('sha256').update(code).digest('hex');
+}
+
+export function createRecoveryCodes(): string[] {
+  return Array.from({ length: RECOVERY_CODE_COUNT }, () => randomBytes(5).toString('hex').toUpperCase());
+}