From 464b8b911fa092093dcbbe1845a5cc52a8b79cdf Mon Sep 17 00:00:00 2001 From: Kyle Pope Date: Thu, 26 Feb 2026 18:39:18 +0800 Subject: [PATCH] Phase 8: Registration flow & MFA enforcement UI MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - backend: add POST /auth/totp/enforce-setup and /auth/totp/enforce-confirm endpoints that operate on mfa_enforce_token (not session cookie), generate TOTP secret/QR/backup codes, verify confirmation code, enable TOTP, clear mfa_enforce_pending flag, and issue a full session cookie - frontend: expand LockScreen to five modes — login, first-run setup, open registration, TOTP challenge, MFA enforcement setup (QR -> verify -> backup codes), and forced password change; all modes share AmbientBackground and the existing card layout; registration visible only when authStatus.registration_open is true Co-Authored-By: Claude Opus 4.6 --- backend/app/routers/totp.py | 112 +++ frontend/src/components/auth/LockScreen.tsx | 854 +++++++++++++++----- 2 files changed, 765 insertions(+), 201 deletions(-) diff --git a/backend/app/routers/totp.py b/backend/app/routers/totp.py index e8b9be7..e147892 100644 --- a/backend/app/routers/totp.py +++ b/backend/app/routers/totp.py @@ -39,6 +39,7 @@ from app.services.auth import ( verify_password_with_upgrade, hash_password, verify_mfa_token, + verify_mfa_enforce_token, create_session_token, ) from app.services.totp import ( @@ -94,6 +95,15 @@ class BackupCodesRegenerateRequest(BaseModel): code: str # Current TOTP code required to regenerate +class EnforceSetupRequest(BaseModel): + mfa_token: str + + +class EnforceConfirmRequest(BaseModel): + mfa_token: str + code: str # 6-digit TOTP code from authenticator app + + # --------------------------------------------------------------------------- # Internal helpers # --------------------------------------------------------------------------- @@ -394,6 +404,108 @@ async def regenerate_backup_codes( return {"backup_codes": plaintext_codes} +@router.post("/totp/enforce-setup") +async def enforce_setup_totp( + data: EnforceSetupRequest, + db: AsyncSession = Depends(get_db), +): + """ + Generate TOTP secret + QR code + backup codes during MFA enforcement. + + Called after login returns mfa_setup_required=True. Uses the mfa_enforce_token + (not a session cookie) because the user is not yet fully authenticated. + + Idempotent: regenerates secret if called again before confirm. + Returns { secret, qr_code_base64, backup_codes }. + """ + user_id = verify_mfa_enforce_token(data.mfa_token) + if user_id is None: + raise HTTPException(status_code=401, detail="Invalid or expired enforcement token — please log in again") + + result = await db.execute(select(User).where(User.id == user_id, User.is_active == True)) + user = result.scalar_one_or_none() + if not user: + raise HTTPException(status_code=401, detail="User not found or inactive") + + if not user.mfa_enforce_pending: + raise HTTPException(status_code=400, detail="MFA enforcement is not pending for this account") + + if user.totp_enabled: + raise HTTPException(status_code=400, detail="TOTP is already enabled for this account") + + # Generate new secret (idempotent — overwrite any unconfirmed secret) + raw_secret = generate_totp_secret() + encrypted_secret = encrypt_totp_secret(raw_secret) + user.totp_secret = encrypted_secret + user.totp_enabled = False # Not enabled until enforce-confirm called + + # Generate backup codes — hash before storage, return plaintext once + plaintext_codes = generate_backup_codes(10) + await _delete_backup_codes(db, user.id) + await _store_backup_codes(db, user.id, plaintext_codes) + + await db.commit() + + uri = get_totp_uri(encrypted_secret, user.username) + qr_base64 = generate_qr_base64(uri) + + return { + "secret": raw_secret, + "qr_code_base64": qr_base64, + "backup_codes": plaintext_codes, + } + + +@router.post("/totp/enforce-confirm") +async def enforce_confirm_totp( + data: EnforceConfirmRequest, + request: Request, + response: Response, + db: AsyncSession = Depends(get_db), +): + """ + Confirm TOTP setup during enforcement, clear the pending flag, issue a full session. + + Must be called after /totp/enforce-setup while totp_enabled is still False. + On success: enables TOTP, clears mfa_enforce_pending, sets session cookie, + returns { authenticated: true }. + """ + user_id = verify_mfa_enforce_token(data.mfa_token) + if user_id is None: + raise HTTPException(status_code=401, detail="Invalid or expired enforcement token — please log in again") + + result = await db.execute(select(User).where(User.id == user_id, User.is_active == True)) + user = result.scalar_one_or_none() + if not user: + raise HTTPException(status_code=401, detail="User not found or inactive") + + if not user.mfa_enforce_pending: + raise HTTPException(status_code=400, detail="MFA enforcement is not pending for this account") + + if not user.totp_secret: + raise HTTPException(status_code=400, detail="TOTP setup not started — call /totp/enforce-setup first") + + if user.totp_enabled: + raise HTTPException(status_code=400, detail="TOTP is already enabled") + + # Verify the confirmation code + matched_window = verify_totp_code(user.totp_secret, data.code) + if matched_window is None: + raise HTTPException(status_code=400, detail="Invalid code — check your authenticator app time sync") + + # Enable TOTP and clear the enforcement flag + user.totp_enabled = True + user.mfa_enforce_pending = False + user.last_login_at = datetime.now() + await db.commit() + + # Issue a full session + token = await _create_full_session(db, user, request) + _set_session_cookie(response, token) + + return {"authenticated": True} + + @router.get("/totp/status") async def totp_status( db: AsyncSession = Depends(get_db), diff --git a/frontend/src/components/auth/LockScreen.tsx b/frontend/src/components/auth/LockScreen.tsx index 961574a..1885512 100644 --- a/frontend/src/components/auth/LockScreen.tsx +++ b/frontend/src/components/auth/LockScreen.tsx @@ -1,15 +1,16 @@ import { useState, FormEvent } from 'react'; import { Navigate } from 'react-router-dom'; import { toast } from 'sonner'; -import { Lock, Loader2 } from 'lucide-react'; +import { AlertTriangle, Copy, Lock, Loader2, ShieldCheck, UserPlus } from 'lucide-react'; import { useAuth } from '@/hooks/useAuth'; -import { getErrorMessage } from '@/lib/api'; +import api, { getErrorMessage } from '@/lib/api'; import { Card, CardContent, CardDescription, CardHeader, CardTitle } from '@/components/ui/card'; import { Input } from '@/components/ui/input'; import { Button } from '@/components/ui/button'; import { Label } from '@/components/ui/label'; import { cn } from '@/lib/utils'; import AmbientBackground from './AmbientBackground'; +import type { TotpSetupResponse } from '@/types'; /** Validates password against backend rules: 12-128 chars, at least one letter + one non-letter. */ function validatePassword(password: string): string | null { @@ -20,63 +21,124 @@ function validatePassword(password: string): string | null { return null; } -export default function LockScreen() { - const { authStatus, isLoading, login, setup, verifyTotp, mfaRequired, isLoginPending, isSetupPending, isTotpPending } = useAuth(); +type ScreenMode = + | 'login' + | 'setup' // first-run admin account creation + | 'register' // open registration + | 'totp' // TOTP challenge after login + | 'mfa_enforce' // forced MFA setup after login/register + | 'force_pw'; // admin-forced password change - // Credentials state (shared across login/setup states) +type MfaEnforceStep = 'qr' | 'verify' | 'backup_codes'; + +export default function LockScreen() { + const { + authStatus, + isLoading, + login, + register, + setup, + verifyTotp, + mfaRequired, + mfaSetupRequired, + mfaToken, + isLoginPending, + isRegisterPending, + isSetupPending, + isTotpPending, + } = useAuth(); + + // ── Shared credential fields ── const [username, setUsername] = useState(''); const [password, setPassword] = useState(''); const [confirmPassword, setConfirmPassword] = useState(''); - // TOTP challenge state + // ── TOTP challenge ── const [totpCode, setTotpCode] = useState(''); const [useBackupCode, setUseBackupCode] = useState(false); - // Lockout handling (HTTP 423) + // ── Registration mode ── + const [mode, setMode] = useState('login'); + + // ── Lockout (HTTP 423) ── const [lockoutMessage, setLockoutMessage] = useState(null); - // Redirect authenticated users immediately - if (!isLoading && authStatus?.authenticated) { + // ── MFA enforcement setup flow ── + const [mfaEnforceStep, setMfaEnforceStep] = useState('qr'); + const [mfaEnforceQr, setMfaEnforceQr] = useState(''); + const [mfaEnforceSecret, setMfaEnforceSecret] = useState(''); + const [mfaEnforceBackupCodes, setMfaEnforceBackupCodes] = useState([]); + const [mfaEnforceCode, setMfaEnforceCode] = useState(''); + const [isMfaEnforceSetupPending, setIsMfaEnforceSetupPending] = useState(false); + const [isMfaEnforceConfirmPending, setIsMfaEnforceConfirmPending] = useState(false); + + // ── Forced password change ── + const [forcedNewPassword, setForcedNewPassword] = useState(''); + const [forcedConfirmPassword, setForcedConfirmPassword] = useState(''); + const [isForcePwPending, setIsForcePwPending] = useState(false); + + // Redirect authenticated users (no pending MFA flows) + if (!isLoading && authStatus?.authenticated && !mfaSetupRequired && mode !== 'force_pw') { return ; } const isSetup = authStatus?.setup_required === true; + const registrationOpen = authStatus?.registration_open === true; + + // Derive active screen — hook-driven states override local mode + const activeMode: ScreenMode = mfaRequired + ? 'totp' + : mfaSetupRequired + ? 'mfa_enforce' + : isSetup + ? 'setup' + : mode; + + // ── Handlers ── const handleCredentialSubmit = async (e: FormEvent) => { e.preventDefault(); setLockoutMessage(null); if (isSetup) { - // Setup mode: validate password then create account - const validationError = validatePassword(password); - if (validationError) { - toast.error(validationError); - return; - } - if (password !== confirmPassword) { - toast.error('Passwords do not match'); - return; - } + const err = validatePassword(password); + if (err) { toast.error(err); return; } + if (password !== confirmPassword) { toast.error('Passwords do not match'); return; } try { await setup({ username, password }); - // useAuth invalidates auth query → Navigate above handles redirect } catch (error) { toast.error(getErrorMessage(error, 'Failed to create account')); } - } else { - // Login mode - try { - await login({ username, password }); - // If mfaRequired becomes true, the TOTP state renders automatically - // If not required, useAuth invalidates auth query → Navigate above handles redirect - } catch (error: any) { - if (error?.response?.status === 423) { - const msg = error.response.data?.detail || 'Account locked. Try again later.'; - setLockoutMessage(msg); - } else { - toast.error(getErrorMessage(error, 'Invalid username or password')); - } + return; + } + + try { + const result = await login({ username, password }); + // must_change_password: backend issued session but UI must gate the app + if ('must_change_password' in result && result.must_change_password) { + setMode('force_pw'); } + // mfaSetupRequired / mfaRequired handled by hook state → activeMode switches automatically + } catch (error: any) { + if (error?.response?.status === 423) { + setLockoutMessage(error.response.data?.detail || 'Account locked. Try again later.'); + } else { + toast.error(getErrorMessage(error, 'Invalid username or password')); + } + } + }; + + const handleRegisterSubmit = async (e: FormEvent) => { + e.preventDefault(); + const err = validatePassword(password); + if (err) { toast.error(err); return; } + if (password !== confirmPassword) { toast.error('Passwords do not match'); return; } + try { + await register({ username, password }); + // On success useAuth invalidates query → Navigate handles redirect + // If mfa_setup_required the hook sets mfaSetupRequired → activeMode switches + } catch (error) { + toast.error(getErrorMessage(error, 'Registration failed')); } }; @@ -84,190 +146,580 @@ export default function LockScreen() { e.preventDefault(); try { await verifyTotp(totpCode); - // useAuth invalidates auth query → Navigate above handles redirect } catch (error) { toast.error(getErrorMessage(error, 'Invalid verification code')); setTotpCode(''); } }; + const handleMfaEnforceStart = async () => { + if (!mfaToken) return; + setIsMfaEnforceSetupPending(true); + try { + const { data } = await api.post('/auth/totp/enforce-setup', { + mfa_token: mfaToken, + }); + setMfaEnforceQr(data.qr_code_base64); + setMfaEnforceSecret(data.secret); + setMfaEnforceBackupCodes(data.backup_codes); + setMfaEnforceStep('qr'); + } catch (error) { + toast.error(getErrorMessage(error, 'Failed to begin MFA setup')); + } finally { + setIsMfaEnforceSetupPending(false); + } + }; + + const handleMfaEnforceConfirm = async () => { + if (!mfaToken || !mfaEnforceCode || mfaEnforceCode.length !== 6) { + toast.error('Enter a 6-digit code from your authenticator app'); + return; + } + setIsMfaEnforceConfirmPending(true); + try { + await api.post('/auth/totp/enforce-confirm', { + mfa_token: mfaToken, + code: mfaEnforceCode, + }); + // Backend issued session — show backup codes then redirect + setMfaEnforceStep('backup_codes'); + } catch (error) { + toast.error(getErrorMessage(error, 'Invalid code — try again')); + setMfaEnforceCode(''); + } finally { + setIsMfaEnforceConfirmPending(false); + } + }; + + const handleCopyBackupCodes = async () => { + try { + await navigator.clipboard.writeText(mfaEnforceBackupCodes.join('\n')); + toast.success('Backup codes copied'); + } catch { + toast.error('Failed to copy — please select and copy manually'); + } + }; + + const handleForcePwSubmit = async (e: FormEvent) => { + e.preventDefault(); + const err = validatePassword(forcedNewPassword); + if (err) { toast.error(err); return; } + if (forcedNewPassword !== forcedConfirmPassword) { + toast.error('Passwords do not match'); + return; + } + setIsForcePwPending(true); + try { + await api.post('/auth/change-password', { + old_password: password, // retained from original login submission + new_password: forcedNewPassword, + }); + toast.success('Password updated — welcome to UMBRA'); + // Auth query still has authenticated:true → Navigate will fire after re-render + setMode('login'); + } catch (error) { + toast.error(getErrorMessage(error, 'Failed to change password')); + } finally { + setIsForcePwPending(false); + } + }; + + // ── Render helpers ── + + const renderTotpChallenge = () => ( + <> + +
+
+
+
+ Two-Factor Authentication + + {useBackupCode ? 'Enter one of your backup codes' : 'Enter the code from your authenticator app'} + +
+
+ + +
+
+ + + setTotpCode( + useBackupCode + ? e.target.value.replace(/[^0-9-]/g, '') + : e.target.value.replace(/\D/g, '') + ) + } + placeholder={useBackupCode ? 'XXXX-XXXX' : '000000'} + autoFocus + autoComplete="one-time-code" + className="text-center text-lg tracking-widest" + /> +
+ + +
+ + + ); + + const renderMfaEnforce = () => { + // Show a loading/start state if QR hasn't been fetched yet + if (!mfaEnforceQr && mfaEnforceStep !== 'backup_codes') { + return ( + <> + +
+
+
+
+ Set Up Two-Factor Authentication + Your account requires MFA before you can continue +
+
+ + +

+ An administrator has required that your account be protected with an authenticator app. + You'll need an app like Google Authenticator, Authy, or 1Password to continue. +

+ + + + ); + } + + if (mfaEnforceStep === 'backup_codes') { + return ( + <> + +
+
+
+
+ Save Your Backup Codes + Store these somewhere safe — they won't be shown again +
+
+ + +

+ These {mfaEnforceBackupCodes.length} codes can each be used once if you lose access to + your authenticator app. MFA is now active on your account. +

+
+ {mfaEnforceBackupCodes.map((code, i) => ( + + {code} + + ))} +
+ + + + + ); + } + + if (mfaEnforceStep === 'qr') { + return ( + <> + +
+
+
+
+ Scan QR Code + Add UMBRA to your authenticator app +
+
+ + +
+ TOTP QR code — scan with your authenticator app +
+

+ Can't scan? Enter this code manually in your app: +

+ + {mfaEnforceSecret} + + + + + ); + } + + // verify step + return ( + <> + +
+
+
+
+ Verify Your Authenticator + Enter the 6-digit code shown in your app +
+
+ + +
+ + setMfaEnforceCode(e.target.value.replace(/\D/g, ''))} + className="text-center tracking-widest text-lg" + autoFocus + autoComplete="one-time-code" + onKeyDown={(e) => { if (e.key === 'Enter') handleMfaEnforceConfirm(); }} + /> +
+ + + + + ); + }; + + const renderLoginOrSetup = () => ( + <> + +
+
+
+
+ {isSetup ? 'Welcome to UMBRA' : 'Sign in'} + + {isSetup ? 'Create your account to get started' : 'Enter your credentials to continue'} + +
+
+ + + {lockoutMessage && ( +
+
+ )} +
+
+ + { setUsername(e.target.value); setLockoutMessage(null); }} + placeholder="Enter username" + required + autoFocus + autoComplete="username" + /> +
+
+ + { setPassword(e.target.value); setLockoutMessage(null); }} + placeholder={isSetup ? 'Create a password' : 'Enter password'} + required + autoComplete={isSetup ? 'new-password' : 'current-password'} + /> +
+ {isSetup && ( +
+ + setConfirmPassword(e.target.value)} + placeholder="Confirm your password" + required + autoComplete="new-password" + /> +

+ Must be 12-128 characters with at least one letter and one non-letter. +

+
+ )} + +
+ + {/* Open registration link — only shown on login screen when enabled */} + {!isSetup && registrationOpen && ( +
+ +
+ )} + + + ); + + const renderRegister = () => ( + <> + +
+
+
+
+ Create Account + Register for access to UMBRA +
+
+ + +
+
+ + setUsername(e.target.value)} + placeholder="Choose a username" + required + autoFocus + autoComplete="username" + /> +
+
+ + setPassword(e.target.value)} + placeholder="Create a password" + required + autoComplete="new-password" + /> +
+
+ + setConfirmPassword(e.target.value)} + placeholder="Confirm your password" + required + autoComplete="new-password" + /> +

+ Must be 12-128 characters with at least one letter and one non-letter. +

+
+ +
+
+ +
+ + + ); + + const renderForcedPasswordChange = () => ( + <> + +
+
+
+
+ Password Change Required + An administrator has reset your password. Please set a new one. +
+
+ + +
+
+ + setForcedNewPassword(e.target.value)} + placeholder="Create a new password" + required + autoFocus + autoComplete="new-password" + /> +
+
+ + setForcedConfirmPassword(e.target.value)} + placeholder="Confirm your new password" + required + autoComplete="new-password" + /> +

+ Must be 12-128 characters with at least one letter and one non-letter. +

+
+ +
+ + + ); + return (
- {/* Wordmark — in flex flow above card */} + {/* Wordmark */} UMBRA {/* Auth card */} - {mfaRequired ? ( - // State C: TOTP challenge - <> - -
-
-
-
- Two-Factor Authentication - - {useBackupCode - ? 'Enter one of your backup codes' - : 'Enter the code from your authenticator app'} - -
-
- - -
-
- - - setTotpCode( - useBackupCode - ? e.target.value.replace(/[^0-9-]/g, '') - : e.target.value.replace(/\D/g, '') - ) - } - placeholder={useBackupCode ? 'XXXX-XXXX' : '000000'} - autoFocus - autoComplete="one-time-code" - className="text-center text-lg tracking-widest" - /> -
- - -
- - - ) : ( - // State A (setup) or State B (login) - <> - -
-
-
-
- {isSetup ? 'Welcome to UMBRA' : 'Sign in'} - - {isSetup - ? 'Create your account to get started' - : 'Enter your credentials to continue'} - -
-
- - - {/* Lockout warning banner */} - {lockoutMessage && ( -
-
- )} - -
-
- - { setUsername(e.target.value); setLockoutMessage(null); }} - placeholder="Enter username" - required - autoFocus - autoComplete="username" - /> -
- -
- - { setPassword(e.target.value); setLockoutMessage(null); }} - placeholder={isSetup ? 'Create a password' : 'Enter password'} - required - autoComplete={isSetup ? 'new-password' : 'current-password'} - /> -
- - {isSetup && ( -
- - setConfirmPassword(e.target.value)} - placeholder="Confirm your password" - required - autoComplete="new-password" - /> -

- Must be 12-128 characters with at least one letter and one non-letter. -

-
- )} - - -
- - - )} + {activeMode === 'totp' && renderTotpChallenge()} + {activeMode === 'mfa_enforce' && renderMfaEnforce()} + {activeMode === 'force_pw' && renderForcedPasswordChange()} + {activeMode === 'register' && renderRegister()} + {(activeMode === 'login' || activeMode === 'setup') && renderLoginOrSetup()}
);