From 27d3340a371198af8cdcd9a55c934b62601c2f66 Mon Sep 17 00:00:00 2001 From: benya Date: Sun, 8 Mar 2026 11:43:51 +0300 Subject: [PATCH] feat(auth): add TOTP 2FA setup and login verification - add user twofa fields and migration - add 2FA setup/enable/disable endpoints - enforce OTP on login when 2FA enabled - add web login OTP field and settings UI --- alembic/versions/0018_user_twofa.py | 28 +++++++++ app/auth/router.py | 34 +++++++++++ app/auth/schemas.py | 11 ++++ app/auth/service.py | 43 +++++++++++++ app/users/models.py | 2 + app/users/schemas.py | 1 + app/utils/totp.py | 44 ++++++++++++++ web/src/api/auth.ts | 22 ++++++- web/src/chat/types.ts | 1 + web/src/components/AuthPanel.tsx | 11 +++- web/src/components/SettingsPanel.tsx | 91 +++++++++++++++++++++++++++- web/src/store/authStore.ts | 6 +- 12 files changed, 287 insertions(+), 7 deletions(-) create mode 100644 alembic/versions/0018_user_twofa.py create mode 100644 app/utils/totp.py diff --git a/alembic/versions/0018_user_twofa.py b/alembic/versions/0018_user_twofa.py new file mode 100644 index 0000000..bcf0848 --- /dev/null +++ b/alembic/versions/0018_user_twofa.py @@ -0,0 +1,28 @@ +"""add user twofa fields + +Revision ID: 0018_user_twofa +Revises: 0017_user_contacts +Create Date: 2026-03-08 23:35:00.000000 +""" + +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa + + +revision: str = "0018_user_twofa" +down_revision: Union[str, Sequence[str], None] = "0017_user_contacts" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + op.add_column("users", sa.Column("twofa_enabled", sa.Boolean(), nullable=False, server_default=sa.text("false"))) + op.add_column("users", sa.Column("twofa_secret", sa.String(length=64), nullable=True)) + + +def downgrade() -> None: + op.drop_column("users", "twofa_secret") + op.drop_column("users", "twofa_enabled") + diff --git a/app/auth/router.py b/app/auth/router.py index 28193db..eb3f194 100644 --- a/app/auth/router.py +++ b/app/auth/router.py @@ -14,9 +14,13 @@ from app.auth.schemas import ( ResetPasswordRequest, TokenResponse, SessionRead, + TwoFactorCodeRequest, + TwoFactorSetupRead, VerifyEmailRequest, ) from app.auth.service import ( + disable_twofa, + enable_twofa, get_current_user, get_email_sender, get_request_metadata, @@ -29,6 +33,7 @@ from app.auth.service import ( request_password_reset, resend_verification_email, reset_password, + setup_twofa, verify_email, ) from app.database.session import get_db @@ -155,3 +160,32 @@ async def revoke_session(jti: str, current_user: User = Depends(get_current_user @router.delete("/sessions", status_code=status.HTTP_204_NO_CONTENT) async def revoke_all_sessions(current_user: User = Depends(get_current_user)) -> None: await revoke_all_user_sessions(user_id=current_user.id) + + +@router.post("/2fa/setup", response_model=TwoFactorSetupRead) +async def setup_2fa( + db: AsyncSession = Depends(get_db), + current_user: User = Depends(get_current_user), +) -> TwoFactorSetupRead: + secret, otpauth_url = await setup_twofa(db, current_user) + return TwoFactorSetupRead(secret=secret, otpauth_url=otpauth_url) + + +@router.post("/2fa/enable", response_model=MessageResponse) +async def enable_2fa( + payload: TwoFactorCodeRequest, + db: AsyncSession = Depends(get_db), + current_user: User = Depends(get_current_user), +) -> MessageResponse: + await enable_twofa(db, current_user, code=payload.code) + return MessageResponse(message="2FA enabled") + + +@router.post("/2fa/disable", response_model=MessageResponse) +async def disable_2fa( + payload: TwoFactorCodeRequest, + db: AsyncSession = Depends(get_db), + current_user: User = Depends(get_current_user), +) -> MessageResponse: + await disable_twofa(db, current_user, code=payload.code) + return MessageResponse(message="2FA disabled") diff --git a/app/auth/schemas.py b/app/auth/schemas.py index 94310d3..fac71a3 100644 --- a/app/auth/schemas.py +++ b/app/auth/schemas.py @@ -13,6 +13,7 @@ class RegisterRequest(BaseModel): class LoginRequest(BaseModel): email: EmailStr password: str = Field(min_length=8, max_length=128) + otp_code: str | None = Field(default=None, min_length=6, max_length=8) class RefreshTokenRequest(BaseModel): @@ -56,6 +57,7 @@ class AuthUserResponse(BaseModel): bio: str | None = None avatar_url: str | None = None email_verified: bool + twofa_enabled: bool created_at: datetime updated_at: datetime @@ -65,3 +67,12 @@ class SessionRead(BaseModel): created_at: datetime ip_address: str | None = None user_agent: str | None = None + + +class TwoFactorSetupRead(BaseModel): + secret: str + otpauth_url: str + + +class TwoFactorCodeRequest(BaseModel): + code: str = Field(min_length=6, max_length=8) diff --git a/app/auth/service.py b/app/auth/service.py index aeb2010..341973d 100644 --- a/app/auth/service.py +++ b/app/auth/service.py @@ -30,6 +30,7 @@ from app.database.session import get_db from app.email.service import EmailDeliveryError, EmailService, get_email_service from app.users.models import User from app.users.repository import create_user, get_user_by_email, get_user_by_id, get_user_by_username +from app.utils.totp import build_otpauth_uri, generate_totp_secret, verify_totp_code from app.utils.security import ( create_access_token, create_refresh_token, @@ -132,6 +133,11 @@ async def login_user( if not user.email_verified: raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="Email not verified") + if user.twofa_enabled: + if not payload.otp_code: + raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="2FA code required") + if not user.twofa_secret or not verify_totp_code(user.twofa_secret, payload.otp_code): + raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid 2FA code") refresh_jti = str(uuid4()) refresh_token = create_refresh_token(str(user.id), jti=refresh_jti) @@ -215,6 +221,43 @@ def get_request_metadata(request: Request) -> tuple[str | None, str | None]: return ip_address, user_agent +async def setup_twofa(db: AsyncSession, user: User) -> tuple[str, str]: + if user.twofa_enabled and user.twofa_secret: + secret = user.twofa_secret + else: + secret = generate_totp_secret() + user.twofa_secret = secret + await db.commit() + await db.refresh(user) + otpauth_url = build_otpauth_uri(secret=secret, account_name=user.email, issuer=settings.app_name) + return secret, otpauth_url + + +async def enable_twofa(db: AsyncSession, user: User, *, code: str) -> None: + if not user.twofa_secret: + raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="2FA is not initialized") + if not verify_totp_code(user.twofa_secret, code): + raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Invalid 2FA code") + user.twofa_enabled = True + await db.commit() + await db.refresh(user) + + +async def disable_twofa(db: AsyncSession, user: User, *, code: str) -> None: + if not user.twofa_enabled or not user.twofa_secret: + user.twofa_enabled = False + user.twofa_secret = None + await db.commit() + await db.refresh(user) + return + if not verify_totp_code(user.twofa_secret, code): + raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Invalid 2FA code") + user.twofa_enabled = False + user.twofa_secret = None + await db.commit() + await db.refresh(user) + + async def request_password_reset( db: AsyncSession, payload: RequestPasswordResetRequest, diff --git a/app/users/models.py b/app/users/models.py index 3bae010..feb5673 100644 --- a/app/users/models.py +++ b/app/users/models.py @@ -24,6 +24,8 @@ class User(Base): bio: Mapped[str | None] = mapped_column(String(500), nullable=True) email_verified: Mapped[bool] = mapped_column(Boolean, default=False, nullable=False, index=True) allow_private_messages: Mapped[bool] = mapped_column(Boolean, default=True, nullable=False, server_default="true") + twofa_enabled: Mapped[bool] = mapped_column(Boolean, default=False, nullable=False, server_default="false") + twofa_secret: Mapped[str | None] = mapped_column(String(64), nullable=True) created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), server_default=func.now(), nullable=False) updated_at: Mapped[datetime] = mapped_column( DateTime(timezone=True), diff --git a/app/users/schemas.py b/app/users/schemas.py index e3fc2c9..8e18a7c 100644 --- a/app/users/schemas.py +++ b/app/users/schemas.py @@ -21,6 +21,7 @@ class UserRead(UserBase): bio: str | None = None email_verified: bool allow_private_messages: bool + twofa_enabled: bool = False created_at: datetime updated_at: datetime diff --git a/app/utils/totp.py b/app/utils/totp.py new file mode 100644 index 0000000..8e0fffc --- /dev/null +++ b/app/utils/totp.py @@ -0,0 +1,44 @@ +import base64 +import hashlib +import hmac +import secrets +import struct +import time +import urllib.parse + + +def generate_totp_secret(length: int = 20) -> str: + raw = secrets.token_bytes(length) + return base64.b32encode(raw).decode("ascii").rstrip("=") + + +def build_otpauth_uri(*, secret: str, account_name: str, issuer: str) -> str: + label = urllib.parse.quote(f"{issuer}:{account_name}") + issuer_q = urllib.parse.quote(issuer) + secret_q = urllib.parse.quote(secret) + return f"otpauth://totp/{label}?secret={secret_q}&issuer={issuer_q}&algorithm=SHA1&digits=6&period=30" + + +def _totp_code(secret: str, for_time: int | None = None, step: int = 30, digits: int = 6) -> str: + now = int(time.time()) if for_time is None else int(for_time) + counter = now // step + padded = secret + "=" * ((8 - len(secret) % 8) % 8) + key = base64.b32decode(padded, casefold=True) + msg = struct.pack(">Q", counter) + digest = hmac.new(key, msg, hashlib.sha1).digest() + offset = digest[-1] & 0x0F + binary = struct.unpack(">I", digest[offset : offset + 4])[0] & 0x7FFFFFFF + return str(binary % (10**digits)).zfill(digits) + + +def verify_totp_code(secret: str, code: str, *, window: int = 1, step: int = 30, digits: int = 6) -> bool: + prepared = "".join(ch for ch in code if ch.isdigit()) + if len(prepared) != digits: + return False + now = int(time.time()) + for delta in range(-window, window + 1): + expected = _totp_code(secret, for_time=now + delta * step, step=step, digits=digits) + if secrets.compare_digest(expected, prepared): + return True + return False + diff --git a/web/src/api/auth.ts b/web/src/api/auth.ts index 5a85094..f6d2b82 100644 --- a/web/src/api/auth.ts +++ b/web/src/api/auth.ts @@ -5,8 +5,8 @@ export async function registerRequest(email: string, name: string, username: str await http.post("/auth/register", { email, name, username, password }); } -export async function loginRequest(email: string, password: string): Promise { - const { data } = await http.post("/auth/login", { email, password }); +export async function loginRequest(email: string, password: string, otpCode?: string): Promise { + const { data } = await http.post("/auth/login", { email, password, otp_code: otpCode || undefined }); return data; } @@ -32,3 +32,21 @@ export async function revokeSession(jti: string): Promise { export async function revokeAllSessions(): Promise { await http.delete("/auth/sessions"); } + +export interface TwoFactorSetupResponse { + secret: string; + otpauth_url: string; +} + +export async function setupTwoFactor(): Promise { + const { data } = await http.post("/auth/2fa/setup"); + return data; +} + +export async function enableTwoFactor(code: string): Promise { + await http.post("/auth/2fa/enable", { code }); +} + +export async function disableTwoFactor(code: string): Promise { + await http.post("/auth/2fa/disable", { code }); +} diff --git a/web/src/chat/types.ts b/web/src/chat/types.ts index 587d5cf..0c6854f 100644 --- a/web/src/chat/types.ts +++ b/web/src/chat/types.ts @@ -77,6 +77,7 @@ export interface AuthUser { bio?: string | null; avatar_url: string | null; email_verified: boolean; + twofa_enabled?: boolean; allow_private_messages: boolean; created_at: string; updated_at: string; diff --git a/web/src/components/AuthPanel.tsx b/web/src/components/AuthPanel.tsx index 3f50c07..d71bafe 100644 --- a/web/src/components/AuthPanel.tsx +++ b/web/src/components/AuthPanel.tsx @@ -12,6 +12,7 @@ export function AuthPanel() { const [name, setName] = useState(""); const [username, setUsername] = useState(""); const [password, setPassword] = useState(""); + const [otpCode, setOtpCode] = useState(""); const [error, setError] = useState(null); const [success, setSuccess] = useState(null); @@ -26,7 +27,7 @@ export function AuthPanel() { setMode("login"); return; } - await login(email, password); + await login(email, password, otpCode.trim() || undefined); } catch { setError("Auth request failed."); } @@ -51,6 +52,14 @@ export function AuthPanel() { )} setPassword(e.target.value)} /> + {mode === "login" ? ( + setOtpCode(e.target.value.replace(/\D/g, "").slice(0, 8))} + /> + ) : null} diff --git a/web/src/components/SettingsPanel.tsx b/web/src/components/SettingsPanel.tsx index e6a843e..9c38c17 100644 --- a/web/src/components/SettingsPanel.tsx +++ b/web/src/components/SettingsPanel.tsx @@ -1,7 +1,7 @@ import { useEffect, useMemo, useState } from "react"; import { createPortal } from "react-dom"; import { listBlockedUsers, updateMyProfile } from "../api/users"; -import { listSessions, revokeAllSessions, revokeSession } from "../api/auth"; +import { disableTwoFactor, enableTwoFactor, listSessions, revokeAllSessions, revokeSession, setupTwoFactor } from "../api/auth"; import type { AuthSession, AuthUser } from "../chat/types"; import { useAuthStore } from "../store/authStore"; import { AppPreferences, getAppPreferences, updateAppPreferences } from "../utils/preferences"; @@ -22,6 +22,10 @@ export function SettingsPanel({ open, onClose }: Props) { const [allowPrivateMessages, setAllowPrivateMessages] = useState(true); const [sessions, setSessions] = useState([]); const [sessionsLoading, setSessionsLoading] = useState(false); + const [twofaCode, setTwofaCode] = useState(""); + const [twofaSecret, setTwofaSecret] = useState(null); + const [twofaUrl, setTwofaUrl] = useState(null); + const [twofaError, setTwofaError] = useState(null); const [profileDraft, setProfileDraft] = useState({ name: "", username: "", @@ -285,6 +289,91 @@ export function SettingsPanel({ open, onClose }: Props) { disabled={savingPrivacy} /> +
+
+

Two-Factor Authentication

+ {me.twofa_enabled ? "Enabled" : "Disabled"} +
+ {!me.twofa_enabled ? ( + <> + + {twofaSecret ? ( +
+

Secret

+

{twofaSecret}

+ {twofaUrl ?

{twofaUrl}

: null} +
+ ) : null} +
+ setTwofaCode(e.target.value.replace(/\D/g, "").slice(0, 8))} + /> + +
+ + ) : ( +
+ setTwofaCode(e.target.value.replace(/\D/g, "").slice(0, 8))} + /> + +
+ )} + {twofaError ?

{twofaError}

: null} +

Active Sessions

diff --git a/web/src/store/authStore.ts b/web/src/store/authStore.ts index 358e92c..142ece0 100644 --- a/web/src/store/authStore.ts +++ b/web/src/store/authStore.ts @@ -8,7 +8,7 @@ interface AuthState { me: AuthUser | null; loading: boolean; setTokens: (accessToken: string, refreshToken: string) => void; - login: (email: string, password: string) => Promise; + login: (email: string, password: string, otpCode?: string) => Promise; loadMe: () => Promise; refresh: () => Promise; logout: () => void; @@ -27,10 +27,10 @@ export const useAuthStore = create((set, get) => ({ localStorage.setItem(REFRESH_KEY, refreshToken); set({ accessToken, refreshToken }); }, - login: async (email, password) => { + login: async (email, password, otpCode) => { set({ loading: true }); try { - const data = await loginRequest(email, password); + const data = await loginRequest(email, password, otpCode); get().setTokens(data.access_token, data.refresh_token); await get().loadMe(); } finally {