vulncheck/frontend/app/mfa-setup/page.tsx

"use client";

/**
 * Forced MFA Enrolment.
 *
 * Lands here when /auth/login returned `mfa_setup_required=true` (or
 * when SAML/OIDC callback redirected with `?token=…` because
 * `mfa_enforce_sso_users=true`). User has NO session yet — the
 * setup_token is the only credential they hold for the next 10 min.
 *
 * Two phases:
 *   1. start  → POST /auth/mfa/forced-setup/start  → receive QR + secret
 *   2. verify → POST /auth/mfa/forced-setup/activate {code} → JWT + redirect
 *
 * Token read via window.location instead of useSearchParams to avoid
 * the Next.js 16 Suspense requirement, which made the page hang on
 * the fallback in production.
 */

import { useEffect, useState } from 'react';
import { QRCodeSVG } from 'qrcode.react';
import TotpAlgorithmNote from '@/components/auth/TotpAlgorithmNote';
import api from '../../lib/api';

type Phase = 'starting' | 'verify' | 'done' | 'error';

export default function ForcedMfaSetupPage() {
    const [setupToken, setSetupToken] = useState('');
    const [phase, setPhase] = useState<Phase>('starting');
    const [secret, setSecret] = useState('');
    const [provisioningUri, setProvisioningUri] = useState('');
    const [algorithm, setAlgorithm] = useState('');
    const [code, setCode] = useState('');
    const [error, setError] = useState('');
    const [loading, setLoading] = useState(false);

    // Phase 1 — pull token from URL, fetch QR once.
    // NOTE: Next.js 16 webpack was stripping the IIFE form (the
    // /auth/mfa/forced-setup/start URL literal disappeared from the
    // bundle). Promise chain stays in.
    useEffect(() => {
        if (typeof window === 'undefined') return;
        const params = new URLSearchParams(window.location.search);
        const token = params.get('token') || '';
        if (!token) {
            setError('Missing setup token. Return to the login page.');
            setPhase('error');
            return;
        }
        setSetupToken(token);
        api.post('/auth/mfa/forced-setup/start', { setup_token: token })
            .then((res) => {
                setSecret(res.data.secret);
                setProvisioningUri(res.data.provisioning_uri);
                setAlgorithm(res.data.algorithm);
                setPhase('verify');
            })
            .catch((err: any) => {
                const detail = err.response?.data?.detail;
                setError(detail || 'Could not start MFA enrolment.');
                setPhase('error');
            });
    }, []);

    const handleActivate = async (e: React.FormEvent) => {
        e.preventDefault();
        setError('');
        setLoading(true);
        try {
            const res = await api.post('/auth/mfa/forced-setup/activate', {
                setup_token: setupToken,
                code,
            });
            if (res.data?.access_token) {
                setPhase('done');
                window.location.href = '/';
                return;
            }
            setError('Activation failed.');
        } catch (err: any) {
            const detail = err.response?.data?.detail;
            setError(detail || 'Invalid or expired code.');
        } finally {
            setLoading(false);
        }
    };

    return (
        <div className="flex min-h-screen flex-col justify-center px-6 py-12 lg:px-8 bg-gray-50">
            <div className="sm:mx-auto sm:w-full sm:max-w-md">
                <img src="/logo.svg" alt="TrueVuln Logo" className="mx-auto h-20 w-20" />
                <h2 className="mt-6 text-center text-2xl font-bold leading-9 tracking-tight text-gray-900 font-mono">
                    Multi-Factor Authentication required
                </h2>
                <p className="text-center text-sm text-gray-600 font-mono mt-2">
                    Your administrator requires TOTP for your role. Enrol now to continue.
                </p>
            </div>

            <div className="mt-8 sm:mx-auto sm:w-full sm:max-w-md bg-white p-6 rounded-md shadow-sm border border-gray-200">
                {phase === 'starting' && (
                    <p className="text-center text-sm text-gray-500 font-mono py-8">
                        Generating your authenticator secret…
                    </p>
                )}

                {phase === 'error' && (
                    <div className="space-y-4">
                        <div className="text-red-600 text-sm font-mono text-center bg-red-50 p-3 rounded">
                            {error}
                        </div>
                        <a
                            href="/login"
                            className="block w-full text-center rounded-md border border-gray-300 px-3 py-2 text-sm font-semibold text-gray-700 hover:bg-gray-50 font-mono"
                        >
                            Back to login
                        </a>
                    </div>
                )}

                {phase === 'verify' && (
                    <>
                        <ol className="list-decimal list-inside text-sm text-gray-700 font-mono space-y-1 mb-4">
                            <li>Install an authenticator app (Google Authenticator, Authy, 1Password, …).</li>
                            <li>Scan the QR code or paste the secret below.</li>
                            <li>Enter the first 6-digit code to finish enrolment.</li>
                        </ol>

                        <div className="flex flex-col items-center mb-4">
                            {provisioningUri ? (
                                <div className="border border-gray-200 rounded p-2 bg-white">
                                    <QRCodeSVG value={provisioningUri} size={220} level="M" />
                                </div>
                            ) : null}
                            <p className="mt-3 text-[11px] uppercase tracking-wider text-gray-500 font-mono">Manual entry secret</p>
                            <code className="block mt-1 text-xs bg-gray-100 px-2 py-1 rounded select-all break-all font-mono text-center">
                                {secret}
                            </code>
                        </div>

                        <div className="mb-4"><TotpAlgorithmNote algorithm={algorithm} /></div>

                        <form className="space-y-4" onSubmit={handleActivate}>
                            <div>
                                <label htmlFor="code" className="block text-sm font-medium text-gray-900 font-mono">
                                    6-digit code
                                </label>
                                <input
                                    id="code"
                                    type="text"
                                    inputMode="numeric"
                                    pattern="[0-9]{6}"
                                    maxLength={6}
                                    autoComplete="one-time-code"
                                    autoFocus
                                    required
                                    value={code}
                                    onChange={(e) => setCode(e.target.value.replace(/\D/g, ''))}
                                    className="mt-2 block w-full rounded-md border-0 py-2 text-center text-xl tracking-[0.5em] font-mono text-gray-900 shadow-sm ring-1 ring-inset ring-gray-300 focus:ring-2 focus:ring-inset focus:ring-truevuln-blue"
                                    placeholder="000000"
                                />
                            </div>

                            {error && (
                                <div className="text-red-600 text-sm font-mono text-center bg-red-50 p-2 rounded">
                                    {error}
                                </div>
                            )}

                            <button
                                type="submit"
                                disabled={loading || code.length !== 6}
                                className="w-full rounded-md bg-truevuln-blue px-3 py-2 text-sm font-semibold text-white shadow-sm hover:bg-blue-600 disabled:opacity-60 font-mono"
                            >
                                {loading ? 'Activating…' : 'Activate MFA & Sign in'}
                            </button>

                            <a
                                href="/login"
                                className="block text-center text-xs text-gray-500 hover:text-gray-700 font-mono"
                            >
                                Cancel — back to login
                            </a>
                        </form>
                    </>
                )}

                {phase === 'done' && (
                    <p className="text-center text-sm text-emerald-600 font-mono py-8">
                        MFA enabled. Signing you in…
                    </p>
                )}
            </div>
        </div>
    );
}