jobtrackingapp/JobTrackerApi/Controllers/TwoFactorController.cs

using System.Security.Cryptography;
using System.Text;
using JobTrackerApi.Data;
using JobTrackerApi.Models;
using JobTrackerApi.Services;
using Microsoft.AspNetCore.Authorization;
using Microsoft.AspNetCore.DataProtection;
using Microsoft.AspNetCore.Identity;
using Microsoft.AspNetCore.Mvc;
using Microsoft.AspNetCore.RateLimiting;
using Microsoft.EntityFrameworkCore;
using OtpNet;
using QRCoder;

namespace JobTrackerApi.Controllers;

// TOTP 2FA (RFC 6238) + recovery codes. Split out from AuthController (already 700+ lines)
// rather than growing it further; shares the session cookie logic via AppSessionIssuer and the
// pending-token handoff via ITwoFactorPendingTokenService.
[ApiController]
[Route("api/auth/2fa")]
public sealed class TwoFactorController : ControllerBase
{
    private const int RecoveryCodeCount = 10;

    private readonly UserManager<ApplicationUser> _users;
    private readonly ITokenService _tokens;
    private readonly JobTrackerContext _db;
    private readonly ITwoFactorPendingTokenService _pending;
    private readonly IDataProtector _protector;

    public TwoFactorController(UserManager<ApplicationUser> users, ITokenService tokens, JobTrackerContext db, ITwoFactorPendingTokenService pending, IDataProtectionProvider protectionProvider)
    {
        _users = users;
        _tokens = tokens;
        _db = db;
        _pending = pending;
        _protector = protectionProvider.CreateProtector("totp-secret-v1");
    }

    public sealed record PasswordConfirmRequest(string CurrentPassword);
    public sealed record SetupResult(string ManualEntryKey, string QrCodeDataUrl);
    public sealed record VerifySetupRequest(string Code);
    public sealed record VerifySetupResult(bool Enabled, IReadOnlyList<string> RecoveryCodes);
    public sealed record StatusResult(bool Enabled, DateTimeOffset? EnabledAtUtc);
    public sealed record RecoveryCodesResult(IReadOnlyList<string> RecoveryCodes);
    public sealed record ChallengeRequest(string PendingToken, string Code);

    [HttpPost("setup")]
    [Authorize(AuthenticationSchemes = "local")]
    [EnableRateLimiting("auth-login")]
    public async Task<IActionResult> Setup([FromBody] PasswordConfirmRequest request, CancellationToken cancellationToken)
    {
        var user = await _users.GetUserAsync(User);
        if (user is null) return Unauthorized();

        if (!await _users.CheckPasswordAsync(user, request.CurrentPassword ?? string.Empty))
        {
            return BadRequest("Current password is incorrect.");
        }

        var secretBytes = KeyGeneration.GenerateRandomKey(20);
        var base32Secret = Base32Encoding.ToString(secretBytes);

        user.TotpPendingSecretEncrypted = _protector.Protect(base32Secret);
        var result = await _users.UpdateAsync(user);
        if (!result.Succeeded)
        {
            return BadRequest(string.Join("; ", result.Errors.Select(e => e.Description)));
        }

        var issuer = "JobTracker";
        var label = Uri.EscapeDataString($"{issuer}:{user.Email}");
        var otpauthUri = $"otpauth://totp/{label}?secret={base32Secret}&issuer={Uri.EscapeDataString(issuer)}&digits=6&period=30";

        using var qrGenerator = new QRCodeGenerator();
        using var qrData = qrGenerator.CreateQrCode(otpauthUri, QRCodeGenerator.ECCLevel.Q);
        var pngQr = new PngByteQRCode(qrData);
        var qrPngBytes = pngQr.GetGraphic(10);
        var qrDataUrl = $"data:image/png;base64,{Convert.ToBase64String(qrPngBytes)}";

        return Ok(new SetupResult(base32Secret, qrDataUrl));
    }

    [HttpPost("verify-setup")]
    [Authorize(AuthenticationSchemes = "local")]
    [EnableRateLimiting("auth-login")]
    public async Task<IActionResult> VerifySetup([FromBody] VerifySetupRequest request, CancellationToken cancellationToken)
    {
        var user = await _users.GetUserAsync(User);
        if (user is null) return Unauthorized();

        if (string.IsNullOrWhiteSpace(user.TotpPendingSecretEncrypted))
        {
            return BadRequest("No pending 2FA setup. Call setup first.");
        }

        var base32Secret = _protector.Unprotect(user.TotpPendingSecretEncrypted);
        if (!VerifyCode(base32Secret, request.Code))
        {
            return Unauthorized();
        }

        user.TotpSecretEncrypted = user.TotpPendingSecretEncrypted;
        user.TotpPendingSecretEncrypted = null;
        user.TwoFactorEnabled = true;
        user.TotpEnabledAtUtc = DateTimeOffset.UtcNow;

        var result = await _users.UpdateAsync(user);
        if (!result.Succeeded)
        {
            return BadRequest(string.Join("; ", result.Errors.Select(e => e.Description)));
        }

        var codes = await RegenerateRecoveryCodesAsync(user.Id, cancellationToken);
        return Ok(new VerifySetupResult(true, codes));
    }

    [HttpPost("disable")]
    [Authorize(AuthenticationSchemes = "local")]
    [EnableRateLimiting("auth-login")]
    public async Task<IActionResult> Disable([FromBody] PasswordConfirmRequest request, CancellationToken cancellationToken)
    {
        var user = await _users.GetUserAsync(User);
        if (user is null) return Unauthorized();

        if (!await _users.CheckPasswordAsync(user, request.CurrentPassword ?? string.Empty))
        {
            return BadRequest("Current password is incorrect.");
        }

        user.TotpSecretEncrypted = null;
        user.TotpPendingSecretEncrypted = null;
        user.TwoFactorEnabled = false;
        user.TotpEnabledAtUtc = null;

        var result = await _users.UpdateAsync(user);
        if (!result.Succeeded)
        {
            return BadRequest(string.Join("; ", result.Errors.Select(e => e.Description)));
        }

        await RemoveAllRecoveryCodesAsync(user.Id, cancellationToken);
        return NoContent();
    }

    [HttpGet("status")]
    [Authorize(AuthenticationSchemes = "local")]
    public async Task<IActionResult> Status()
    {
        var user = await _users.GetUserAsync(User);
        if (user is null) return Unauthorized();

        return Ok(new StatusResult(user.TwoFactorEnabled, user.TotpEnabledAtUtc));
    }

    [HttpPost("recovery-codes/regenerate")]
    [Authorize(AuthenticationSchemes = "local")]
    [EnableRateLimiting("auth-login")]
    public async Task<IActionResult> RegenerateRecoveryCodes([FromBody] PasswordConfirmRequest request, CancellationToken cancellationToken)
    {
        var user = await _users.GetUserAsync(User);
        if (user is null) return Unauthorized();

        if (!user.TwoFactorEnabled)
        {
            return BadRequest("Two-factor authentication is not enabled.");
        }

        if (!await _users.CheckPasswordAsync(user, request.CurrentPassword ?? string.Empty))
        {
            return BadRequest("Current password is incorrect.");
        }

        var codes = await RegenerateRecoveryCodesAsync(user.Id, cancellationToken);
        return Ok(new RecoveryCodesResult(codes));
    }

    [HttpPost("challenge")]
    [AllowAnonymous]
    [EnableRateLimiting("auth-2fa-challenge")]
    public async Task<IActionResult> Challenge([FromBody] ChallengeRequest request, CancellationToken cancellationToken)
    {
        var pendingToken = (request.PendingToken ?? string.Empty).Trim();
        var code = (request.Code ?? string.Empty).Trim();
        if (pendingToken.Length == 0 || code.Length == 0) return Unauthorized();

        // Peek without consuming: only burn the pending token once the code actually checks out,
        // so a mistyped code doesn't force the user back through password login.
        var session = _pending.Resolve(pendingToken, consume: false);
        if (session is null) return Unauthorized();

        var user = await _users.FindByIdAsync(session.UserId);
        if (user is null || !user.TwoFactorEnabled || string.IsNullOrWhiteSpace(user.TotpSecretEncrypted))
        {
            return Unauthorized();
        }

        var base32Secret = _protector.Unprotect(user.TotpSecretEncrypted);
        var verified = VerifyCode(base32Secret, code) || await TryConsumeRecoveryCodeAsync(user.Id, code, cancellationToken);
        if (!verified) return Unauthorized();

        _pending.Resolve(pendingToken, consume: true);
        await AppSessionIssuer.IssueAsync(Request, Response, _tokens, user, session.RememberMe, cancellationToken);
        return Ok(new AuthController.AuthSessionResult(true, "local"));
    }

    private static bool VerifyCode(string base32Secret, string? code)
    {
        code = (code ?? string.Empty).Trim();
        if (code.Length == 0) return false;

        var totp = new Totp(Base32Encoding.ToBytes(base32Secret));
        // +-1 step (30s) of drift, the RFC 6238 standard tolerance for clock skew between the
        // authenticator app and the server.
        return totp.VerifyTotp(code, out _, new VerificationWindow(1, 1));
    }

    private async Task<bool> TryConsumeRecoveryCodeAsync(string userId, string code, CancellationToken cancellationToken)
    {
        var hash = HashRecoveryCode(code);
        var match = await _db.TwoFactorRecoveryCodes
            .IgnoreQueryFilters()
            .FirstOrDefaultAsync(x => x.UserId == userId && x.CodeHash == hash && x.UsedAtUtc == null, cancellationToken);
        if (match is null) return false;

        match.UsedAtUtc = DateTimeOffset.UtcNow;
        await _db.SaveChangesAsync(cancellationToken);
        return true;
    }

    private async Task<IReadOnlyList<string>> RegenerateRecoveryCodesAsync(string userId, CancellationToken cancellationToken)
    {
        await RemoveAllRecoveryCodesAsync(userId, cancellationToken);

        var now = DateTimeOffset.UtcNow;
        var plainCodes = new List<string>(RecoveryCodeCount);
        var rows = new List<TwoFactorRecoveryCode>(RecoveryCodeCount);
        for (var i = 0; i < RecoveryCodeCount; i++)
        {
            var plain = GenerateRecoveryCode();
            plainCodes.Add(plain);
            rows.Add(new TwoFactorRecoveryCode { UserId = userId, CodeHash = HashRecoveryCode(plain), CreatedAtUtc = now });
        }

        _db.TwoFactorRecoveryCodes.AddRange(rows);
        await _db.SaveChangesAsync(cancellationToken);
        return plainCodes;
    }

    private async Task RemoveAllRecoveryCodesAsync(string userId, CancellationToken cancellationToken)
    {
        var existing = await _db.TwoFactorRecoveryCodes.IgnoreQueryFilters().Where(x => x.UserId == userId).ToListAsync(cancellationToken);
        if (existing.Count == 0) return;
        _db.TwoFactorRecoveryCodes.RemoveRange(existing);
        await _db.SaveChangesAsync(cancellationToken);
    }

    private static string GenerateRecoveryCode()
    {
        var hex = Convert.ToHexString(RandomNumberGenerator.GetBytes(5)).ToLowerInvariant(); // 10 hex chars, 40 bits
        return $"{hex[..5]}-{hex[5..]}";
    }

    // ponytail: recovery codes are already random high-entropy tokens (not user-chosen
    // passwords), so a plain SHA-256 hash is sufficient -- no per-code salt or PBKDF2 needed.
    private static string HashRecoveryCode(string code)
    {
        var normalized = code.Trim().ToLowerInvariant();
        return Convert.ToHexString(SHA256.HashData(Encoding.UTF8.GetBytes(normalized))).ToLowerInvariant();
    }
}