admin/Azaion.Services/MfaService.cs

using System.IdentityModel.Tokens.Jwt;
using System.Security.Claims;
using System.Security.Cryptography;
using System.Text.Json;
using Azaion.Common;
using Azaion.Common.Configs;
using Azaion.Common.Database;
using Azaion.Common.Entities;
using Azaion.Common.Requests;
using LinqToDB;
using LinqToDB.Data;
using Microsoft.AspNetCore.DataProtection;
using Microsoft.Extensions.Options;
using Microsoft.IdentityModel.Tokens;
using OtpNet;
using QRCoder;

namespace Azaion.Services;

/// <summary>
/// AZ-534 — RFC 6238 TOTP enrollment + login validation, with single-use recovery codes.
/// MfaSecret is encrypted at rest via <see cref="IDataProtector"/>; recovery codes are
/// stored as SHA-256 hashes (high-entropy secrets need a fast hash, not Argon2id —
/// same reasoning the refresh-token store uses).
/// </summary>
public interface IMfaService
{
    Task<MfaEnrollResponse> Enroll(Guid userId, string password, CancellationToken ct = default);
    Task                    Confirm(Guid userId, string code, CancellationToken ct = default);
    Task                    Disable(Guid userId, string password, string code, CancellationToken ct = default);

    /// <summary>
    /// Issued at /login when the user has MFA enabled — a 5-minute JWT (aud=azaion-mfa-step2)
    /// the client carries to /login/mfa for the second-factor verification.
    /// </summary>
    string                  IssueMfaStepToken(Guid userId);

    /// <summary>
    /// Decode the step-1 token, returning the userId. Throws BusinessException(InvalidMfaToken)
    /// on bad signature, audience mismatch, or expired token.
    /// </summary>
    Guid                    ValidateMfaStepToken(string token);

    /// <summary>
    /// AZ-534 AC-3 + AC-4 — second-factor verification at login. Returns the
    /// <c>amr</c> values the access token should carry (always includes <c>"pwd"</c>
    /// and <c>"mfa"</c>; <c>"recovery"</c> is added when a recovery code was used).
    /// </summary>
    Task<string[]>          VerifyForLogin(Guid userId, string code, CancellationToken ct = default);
}

public class MfaService(
    IDbFactory dbFactory,
    IUserService userService,
    IDataProtectionProvider dataProtectionProvider,
    IJwtSigningKeyProvider signingKeys,
    IOptions<JwtConfig> jwtConfig,
    IOptions<AuthConfig> authConfig,
    IAuditLog auditLog) : IMfaService
{
    private const string MfaSecretPurpose = "Azaion.Mfa.Secret.v1";
    private const string MfaStepAudience  = "azaion-mfa-step2";
    private const int    MfaStepLifetimeSeconds = 300; // 5 min — matches AC-3
    private const int    SecretBytes      = 20;        // 160 bits — RFC 6238 §3
    private const int    RecoveryCodeCount = 10;
    private const int    RecoveryCodeBytes = 10;       // base32(10) = 16 chars (≥12 per AC-1)

    private readonly IDataProtector _protector = dataProtectionProvider.CreateProtector(MfaSecretPurpose);
    private readonly JwtConfig      _jwt       = jwtConfig.Value;
    private readonly AuthConfig     _auth      = authConfig.Value;

    public async Task<MfaEnrollResponse> Enroll(Guid userId, string password, CancellationToken ct = default)
    {
        var user = await userService.GetById(userId, ct)
            ?? throw new BusinessException(ExceptionEnum.NoEmailFound);

        // Re-auth with password — AC-1 requires this to defend a stolen access token
        // from being usable to silently flip the user into MFA.
        var verify = Security.VerifyPassword(password, user.PasswordHash);
        if (!verify.Valid)
            throw new BusinessException(ExceptionEnum.WrongPassword);

        if (user.MfaEnabled)
            throw new BusinessException(ExceptionEnum.MfaAlreadyEnabled);

        var secretBytes  = KeyGeneration.GenerateRandomKey(SecretBytes);
        var secretBase32 = Base32Encoding.ToString(secretBytes); // 32 chars (per AC-1)

        var otpAuthUrl = new OtpUri(
            schema:    OtpType.Totp,
            secret:    secretBase32,
            user:      user.Email,
            issuer:    _jwt.Issuer,
            algorithm: OtpHashMode.Sha1,
            digits:    6,
            period:    30).ToString();

        var qrPng = GenerateQrPng(otpAuthUrl);

        var recoveryPlain = new string[RecoveryCodeCount];
        var recoveryStore = new RecoveryCodeStore[RecoveryCodeCount];
        for (var i = 0; i < RecoveryCodeCount; i++)
        {
            var code = Base32Encoding.ToString(KeyGeneration.GenerateRandomKey(RecoveryCodeBytes));
            recoveryPlain[i] = code;
            recoveryStore[i] = new RecoveryCodeStore { Hash = HashRecoveryCode(code), UsedAt = null };
        }

        var encryptedSecret = _protector.Protect(secretBase32);
        var recoveryJson    = JsonSerializer.Serialize(recoveryStore);

        await dbFactory.RunAdmin(async db =>
            await db.Users.UpdateAsync(
                u => u.Id == userId,
                u => new User
                {
                    MfaSecret        = encryptedSecret,
                    MfaRecoveryCodes = recoveryJson,
                    MfaEnabled       = false, // confirm step flips this true
                    MfaEnrolledAt    = null
                },
                token: ct));

        await auditLog.RecordMfaEnroll(user.Email, ct);

        return new MfaEnrollResponse
        {
            Secret        = secretBase32,
            OtpAuthUrl    = otpAuthUrl,
            QrPngBase64   = qrPng,
            RecoveryCodes = recoveryPlain
        };
    }

    public async Task Confirm(Guid userId, string code, CancellationToken ct = default)
    {
        var user = await userService.GetById(userId, ct)
            ?? throw new BusinessException(ExceptionEnum.NoEmailFound);

        if (user.MfaEnabled)
            throw new BusinessException(ExceptionEnum.MfaAlreadyEnabled);

        if (string.IsNullOrEmpty(user.MfaSecret))
            throw new BusinessException(ExceptionEnum.MfaNotEnrolling);

        var secret = _protector.Unprotect(user.MfaSecret);
        if (!VerifyTotpCode(secret, code, lastUsedWindow: null, out _))
            throw new BusinessException(ExceptionEnum.InvalidMfaCode);

        await dbFactory.RunAdmin(async db =>
            await db.Users.UpdateAsync(
                u => u.Id == userId,
                u => new User
                {
                    MfaEnabled    = true,
                    MfaEnrolledAt = DateTime.UtcNow
                },
                token: ct));

        await auditLog.RecordMfaConfirm(user.Email, ct);
    }

    public async Task Disable(Guid userId, string password, string code, CancellationToken ct = default)
    {
        var user = await userService.GetById(userId, ct)
            ?? throw new BusinessException(ExceptionEnum.NoEmailFound);

        if (!user.MfaEnabled)
            throw new BusinessException(ExceptionEnum.MfaNotEnabled);

        var verify = Security.VerifyPassword(password, user.PasswordHash);
        if (!verify.Valid)
            throw new BusinessException(ExceptionEnum.WrongPassword);

        var secret = _protector.Unprotect(user.MfaSecret!);
        if (!VerifyTotpCode(secret, code, lastUsedWindow: null, out _))
            throw new BusinessException(ExceptionEnum.InvalidMfaCode);

        // Raw SQL: setting mfa_recovery_codes (jsonb) to NULL via the LinqToDB UPDATE
        // expression sends an untyped NULL literal that Postgres parses as text and
        // rejects (42804). A small parameterized SQL avoids the type-inference dance.
        await dbFactory.RunAdmin(async db =>
            await db.ExecuteAsync(
                @"UPDATE public.users
                  SET mfa_enabled = false,
                      mfa_secret = NULL,
                      mfa_recovery_codes = NULL::jsonb,
                      mfa_enrolled_at = NULL,
                      mfa_last_used_window = NULL
                  WHERE id = @id",
                new DataParameter("id", userId, DataType.Guid)));

        await auditLog.RecordMfaDisable(user.Email, ct);
    }

    public string IssueMfaStepToken(Guid userId)
    {
        var active = signingKeys.Active;
        var creds  = new SigningCredentials(active.SecurityKey, SecurityAlgorithms.EcdsaSha256);
        var expires = DateTime.UtcNow.AddSeconds(MfaStepLifetimeSeconds);

        var descriptor = new SecurityTokenDescriptor
        {
            Subject = new ClaimsIdentity([
                new Claim(ClaimTypes.NameIdentifier, userId.ToString()),
                new Claim("token_use", "mfa_step")
            ]),
            Expires            = expires,
            Issuer             = _jwt.Issuer,
            // AZ-534 — narrow audience: this token is ONLY usable at /login/mfa.
            // The main JwtBearer middleware accepts _jwt.Audience and rejects this one.
            Audience           = MfaStepAudience,
            SigningCredentials = creds
        };

        var handler = new JwtSecurityTokenHandler();
        return handler.WriteToken(handler.CreateToken(descriptor));
    }

    public Guid ValidateMfaStepToken(string token)
    {
        try
        {
            var handler = new JwtSecurityTokenHandler();
            var principal = handler.ValidateToken(token, new TokenValidationParameters
            {
                ValidateIssuer           = true,
                ValidateAudience         = true,
                ValidateLifetime         = true,
                ValidateIssuerSigningKey = true,
                ValidIssuer              = _jwt.Issuer,
                ValidAudience            = MfaStepAudience,
                ValidAlgorithms          = [SecurityAlgorithms.EcdsaSha256],
                IssuerSigningKeyResolver = (_, _, _, _) =>
                    signingKeys.All.Select(k => (SecurityKey)k.SecurityKey)
            }, out _);

            var sub = principal.FindFirst(ClaimTypes.NameIdentifier)?.Value
                ?? throw new BusinessException(ExceptionEnum.InvalidMfaToken);
            return Guid.Parse(sub);
        }
        catch (BusinessException) { throw; }
        catch (Exception)
        {
            throw new BusinessException(ExceptionEnum.InvalidMfaToken);
        }
    }

    public async Task<string[]> VerifyForLogin(Guid userId, string code, CancellationToken ct = default)
    {
        var user = await userService.GetById(userId, ct)
            ?? throw new BusinessException(ExceptionEnum.InvalidCredentials);

        if (!user.MfaEnabled || string.IsNullOrEmpty(user.MfaSecret))
            throw new BusinessException(ExceptionEnum.MfaNotEnabled);

        // AZ-557 — active lockout from EITHER the password or the MFA side rejects
        // the request before the TOTP verify runs, with the same wire shape the
        // password path uses (`InvalidCredentials` + Retry-After).
        if (user.LockoutUntil is { } until && until > DateTime.UtcNow)
        {
            var remaining = (int)Math.Ceiling((until - DateTime.UtcNow).TotalSeconds);
            throw new BusinessException(ExceptionEnum.InvalidCredentials, Math.Max(remaining, 1));
        }

        // AZ-557 — per-account sliding-window rate limit applies to MFA failures too
        // (CountRecentFailedLogins counts login_failed + mfa_login_failed). Without
        // this an attacker with a leaked password could brute-force the 6-digit TOTP
        // from rotating IPs without ever tripping the per-account throttle.
        var recentFailures = await auditLog.CountRecentFailedLogins(
            user.Email, _auth.RateLimit.PerAccountWindowSeconds, ct);
        if (recentFailures >= _auth.RateLimit.PerAccountPermitLimit)
            throw new BusinessException(ExceptionEnum.InvalidCredentials, _auth.RateLimit.PerAccountWindowSeconds);

        var secret = _protector.Unprotect(user.MfaSecret);
        if (VerifyTotpCode(secret, code, user.MfaLastUsedWindow, out var window))
        {
            // Persist last-used window so a re-presented code in the same 30 s
            // step is rejected even if the attacker presents it before the next step.
            await dbFactory.RunAdmin(async db =>
                await db.Users.UpdateAsync(
                    u => u.Id == userId,
                    u => new User { MfaLastUsedWindow = window },
                    token: ct));
            // AZ-557 — TOTP success also resets the failure counter so a user who
            // fat-fingered a few codes before getting it right doesn't drift toward
            // lockout. Mirrors the password-side reset in RegisterSuccessfulLogin.
            await dbFactory.RunAdmin(async db =>
                await db.Users.UpdateAsync(
                    u => u.Id == userId,
                    u => new User { FailedLoginCount = 0, LockoutUntil = null },
                    token: ct));
            await auditLog.RecordMfaLoginSuccess(user.Email, ct);
            return ["pwd", "mfa"];
        }

        // TOTP failed — try recovery code (single-use). Recovery codes are
        // high-entropy and intentionally NOT counted by the lockout pipeline; a
        // locked-out user can still escape via a recovery code.
        if (await TryConsumeRecoveryCode(user, code, ct))
        {
            await dbFactory.RunAdmin(async db =>
                await db.Users.UpdateAsync(
                    u => u.Id == user.Id,
                    u => new User { FailedLoginCount = 0, LockoutUntil = null },
                    token: ct));
            await auditLog.RecordMfaRecoveryUsed(user.Email, ct);
            return ["pwd", "mfa", "recovery"];
        }

        // AZ-557 — feed the shared failure-accounting helper. It records the audit
        // row (mfa_login_failed), bumps failed_login_count, and on threshold-crossing
        // throws InvalidCredentials + Retry-After (which we let propagate). If it
        // does NOT throw, we fall through and throw the bare InvalidCredentials so
        // the wire response is uniform with the password path.
        await userService.RegisterMfaFailedLogin(user, ct);
        throw new BusinessException(ExceptionEnum.InvalidCredentials);
    }

    private static bool VerifyTotpCode(string secretBase32, string code, long? lastUsedWindow, out long matchedWindow)
    {
        matchedWindow = 0;
        var totp = new Totp(Base32Encoding.ToBytes(secretBase32));
        if (!totp.VerifyTotp(code, out matchedWindow, VerificationWindow.RfcSpecifiedNetworkDelay))
            return false;
        if (lastUsedWindow.HasValue && matchedWindow <= lastUsedWindow.Value)
            return false; // replay within or before the last accepted window
        return true;
    }

    private async Task<bool> TryConsumeRecoveryCode(User user, string code, CancellationToken ct)
    {
        if (string.IsNullOrEmpty(user.MfaRecoveryCodes)) return false;

        var codes = JsonSerializer.Deserialize<RecoveryCodeStore[]>(user.MfaRecoveryCodes)
            ?? Array.Empty<RecoveryCodeStore>();
        var candidateHash = HashRecoveryCode(code);

        var matchIdx = -1;
        for (var i = 0; i < codes.Length; i++)
        {
            if (codes[i].UsedAt != null) continue;
            if (CryptographicOperations.FixedTimeEquals(
                    System.Text.Encoding.ASCII.GetBytes(codes[i].Hash),
                    System.Text.Encoding.ASCII.GetBytes(candidateHash)))
            {
                matchIdx = i;
                break;
            }
        }
        if (matchIdx < 0) return false;

        codes[matchIdx] = codes[matchIdx] with { UsedAt = DateTime.UtcNow };
        var updated = JsonSerializer.Serialize(codes);

        await dbFactory.RunAdmin(async db =>
            await db.Users.UpdateAsync(
                // Conditional update on the prior JSON to avoid a race where two
                // concurrent /login/mfa calls both consume the same code.
                u => u.Id == user.Id && u.MfaRecoveryCodes == user.MfaRecoveryCodes,
                u => new User { MfaRecoveryCodes = updated },
                token: ct));
        return true;
    }

    private static string GenerateQrPng(string text)
    {
        using var generator = new QRCodeGenerator();
        using var data      = generator.CreateQrCode(text, QRCodeGenerator.ECCLevel.M);
        var pngBytes        = new PngByteQRCode(data).GetGraphic(pixelsPerModule: 6);
        return Convert.ToBase64String(pngBytes);
    }

    private static string HashRecoveryCode(string code)
    {
        var bytes  = System.Text.Encoding.UTF8.GetBytes(code);
        var digest = SHA256.HashData(bytes);
        return Convert.ToHexString(digest);
    }

    private sealed record RecoveryCodeStore
    {
        public string    Hash    { get; init; } = "";
        public DateTime? UsedAt  { get; init; }
    }
}