PowderCoatingLogix/src/PowderCoating.Web/Controllers/TwoFactorSetupController.cs

using Microsoft.AspNetCore.Authorization;
using Microsoft.AspNetCore.Identity;
using Microsoft.AspNetCore.Mvc;
using PowderCoating.Core.Entities;
using QRCoder;

namespace PowderCoating.Web.Controllers;

/// <summary>
/// TOTP-based 2FA setup and management for SuperAdmin accounts.
/// Accessible to any authenticated user — SuperAdmins are redirected here
/// by EnforceSuperAdmin2FAFilter when 2FA is not yet configured.
/// </summary>
[Authorize]
public class TwoFactorSetupController : Controller
{
    private readonly UserManager<ApplicationUser> _userManager;
    private readonly SignInManager<ApplicationUser> _signInManager;
    private readonly IHostEnvironment _env;

    public TwoFactorSetupController(
        UserManager<ApplicationUser> userManager,
        SignInManager<ApplicationUser> signInManager,
        IHostEnvironment env)
    {
        _userManager = userManager;
        _signInManager = signInManager;
        _env = env;
    }

    /// <summary>
    /// Returns the TOTP issuer name embedded in the authenticator URI.
    /// Non-production environments append the environment name (e.g. "[Development]")
    /// so that QR codes scanned against a staging server show up separately in the
    /// authenticator app and cannot be accidentally used against production.
    /// </summary>
    private string AppName => _env.IsProduction()
        ? "Powder Coating Logix"
        : $"Powder Coating Logix [{_env.EnvironmentName}]";

    /// <summary>
    /// Renders the 2FA management overview page showing the current user's
    /// two-factor status, whether an authenticator key has been registered, and
    /// whether they hold the SuperAdmin role (which affects which actions are available).
    /// </summary>
    // GET: /TwoFactorSetup
    public async Task<IActionResult> Index()
    {
        var user = await _userManager.GetUserAsync(User);
        if (user == null) return Challenge();

        var isSuperAdmin = await _userManager.IsInRoleAsync(user, "SuperAdmin");

        ViewBag.IsSuperAdmin = isSuperAdmin;
        ViewBag.TwoFactorEnabled = user.TwoFactorEnabled;
        ViewBag.HasAuthenticator = await _userManager.GetAuthenticatorKeyAsync(user) != null;

        return View();
    }

    /// <summary>
    /// Returns the Setup page with a QR code and manual entry key for the user's
    /// authenticator app.  Only generates a new secret key if one does not already
    /// exist — this prevents the key from being silently reset when background
    /// notification-poll requests follow a filter redirect back to this page,
    /// which would invalidate a code the user just scanned.
    /// </summary>
    // GET: /TwoFactorSetup/Setup
    public async Task<IActionResult> Setup()
    {
        var user = await _userManager.GetUserAsync(User);
        if (user == null) return Challenge();

        // Only generate a new key if none exists — avoids resetting the key when
        // background fetch requests (e.g. notification polling) follow filter redirects
        // back to this page and silently invalidate the code the user just scanned.
        var key = await _userManager.GetAuthenticatorKeyAsync(user);
        if (string.IsNullOrEmpty(key))
        {
            await _userManager.ResetAuthenticatorKeyAsync(user);
            key = await _userManager.GetAuthenticatorKeyAsync(user) ?? string.Empty;
        }

        var appName = AppName;
        var email = user.Email ?? user.UserName ?? "user";
        var totpUri = $"otpauth://totp/{Uri.EscapeDataString(appName)}:{Uri.EscapeDataString(email)}?secret={key}&issuer={Uri.EscapeDataString(appName)}&algorithm=SHA1&digits=6&period=30";

        var qrBase64 = GenerateQrBase64(totpUri);

        ViewBag.SharedKey = FormatKey(key);
        ViewBag.QrCodeBase64 = qrBase64;
        ViewBag.TotpUri = totpUri;

        return View();
    }

    /// <summary>
    /// Verifies the TOTP code entered by the user and, on success, enables
    /// two-factor authentication on their account.  Strips spaces and dashes before
    /// verification because authenticator apps often display codes with separators.
    /// Calls <c>RefreshSignInAsync</c> so the authentication cookie reflects the
    /// updated 2FA state without requiring a new login.
    /// </summary>
    // POST: /TwoFactorSetup/Setup
    [HttpPost, ValidateAntiForgeryToken]
    public async Task<IActionResult> Setup(string verificationCode)
    {
        var user = await _userManager.GetUserAsync(User);
        if (user == null) return Challenge();

        var code = verificationCode?.Replace(" ", "").Replace("-", "") ?? string.Empty;
        var isValid = await _userManager.VerifyTwoFactorTokenAsync(
            user, _userManager.Options.Tokens.AuthenticatorTokenProvider, code);

        if (!isValid)
        {
            // Re-show the setup page with current key (don't reset again)
            var key = await _userManager.GetAuthenticatorKeyAsync(user) ?? string.Empty;
            var appName = AppName;
            var email = user.Email ?? user.UserName ?? "user";
            var totpUri = $"otpauth://totp/{Uri.EscapeDataString(appName)}:{Uri.EscapeDataString(email)}?secret={key}&issuer={Uri.EscapeDataString(appName)}&algorithm=SHA1&digits=6&period=30";

            ViewBag.SharedKey = FormatKey(key);
            ViewBag.QrCodeBase64 = GenerateQrBase64(totpUri);
            ViewBag.TotpUri = totpUri;
            ViewBag.Error = "Verification code was incorrect. Make sure your authenticator app is synced and try again.";
            return View();
        }

        await _userManager.SetTwoFactorEnabledAsync(user, true);
        await _signInManager.RefreshSignInAsync(user);

        TempData["Success"] = "Two-factor authentication has been enabled on your account.";
        return RedirectToAction(nameof(Index));
    }

    /// <summary>
    /// Disables two-factor authentication and resets the authenticator key.
    /// SuperAdmins must supply a valid current TOTP code before 2FA can be
    /// removed — this prevents an attacker with a hijacked session from weakening
    /// a privileged account's security posture without physical access to the device.
    /// Non-SuperAdmin users can disable 2FA without a code (standard Identity behaviour).
    /// The authenticator key is reset (not just disabled) so that a re-enabled key
    /// requires fresh scanning of a new QR code.
    /// </summary>
    // POST: /TwoFactorSetup/Disable
    [HttpPost, ValidateAntiForgeryToken]
    public async Task<IActionResult> Disable(string confirmationCode)
    {
        var user = await _userManager.GetUserAsync(User);
        if (user == null) return Challenge();

        var isSuperAdmin = await _userManager.IsInRoleAsync(user, "SuperAdmin");

        // SuperAdmins cannot disable 2FA without a valid TOTP code
        if (isSuperAdmin)
        {
            var code = confirmationCode?.Replace(" ", "").Replace("-", "") ?? string.Empty;
            var isValid = await _userManager.VerifyTwoFactorTokenAsync(
                user, _userManager.Options.Tokens.AuthenticatorTokenProvider, code);

            if (!isValid)
            {
                TempData["Error"] = "Verification code was incorrect. 2FA has not been disabled.";
                return RedirectToAction(nameof(Index));
            }
        }

        await _userManager.SetTwoFactorEnabledAsync(user, false);
        await _userManager.ResetAuthenticatorKeyAsync(user);
        await _signInManager.RefreshSignInAsync(user);

        TempData["Success"] = "Two-factor authentication has been disabled.";
        return RedirectToAction(nameof(Index));
    }

    // ── Helpers ──────────────────────────────────────────────────────────────

    /// <summary>
    /// Generates a PNG QR code for the given <paramref name="content"/> string
    /// and returns it as a Base64-encoded string suitable for embedding in an
    /// <c>&lt;img src="data:image/png;base64,…"&gt;</c> element.
    /// Uses error-correction level Q (≈25 % recovery capacity) — a balanced
    /// choice that keeps the code scannable even if partially obscured, without
    /// making the image too dense for typical phone cameras.
    /// </summary>
    private static string GenerateQrBase64(string content)
    {
        using var qrGenerator = new QRCodeGenerator();
        var qrData = qrGenerator.CreateQrCode(content, QRCodeGenerator.ECCLevel.Q);
        using var qrCode = new PngByteQRCode(qrData);
        var bytes = qrCode.GetGraphic(6);
        return Convert.ToBase64String(bytes);
    }

    /// <summary>
    /// Formats a raw Base32 authenticator key into space-separated groups of four
    /// uppercase characters (e.g. <c>ABCD EFGH IJKL …</c>) to improve readability
    /// for users who prefer to type the key manually instead of scanning the QR code.
    /// </summary>
    private static string FormatKey(string key)
    {
        // Group into chunks of 4 for readability: XXXX XXXX XXXX ...
        var sb = new System.Text.StringBuilder();
        for (int i = 0; i < key.Length; i++)
        {
            if (i > 0 && i % 4 == 0) sb.Append(' ');
            sb.Append(char.ToUpperInvariant(key[i]));
        }
        return sb.ToString();
    }
}