import React, { useCallback, useState } from 'react'; import { useSetRecoilState } from 'recoil'; import { SmartphoneIcon } from 'lucide-react'; import { motion, AnimatePresence } from 'framer-motion'; import { OGDialog, useToastContext, OGDialogContent, OGDialogHeader, OGDialogTitle, Progress, } from '@librechat/client'; import type { TUser, TVerify2FARequest } from 'librechat-data-provider'; import { useConfirmTwoFactorMutation, useDisableTwoFactorMutation, useEnableTwoFactorMutation, useVerifyTwoFactorMutation, } from '~/data-provider'; import { SetupPhase, QRPhase, VerifyPhase, BackupPhase, DisablePhase } from './TwoFactorPhases'; import { DisableTwoFactorToggle } from './DisableTwoFactorToggle'; import { useAuthContext, useLocalize } from '~/hooks'; import store from '~/store'; export type Phase = 'setup' | 'qr' | 'verify' | 'backup' | 'disable'; const phaseVariants = { initial: { opacity: 0, scale: 0.95 }, animate: { opacity: 1, scale: 1, transition: { duration: 0.3, ease: 'easeOut' } }, exit: { opacity: 0, scale: 0.95, transition: { duration: 0.3, ease: 'easeIn' } }, }; const TwoFactorAuthentication: React.FC = () => { const localize = useLocalize(); const { user } = useAuthContext(); const setUser = useSetRecoilState(store.user); const { showToast } = useToastContext(); const [secret, setSecret] = useState(''); const [otpauthUrl, setOtpauthUrl] = useState(''); const [downloaded, setDownloaded] = useState(false); const [backupCodes, setBackupCodes] = useState([]); const [_disableToken, setDisableToken] = useState(''); const [isDialogOpen, setDialogOpen] = useState(false); const [verificationToken, setVerificationToken] = useState(''); const [phase, setPhase] = useState(user?.twoFactorEnabled ? 'disable' : 'setup'); const { mutate: confirm2FAMutate } = useConfirmTwoFactorMutation(); const { mutate: enable2FAMutate, isLoading: isGenerating } = useEnableTwoFactorMutation(); const { mutate: verify2FAMutate, isLoading: isVerifying } = useVerifyTwoFactorMutation(); const { mutate: disable2FAMutate, isLoading: isDisabling } = useDisableTwoFactorMutation(); const steps = ['Setup', 'Scan QR', 'Verify', 'Backup']; const phasesLabel: Record = { setup: 'Setup', qr: 'Scan QR', verify: 'Verify', backup: 'Backup', disable: '', }; const currentStep = steps.indexOf(phasesLabel[phase]); const resetState = useCallback(() => { if (user?.twoFactorEnabled && otpauthUrl) { disable2FAMutate(undefined, { onError: () => showToast({ message: localize('com_ui_2fa_disable_error'), status: 'error' }), }); } setOtpauthUrl(''); setSecret(''); setBackupCodes([]); setVerificationToken(''); setDisableToken(''); setPhase(user?.twoFactorEnabled ? 'disable' : 'setup'); setDownloaded(false); }, [user, otpauthUrl, disable2FAMutate, localize, showToast]); const handleGenerateQRCode = useCallback(() => { enable2FAMutate(undefined, { onSuccess: ({ otpauthUrl, backupCodes }) => { setOtpauthUrl(otpauthUrl); setSecret(otpauthUrl.split('secret=')[1].split('&')[0]); setBackupCodes(backupCodes); setPhase('qr'); }, onError: () => showToast({ message: localize('com_ui_2fa_generate_error'), status: 'error' }), }); }, [enable2FAMutate, localize, showToast]); const handleVerify = useCallback(() => { if (!verificationToken) { return; } verify2FAMutate( { token: verificationToken }, { onSuccess: () => { showToast({ message: localize('com_ui_2fa_verified') }); confirm2FAMutate( { token: verificationToken }, { onSuccess: () => setPhase('backup'), onError: () => showToast({ message: localize('com_ui_2fa_invalid'), status: 'error' }), }, ); }, onError: () => showToast({ message: localize('com_ui_2fa_invalid'), status: 'error' }), }, ); }, [verificationToken, verify2FAMutate, confirm2FAMutate, localize, showToast]); const handleDownload = useCallback(() => { if (!backupCodes.length) { return; } const blob = new Blob([backupCodes.join('\n')], { type: 'text/plain;charset=utf-8' }); const url = URL.createObjectURL(blob); const a = document.createElement('a'); a.href = url; a.download = 'backup-codes.txt'; a.click(); URL.revokeObjectURL(url); setDownloaded(true); }, [backupCodes]); const handleConfirm = useCallback(() => { setDialogOpen(false); setPhase('disable'); showToast({ message: localize('com_ui_2fa_enabled') }); setUser( (prev) => ({ ...prev, backupCodes: backupCodes.map((code) => ({ code, codeHash: code, used: false, usedAt: null, })), twoFactorEnabled: true, }) as TUser, ); }, [setUser, localize, showToast, backupCodes]); const handleDisableVerify = useCallback( (token: string, useBackup: boolean) => { // Validate: if not using backup, ensure token has at least 6 digits; // if using backup, ensure backup code has at least 8 characters. if (!useBackup && token.trim().length < 6) { return; } if (useBackup && token.trim().length < 8) { return; } const payload: TVerify2FARequest = {}; if (useBackup) { payload.backupCode = token.trim(); } else { payload.token = token.trim(); } disable2FAMutate(payload, { onSuccess: () => { showToast({ message: localize('com_ui_2fa_disabled') }); setDialogOpen(false); setUser( (prev) => ({ ...prev, totpSecret: '', backupCodes: [], twoFactorEnabled: false, }) as TUser, ); setPhase('setup'); setOtpauthUrl(''); }, onError: () => showToast({ message: localize('com_ui_2fa_invalid'), status: 'error' }), }); }, [disable2FAMutate, showToast, localize, setUser], ); return ( { setDialogOpen(open); if (!open) { resetState(); } }} > setDialogOpen(true)} disabled={isVerifying || isDisabling || isGenerating} /> {user?.twoFactorEnabled && phase !== 'disable' && (
{steps.map((step, index) => ( = index ? 'var(--text-primary)' : 'var(--text-tertiary)', }} className="font-medium" > {step} ))}
)}
{phase === 'setup' && ( setPhase('qr')} onError={(error) => showToast({ message: error.message, status: 'error' })} /> )} {phase === 'qr' && ( setPhase('verify')} onError={(error) => showToast({ message: error.message, status: 'error' })} /> )} {phase === 'verify' && ( showToast({ message: error.message, status: 'error' })} /> )} {phase === 'backup' && ( showToast({ message: error.message, status: 'error' })} /> )} {phase === 'disable' && ( showToast({ message: error.message, status: 'error' })} /> )}
); }; export default React.memo(TwoFactorAuthentication);