Phase 8: Registration flow & MFA enforcement UI

- 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 <noreply@anthropic.com>

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),

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() {
+type ScreenMode =
-  const { authStatus, isLoading, login, setup, verifyTotp, mfaRequired, isLoginPending, isSetupPending, isTotpPending } = useAuth();
+  | '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<ScreenMode>('login');
+
+  // ── Lockout (HTTP 423) ──
   const [lockoutMessage, setLockoutMessage] = useState<string | null>(null);
 
-  // Redirect authenticated users immediately
+  // ── MFA enforcement setup flow ──
-  if (!isLoading && authStatus?.authenticated) {
+  const [mfaEnforceStep, setMfaEnforceStep] = useState<MfaEnforceStep>('qr');
+  const [mfaEnforceQr, setMfaEnforceQr] = useState('');
+  const [mfaEnforceSecret, setMfaEnforceSecret] = useState('');
+  const [mfaEnforceBackupCodes, setMfaEnforceBackupCodes] = useState<string[]>([]);
+  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 <Navigate to="/dashboard" replace />;
   }
 
   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 err = validatePassword(password);
-      const validationError = validatePassword(password);
+      if (err) { toast.error(err); return; }
-      if (validationError) {
+      if (password !== confirmPassword) { toast.error('Passwords do not match'); return; }
-        toast.error(validationError);
-        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 {
+      return;
-      // Login mode
+    }
-      try {
+
-        await login({ username, password });
+    try {
-        // If mfaRequired becomes true, the TOTP state renders automatically
+      const result = await login({ username, password });
-        // If not required, useAuth invalidates auth query → Navigate above handles redirect
+      // must_change_password: backend issued session but UI must gate the app
-      } catch (error: any) {
+      if ('must_change_password' in result && result.must_change_password) {
-        if (error?.response?.status === 423) {
+        setMode('force_pw');
-          const msg = error.response.data?.detail || 'Account locked. Try again later.';
-          setLockoutMessage(msg);
-        } else {
-          toast.error(getErrorMessage(error, 'Invalid username or password'));
-        }
       }
+      // 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<TotpSetupResponse>('/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 = () => (
+    <>
+      <CardHeader>
+        <div className="flex items-center gap-3">
+          <div className="p-1.5 rounded-md bg-accent/10">
+            <Lock className="h-4 w-4 text-accent" aria-hidden="true" />
+          </div>
+          <div>
+            <CardTitle>Two-Factor Authentication</CardTitle>
+            <CardDescription>
+              {useBackupCode ? 'Enter one of your backup codes' : 'Enter the code from your authenticator app'}
+            </CardDescription>
+          </div>
+        </div>
+      </CardHeader>
+      <CardContent>
+        <form onSubmit={handleTotpSubmit} className="space-y-4">
+          <div className="space-y-2">
+            <Label htmlFor="totp-code">
+              {useBackupCode ? 'Backup Code' : 'Authenticator Code'}
+            </Label>
+            <Input
+              id="totp-code"
+              type="text"
+              inputMode={useBackupCode ? 'text' : 'numeric'}
+              pattern={useBackupCode ? undefined : '[0-9]*'}
+              maxLength={useBackupCode ? 9 : 6}
+              value={totpCode}
+              onChange={(e) =>
+                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"
+            />
+          </div>
+          <Button type="submit" className="w-full" disabled={isTotpPending}>
+            {isTotpPending ? (
+              <><Loader2 className="h-4 w-4 animate-spin" />Verifying</>
+            ) : (
+              'Verify'
+            )}
+          </Button>
+          <button
+            type="button"
+            onClick={() => { setUseBackupCode(!useBackupCode); setTotpCode(''); }}
+            className="w-full text-center text-xs text-muted-foreground hover:text-foreground transition-colors"
+          >
+            {useBackupCode ? 'Use authenticator app instead' : 'Use a backup code instead'}
+          </button>
+        </form>
+      </CardContent>
+    </>
+  );
+
+  const renderMfaEnforce = () => {
+    // Show a loading/start state if QR hasn't been fetched yet
+    if (!mfaEnforceQr && mfaEnforceStep !== 'backup_codes') {
+      return (
+        <>
+          <CardHeader>
+            <div className="flex items-center gap-3">
+              <div className="p-1.5 rounded-md bg-amber-500/10">
+                <ShieldCheck className="h-4 w-4 text-amber-400" aria-hidden="true" />
+              </div>
+              <div>
+                <CardTitle>Set Up Two-Factor Authentication</CardTitle>
+                <CardDescription>Your account requires MFA before you can continue</CardDescription>
+              </div>
+            </div>
+          </CardHeader>
+          <CardContent>
+            <p className="text-sm text-muted-foreground mb-4">
+              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.
+            </p>
+            <Button className="w-full" onClick={handleMfaEnforceStart} disabled={isMfaEnforceSetupPending}>
+              {isMfaEnforceSetupPending ? (
+                <><Loader2 className="h-4 w-4 animate-spin" />Generating QR Code</>
+              ) : (
+                'Begin Setup'
+              )}
+            </Button>
+          </CardContent>
+        </>
+      );
+    }
+
+    if (mfaEnforceStep === 'backup_codes') {
+      return (
+        <>
+          <CardHeader>
+            <div className="flex items-center gap-3">
+              <div className="p-1.5 rounded-md bg-amber-500/10">
+                <AlertTriangle className="h-4 w-4 text-amber-400" aria-hidden="true" />
+              </div>
+              <div>
+                <CardTitle>Save Your Backup Codes</CardTitle>
+                <CardDescription>Store these somewhere safe — they won't be shown again</CardDescription>
+              </div>
+            </div>
+          </CardHeader>
+          <CardContent className="space-y-4">
+            <p className="text-xs text-muted-foreground">
+              These {mfaEnforceBackupCodes.length} codes can each be used once if you lose access to
+              your authenticator app. MFA is now active on your account.
+            </p>
+            <div className="grid grid-cols-2 gap-2 bg-secondary rounded-md p-3">
+              {mfaEnforceBackupCodes.map((code, i) => (
+                <code key={i} className="text-xs font-mono text-foreground text-center py-0.5">
+                  {code}
+                </code>
+              ))}
+            </div>
+            <Button
+              variant="outline"
+              size="sm"
+              onClick={handleCopyBackupCodes}
+              className="w-full gap-2"
+            >
+              <Copy className="h-4 w-4" />
+              Copy All Codes
+            </Button>
+            <Button
+              className="w-full"
+              onClick={() => {
+                // Session is already issued — redirect to app
+                window.location.href = '/dashboard';
+              }}
+            >
+              I've saved my backup codes — Enter UMBRA
+            </Button>
+          </CardContent>
+        </>
+      );
+    }
+
+    if (mfaEnforceStep === 'qr') {
+      return (
+        <>
+          <CardHeader>
+            <div className="flex items-center gap-3">
+              <div className="p-1.5 rounded-md bg-amber-500/10">
+                <ShieldCheck className="h-4 w-4 text-amber-400" aria-hidden="true" />
+              </div>
+              <div>
+                <CardTitle>Scan QR Code</CardTitle>
+                <CardDescription>Add UMBRA to your authenticator app</CardDescription>
+              </div>
+            </div>
+          </CardHeader>
+          <CardContent className="space-y-4">
+            <div className="flex justify-center">
+              <img
+                src={`data:image/png;base64,${mfaEnforceQr}`}
+                alt="TOTP QR code — scan with your authenticator app"
+                className="h-44 w-44 rounded-md border border-border"
+              />
+            </div>
+            <p className="text-xs text-muted-foreground text-center">
+              Can't scan? Enter this code manually in your app:
+            </p>
+            <code className="block text-center text-xs font-mono bg-secondary px-3 py-2 rounded-md tracking-widest break-all">
+              {mfaEnforceSecret}
+            </code>
+            <Button className="w-full" onClick={() => setMfaEnforceStep('verify')}>
+              Next: Verify Code
+            </Button>
+          </CardContent>
+        </>
+      );
+    }
+
+    // verify step
+    return (
+      <>
+        <CardHeader>
+          <div className="flex items-center gap-3">
+            <div className="p-1.5 rounded-md bg-amber-500/10">
+              <ShieldCheck className="h-4 w-4 text-amber-400" aria-hidden="true" />
+            </div>
+            <div>
+              <CardTitle>Verify Your Authenticator</CardTitle>
+              <CardDescription>Enter the 6-digit code shown in your app</CardDescription>
+            </div>
+          </div>
+        </CardHeader>
+        <CardContent className="space-y-4">
+          <div className="space-y-2">
+            <Label htmlFor="enforce-code">Verification Code</Label>
+            <Input
+              id="enforce-code"
+              type="text"
+              inputMode="numeric"
+              maxLength={6}
+              placeholder="000000"
+              value={mfaEnforceCode}
+              onChange={(e) => 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(); }}
+            />
+          </div>
+          <Button
+            className="w-full"
+            onClick={handleMfaEnforceConfirm}
+            disabled={isMfaEnforceConfirmPending}
+          >
+            {isMfaEnforceConfirmPending ? (
+              <><Loader2 className="h-4 w-4 animate-spin" />Verifying</>
+            ) : (
+              'Verify & Enable MFA'
+            )}
+          </Button>
+          <button
+            type="button"
+            onClick={() => setMfaEnforceStep('qr')}
+            className="w-full text-center text-xs text-muted-foreground hover:text-foreground transition-colors"
+          >
+            Back to QR code
+          </button>
+        </CardContent>
+      </>
+    );
+  };
+
+  const renderLoginOrSetup = () => (
+    <>
+      <CardHeader>
+        <div className="flex items-center gap-3">
+          <div className="p-1.5 rounded-md bg-accent/10">
+            <Lock className="h-4 w-4 text-accent" aria-hidden="true" />
+          </div>
+          <div>
+            <CardTitle>{isSetup ? 'Welcome to UMBRA' : 'Sign in'}</CardTitle>
+            <CardDescription>
+              {isSetup ? 'Create your account to get started' : 'Enter your credentials to continue'}
+            </CardDescription>
+          </div>
+        </div>
+      </CardHeader>
+      <CardContent>
+        {lockoutMessage && (
+          <div
+            role="alert"
+            className={cn(
+              'flex items-center gap-2 rounded-md border border-red-500/30',
+              'bg-red-500/10 px-3 py-2 mb-4'
+            )}
+          >
+            <Lock className="h-4 w-4 text-red-400 shrink-0" aria-hidden="true" />
+            <p className="text-xs text-red-400">{lockoutMessage}</p>
+          </div>
+        )}
+        <form onSubmit={handleCredentialSubmit} className="space-y-4">
+          <div className="space-y-2">
+            <Label htmlFor="username" required>Username</Label>
+            <Input
+              id="username"
+              type="text"
+              value={username}
+              onChange={(e) => { setUsername(e.target.value); setLockoutMessage(null); }}
+              placeholder="Enter username"
+              required
+              autoFocus
+              autoComplete="username"
+            />
+          </div>
+          <div className="space-y-2">
+            <Label htmlFor="password" required>Password</Label>
+            <Input
+              id="password"
+              type="password"
+              value={password}
+              onChange={(e) => { setPassword(e.target.value); setLockoutMessage(null); }}
+              placeholder={isSetup ? 'Create a password' : 'Enter password'}
+              required
+              autoComplete={isSetup ? 'new-password' : 'current-password'}
+            />
+          </div>
+          {isSetup && (
+            <div className="space-y-2">
+              <Label htmlFor="confirm-password" required>Confirm Password</Label>
+              <Input
+                id="confirm-password"
+                type="password"
+                value={confirmPassword}
+                onChange={(e) => setConfirmPassword(e.target.value)}
+                placeholder="Confirm your password"
+                required
+                autoComplete="new-password"
+              />
+              <p className="text-xs text-muted-foreground">
+                Must be 12-128 characters with at least one letter and one non-letter.
+              </p>
+            </div>
+          )}
+          <Button
+            type="submit"
+            className="w-full"
+            disabled={isLoginPending || isSetupPending || !!lockoutMessage}
+          >
+            {isLoginPending || isSetupPending ? (
+              <><Loader2 className="h-4 w-4 animate-spin" />Please wait</>
+            ) : isSetup ? (
+              'Create Account'
+            ) : (
+              'Sign in'
+            )}
+          </Button>
+        </form>
+
+        {/* Open registration link — only shown on login screen when enabled */}
+        {!isSetup && registrationOpen && (
+          <div className="mt-4 text-center">
+            <button
+              type="button"
+              onClick={() => {
+                setMode('register');
+                setUsername('');
+                setPassword('');
+                setConfirmPassword('');
+                setLockoutMessage(null);
+              }}
+              className="text-xs text-muted-foreground hover:text-foreground transition-colors"
+            >
+              Don't have an account?{' '}
+              <span className="text-accent hover:underline">Create one</span>
+            </button>
+          </div>
+        )}
+      </CardContent>
+    </>
+  );
+
+  const renderRegister = () => (
+    <>
+      <CardHeader>
+        <div className="flex items-center gap-3">
+          <div className="p-1.5 rounded-md bg-accent/10">
+            <UserPlus className="h-4 w-4 text-accent" aria-hidden="true" />
+          </div>
+          <div>
+            <CardTitle>Create Account</CardTitle>
+            <CardDescription>Register for access to UMBRA</CardDescription>
+          </div>
+        </div>
+      </CardHeader>
+      <CardContent>
+        <form onSubmit={handleRegisterSubmit} className="space-y-4">
+          <div className="space-y-2">
+            <Label htmlFor="reg-username" required>Username</Label>
+            <Input
+              id="reg-username"
+              type="text"
+              value={username}
+              onChange={(e) => setUsername(e.target.value)}
+              placeholder="Choose a username"
+              required
+              autoFocus
+              autoComplete="username"
+            />
+          </div>
+          <div className="space-y-2">
+            <Label htmlFor="reg-password" required>Password</Label>
+            <Input
+              id="reg-password"
+              type="password"
+              value={password}
+              onChange={(e) => setPassword(e.target.value)}
+              placeholder="Create a password"
+              required
+              autoComplete="new-password"
+            />
+          </div>
+          <div className="space-y-2">
+            <Label htmlFor="reg-confirm-password" required>Confirm Password</Label>
+            <Input
+              id="reg-confirm-password"
+              type="password"
+              value={confirmPassword}
+              onChange={(e) => setConfirmPassword(e.target.value)}
+              placeholder="Confirm your password"
+              required
+              autoComplete="new-password"
+            />
+            <p className="text-xs text-muted-foreground">
+              Must be 12-128 characters with at least one letter and one non-letter.
+            </p>
+          </div>
+          <Button type="submit" className="w-full" disabled={isRegisterPending}>
+            {isRegisterPending ? (
+              <><Loader2 className="h-4 w-4 animate-spin" />Creating account</>
+            ) : (
+              'Create Account'
+            )}
+          </Button>
+        </form>
+        <div className="mt-4 text-center">
+          <button
+            type="button"
+            onClick={() => {
+              setMode('login');
+              setUsername('');
+              setPassword('');
+              setConfirmPassword('');
+            }}
+            className="text-xs text-muted-foreground hover:text-foreground transition-colors"
+          >
+            Already have an account?{' '}
+            <span className="text-accent hover:underline">Sign in</span>
+          </button>
+        </div>
+      </CardContent>
+    </>
+  );
+
+  const renderForcedPasswordChange = () => (
+    <>
+      <CardHeader>
+        <div className="flex items-center gap-3">
+          <div className="p-1.5 rounded-md bg-amber-500/10">
+            <Lock className="h-4 w-4 text-amber-400" aria-hidden="true" />
+          </div>
+          <div>
+            <CardTitle>Password Change Required</CardTitle>
+            <CardDescription>An administrator has reset your password. Please set a new one.</CardDescription>
+          </div>
+        </div>
+      </CardHeader>
+      <CardContent>
+        <form onSubmit={handleForcePwSubmit} className="space-y-4">
+          <div className="space-y-2">
+            <Label htmlFor="force-new-pw" required>New Password</Label>
+            <Input
+              id="force-new-pw"
+              type="password"
+              value={forcedNewPassword}
+              onChange={(e) => setForcedNewPassword(e.target.value)}
+              placeholder="Create a new password"
+              required
+              autoFocus
+              autoComplete="new-password"
+            />
+          </div>
+          <div className="space-y-2">
+            <Label htmlFor="force-confirm-pw" required>Confirm New Password</Label>
+            <Input
+              id="force-confirm-pw"
+              type="password"
+              value={forcedConfirmPassword}
+              onChange={(e) => setForcedConfirmPassword(e.target.value)}
+              placeholder="Confirm your new password"
+              required
+              autoComplete="new-password"
+            />
+            <p className="text-xs text-muted-foreground">
+              Must be 12-128 characters with at least one letter and one non-letter.
+            </p>
+          </div>
+          <Button type="submit" className="w-full" disabled={isForcePwPending}>
+            {isForcePwPending ? (
+              <><Loader2 className="h-4 w-4 animate-spin" />Saving</>
+            ) : (
+              'Set New Password'
+            )}
+          </Button>
+        </form>
+      </CardContent>
+    </>
+  );
+
   return (
     <div className="relative flex min-h-screen flex-col items-center justify-center bg-background p-4 overflow-hidden">
       <AmbientBackground />
 
-      {/* Wordmark — in flex flow above card */}
+      {/* Wordmark */}
       <span className="font-heading text-5xl sm:text-6xl font-bold tracking-tight text-accent mb-10 relative z-10 animate-slide-up">
         UMBRA
       </span>
 
       {/* Auth card */}
       <Card className="w-full max-w-sm relative z-10 border-border/80 animate-slide-up">
-        {mfaRequired ? (
+        {activeMode === 'totp' && renderTotpChallenge()}
-          // State C: TOTP challenge
+        {activeMode === 'mfa_enforce' && renderMfaEnforce()}
-          <>
+        {activeMode === 'force_pw' && renderForcedPasswordChange()}
-            <CardHeader>
+        {activeMode === 'register' && renderRegister()}
-              <div className="flex items-center gap-3">
+        {(activeMode === 'login' || activeMode === 'setup') && renderLoginOrSetup()}
-                <div className="p-1.5 rounded-md bg-accent/10">
-                  <Lock className="h-4 w-4 text-accent" aria-hidden="true" />
-                </div>
-                <div>
-                  <CardTitle>Two-Factor Authentication</CardTitle>
-                  <CardDescription>
-                    {useBackupCode
-                      ? 'Enter one of your backup codes'
-                      : 'Enter the code from your authenticator app'}
-                  </CardDescription>
-                </div>
-              </div>
-            </CardHeader>
-            <CardContent>
-              <form onSubmit={handleTotpSubmit} className="space-y-4">
-                <div className="space-y-2">
-                  <Label htmlFor="totp-code">
-                    {useBackupCode ? 'Backup Code' : 'Authenticator Code'}
-                  </Label>
-                  <Input
-                    id="totp-code"
-                    type="text"
-                    inputMode={useBackupCode ? 'text' : 'numeric'}
-                    pattern={useBackupCode ? undefined : '[0-9]*'}
-                    maxLength={useBackupCode ? 9 : 6}
-                    value={totpCode}
-                    onChange={(e) =>
-                      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"
-                  />
-                </div>
-                <Button type="submit" className="w-full" disabled={isTotpPending}>
-                  {isTotpPending ? (
-                    <>
-                      <Loader2 className="h-4 w-4 animate-spin" />
-                      Verifying
-                    </>
-                  ) : (
-                    'Verify'
-                  )}
-                </Button>
-                <button
-                  type="button"
-                  onClick={() => {
-                    setUseBackupCode(!useBackupCode);
-                    setTotpCode('');
-                  }}
-                  className="w-full text-center text-xs text-muted-foreground hover:text-foreground transition-colors"
-                >
-                  {useBackupCode ? 'Use authenticator app instead' : 'Use a backup code instead'}
-                </button>
-              </form>
-            </CardContent>
-          </>
-        ) : (
-          // State A (setup) or State B (login)
-          <>
-            <CardHeader>
-              <div className="flex items-center gap-3">
-                <div className="p-1.5 rounded-md bg-accent/10">
-                  <Lock className="h-4 w-4 text-accent" aria-hidden="true" />
-                </div>
-                <div>
-                  <CardTitle>{isSetup ? 'Welcome to UMBRA' : 'Sign in'}</CardTitle>
-                  <CardDescription>
-                    {isSetup
-                      ? 'Create your account to get started'
-                      : 'Enter your credentials to continue'}
-                  </CardDescription>
-                </div>
-              </div>
-            </CardHeader>
-            <CardContent>
-              {/* Lockout warning banner */}
-              {lockoutMessage && (
-                <div
-                  role="alert"
-                  className={cn(
-                    'flex items-center gap-2 rounded-md border border-red-500/30',
-                    'bg-red-500/10 px-3 py-2 mb-4'
-                  )}
-                >
-                  <Lock className="h-4 w-4 text-red-400 shrink-0" aria-hidden="true" />
-                  <p className="text-xs text-red-400">{lockoutMessage}</p>
-                </div>
-              )}
-
-              <form onSubmit={handleCredentialSubmit} className="space-y-4">
-                <div className="space-y-2">
-                  <Label htmlFor="username" required>Username</Label>
-                  <Input
-                    id="username"
-                    type="text"
-                    value={username}
-                    onChange={(e) => { setUsername(e.target.value); setLockoutMessage(null); }}
-                    placeholder="Enter username"
-                    required
-                    autoFocus
-                    autoComplete="username"
-                  />
-                </div>
-
-                <div className="space-y-2">
-                  <Label htmlFor="password" required>Password</Label>
-                  <Input
-                    id="password"
-                    type="password"
-                    value={password}
-                    onChange={(e) => { setPassword(e.target.value); setLockoutMessage(null); }}
-                    placeholder={isSetup ? 'Create a password' : 'Enter password'}
-                    required
-                    autoComplete={isSetup ? 'new-password' : 'current-password'}
-                  />
-                </div>
-
-                {isSetup && (
-                  <div className="space-y-2">
-                    <Label htmlFor="confirm-password" required>Confirm Password</Label>
-                    <Input
-                      id="confirm-password"
-                      type="password"
-                      value={confirmPassword}
-                      onChange={(e) => setConfirmPassword(e.target.value)}
-                      placeholder="Confirm your password"
-                      required
-                      autoComplete="new-password"
-                    />
-                    <p className="text-xs text-muted-foreground">
-                      Must be 12-128 characters with at least one letter and one non-letter.
-                    </p>
-                  </div>
-                )}
-
-                <Button
-                  type="submit"
-                  className="w-full"
-                  disabled={isLoginPending || isSetupPending || !!lockoutMessage}
-                >
-                  {isLoginPending || isSetupPending ? (
-                    <>
-                      <Loader2 className="h-4 w-4 animate-spin" />
-                      Please wait
-                    </>
-                  ) : isSetup ? (
-                    'Create Account'
-                  ) : (
-                    'Sign in'
-                  )}
-                </Button>
-              </form>
-            </CardContent>
-          </>
-        )}
       </Card>
     </div>
   );