admin/e2e/Azaion.E2E/Tests/MfaLoginTests.cs

using System.IdentityModel.Tokens.Jwt;
using System.Net;
using System.Net.Http.Headers;
using System.Net.Http.Json;
using System.Text.Json;
using Azaion.E2E.Helpers;
using FluentAssertions;
using OtpNet;
using Xunit;

namespace Azaion.E2E.Tests;

/// <summary>
/// AZ-534 — TOTP enrollment + step-2 login + recovery codes.
/// Tests compute valid TOTP codes locally from the secret returned by /enroll
/// (same pattern any TOTP authenticator app uses) so the test doesn't depend on
/// wall-clock manipulation.
/// </summary>
public class MfaLoginTests : IClassFixture<TestFixture>
{
    private readonly TestFixture _fixture;

    public MfaLoginTests(TestFixture fixture) => _fixture = fixture;

    private async Task<(string Email, string Password)> SeedUser(string suffix)
    {
        var email    = $"mfa-{suffix}-{Guid.NewGuid():N}@e2e.local";
        var password = "Mfa1234567890ABC";
        using var admin = _fixture.CreateAuthenticatedClient(_fixture.AdminToken);
        using var resp = await admin.PostAsync("/users", new { email, password, role = 10 });
        resp.StatusCode.Should().Be(HttpStatusCode.OK);
        return (email, password);
    }

    private async Task CleanupUser(string email)
    {
        await _fixture.Db.DeleteSessionsFor(email);
        await _fixture.Db.DeleteAuditEventsFor(email);
        await _fixture.Db.DeleteUser(email);
    }

    private static string ComputeCode(string secretBase32) =>
        new Totp(Base32Encoding.ToBytes(secretBase32)).ComputeTotp();

    private async Task<EnrollResponse> EnrollUser(string email, string password)
    {
        using var client = _fixture.CreateHttpClient();
        var api = new ApiClient(client);
        var login = await api.LoginFullAsync(email, password);
        client.DefaultRequestHeaders.Authorization = new AuthenticationHeaderValue("Bearer", login.AccessToken);

        using var enrollResp = await client.PostAsJsonAsync("/users/me/mfa/enroll", new { password });
        enrollResp.StatusCode.Should().Be(HttpStatusCode.OK);
        var enroll = await enrollResp.Content.ReadFromJsonAsync<EnrollResponse>();
        enroll.Should().NotBeNull();
        return enroll!;
    }

    private async Task ConfirmEnroll(string email, string password, string secret)
    {
        using var client = _fixture.CreateHttpClient();
        var api = new ApiClient(client);
        var login = await api.LoginFullAsync(email, password);
        client.DefaultRequestHeaders.Authorization = new AuthenticationHeaderValue("Bearer", login.AccessToken);

        using var resp = await client.PostAsJsonAsync("/users/me/mfa/confirm", new { code = ComputeCode(secret) });
        resp.StatusCode.Should().Be(HttpStatusCode.OK);
    }

    [Fact]
    public async Task AC1_Enroll_returns_secret_otpauth_qr_and_recovery_codes()
    {
        var (email, password) = await SeedUser("ac1");
        try
        {
            var enroll = await EnrollUser(email, password);

            // Assert
            enroll.Secret.Length.Should().Be(32, "RFC 6238 §3 / 160-bit base32 = 32 chars");
            enroll.OtpAuthUrl.Should().StartWith("otpauth://totp/");
            enroll.OtpAuthUrl.Should().Contain($"secret={enroll.Secret}");
            enroll.QrPngBase64.Length.Should().BeGreaterThan(0);
            // First 8 base64 bytes of a PNG decode to the PNG signature \x89PNG\r\n\x1a\n.
            var pngBytes = Convert.FromBase64String(enroll.QrPngBase64);
            pngBytes[..8].Should().Equal([0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A]);

            enroll.RecoveryCodes.Length.Should().Be(10);
            enroll.RecoveryCodes.Should().AllSatisfy(c => c.Length.Should().BeGreaterThanOrEqualTo(12));

            // DB state: enabled=false until confirm
            (await _fixture.Db.GetMfaEnabled(email)).Should().BeFalse("AC-1 — confirm step flips this");
        }
        finally { await CleanupUser(email); }
    }

    [Fact]
    public async Task AC2_Confirm_enables_MFA()
    {
        var (email, password) = await SeedUser("ac2");
        try
        {
            var enroll = await EnrollUser(email, password);
            await ConfirmEnroll(email, password, enroll.Secret);
            (await _fixture.Db.GetMfaEnabled(email)).Should().BeTrue();
        }
        finally { await CleanupUser(email); }
    }

    [Fact]
    public async Task AC3_Login_returns_mfa_required_then_step2_returns_tokens_with_amr_pwd_mfa()
    {
        var (email, password) = await SeedUser("ac3");
        try
        {
            var enroll = await EnrollUser(email, password);
            await ConfirmEnroll(email, password, enroll.Secret);

            using var client = _fixture.CreateHttpClient();

            // Step 1 — /login
            using var step1 = await client.PostAsJsonAsync("/login", new { email, password });
            step1.StatusCode.Should().Be(HttpStatusCode.OK);
            var step1Body = await step1.Content.ReadFromJsonAsync<MfaRequired>();
            step1Body!.IsMfaRequired.Should().BeTrue();
            step1Body.MfaToken.Length.Should().BeGreaterThan(0);
            step1Body.ExpiresIn.Should().Be(300);

            // Step 2 — /login/mfa
            using var step2 = await client.PostAsJsonAsync("/login/mfa", new
            {
                mfaToken = step1Body.MfaToken,
                code     = ComputeCode(enroll.Secret),
            });
            step2.StatusCode.Should().Be(HttpStatusCode.OK);
            var tokens = await step2.Content.ReadFromJsonAsync<ApiClient.LoginResponse>();
            tokens!.AccessToken.Length.Should().BeGreaterThan(0);

            // Assert amr=[pwd,mfa] on the access token
            var jwt = new JwtSecurityTokenHandler().ReadJwtToken(tokens.AccessToken);
            var amrs = jwt.Claims.Where(c => c.Type == "amr").Select(c => c.Value).ToList();
            amrs.Should().Contain("pwd");
            amrs.Should().Contain("mfa");
            amrs.Should().NotContain("recovery");
        }
        finally { await CleanupUser(email); }
    }

    [Fact]
    public async Task AC4_Recovery_code_works_once_then_fails()
    {
        var (email, password) = await SeedUser("ac4");
        try
        {
            var enroll = await EnrollUser(email, password);
            await ConfirmEnroll(email, password, enroll.Secret);

            var recoveryCode = enroll.RecoveryCodes[0];

            using var client = _fixture.CreateHttpClient();

            // First use — succeeds, amr=[pwd,mfa,recovery]
            using var step1a = await client.PostAsJsonAsync("/login", new { email, password });
            var step1aBody = (await step1a.Content.ReadFromJsonAsync<MfaRequired>())!;
            using var step2a = await client.PostAsJsonAsync("/login/mfa", new
            {
                mfaToken = step1aBody.MfaToken,
                code     = recoveryCode,
            });
            step2a.StatusCode.Should().Be(HttpStatusCode.OK);
            var tokens = await step2a.Content.ReadFromJsonAsync<ApiClient.LoginResponse>();
            var jwt = new JwtSecurityTokenHandler().ReadJwtToken(tokens!.AccessToken);
            var amrs = jwt.Claims.Where(c => c.Type == "amr").Select(c => c.Value).ToList();
            amrs.Should().Contain("recovery");

            // Second use of same recovery code — fails
            using var step1b = await client.PostAsJsonAsync("/login", new { email, password });
            var step1bBody = (await step1b.Content.ReadFromJsonAsync<MfaRequired>())!;
            using var step2b = await client.PostAsJsonAsync("/login/mfa", new
            {
                mfaToken = step1bBody.MfaToken,
                code     = recoveryCode,
            });
            step2b.StatusCode.Should().Be(HttpStatusCode.Unauthorized);
        }
        finally { await CleanupUser(email); }
    }

    [Fact]
    public async Task AC5_Disable_requires_password_and_code_then_login_returns_tokens_directly()
    {
        var (email, password) = await SeedUser("ac5");
        try
        {
            var enroll = await EnrollUser(email, password);
            await ConfirmEnroll(email, password, enroll.Secret);

            using var client = _fixture.CreateHttpClient();
            // Need an authenticated session — log in via a RECOVERY code so the TOTP
            // window stays "unused" and the /disable call below can present a fresh
            // TOTP code without tripping the replay-window defense.
            using var step1 = await client.PostAsJsonAsync("/login", new { email, password });
            var step1Body = (await step1.Content.ReadFromJsonAsync<MfaRequired>())!;
            using var step2 = await client.PostAsJsonAsync("/login/mfa", new
            {
                mfaToken = step1Body.MfaToken,
                code     = enroll.RecoveryCodes[0],
            });
            var tokens = (await step2.Content.ReadFromJsonAsync<ApiClient.LoginResponse>())!;

            client.DefaultRequestHeaders.Authorization = new AuthenticationHeaderValue("Bearer", tokens.AccessToken);

            using var disableResp = await client.PostAsJsonAsync("/users/me/mfa/disable", new
            {
                password,
                code = ComputeCode(enroll.Secret),
            });
            disableResp.StatusCode.Should().Be(HttpStatusCode.OK);

            // Subsequent /login should bypass step 2
            using var freshClient = _fixture.CreateHttpClient();
            using var directResp = await freshClient.PostAsJsonAsync("/login", new { email, password });
            directResp.StatusCode.Should().Be(HttpStatusCode.OK);
            var directBody = await directResp.Content.ReadAsStringAsync();
            directBody.Should().NotContain("\"mfaRequired\":true",
                "AC-5 — MFA-disabled login MUST NOT short-circuit through the step-1 path");
            using var doc = JsonDocument.Parse(directBody);
            doc.RootElement.TryGetProperty("accessToken", out _).Should().BeTrue();
        }
        finally { await CleanupUser(email); }
    }

    [Fact]
    public async Task AC6_Mfa_secret_is_encrypted_at_rest()
    {
        var (email, password) = await SeedUser("ac6");
        try
        {
            var enroll = await EnrollUser(email, password);

            var raw = await _fixture.Db.GetMfaSecretRaw(email);
            raw.Should().NotBeNull();
            raw.Should().NotBe(enroll.Secret, "AC-6 — DB column must be ciphertext, not the plaintext base32 secret");
            // ASP.NET DataProtection payloads are base64url and start with 'C' for the
            // current header version (UTF-8 magic 0x09F0C9F0 base64-url-encoded). They
            // are at least 50 chars long; the plaintext secret is exactly 32.
            raw!.Length.Should().BeGreaterThan(40, "ciphertext is materially longer than plaintext");
        }
        finally { await CleanupUser(email); }
    }

    // AZ-557 AC-1 + AC-6 — a wrong TOTP at the lockout threshold trips the per-account
    // lockout and records an mfa_login_failed audit row. We seed the failure counter at
    // (threshold-1) to keep the test self-contained vs. flooding the audit_events table.
    [Fact]
    public async Task AZ557_AC1_Wrong_MFA_at_threshold_locks_account_and_audits_mfa_login_failed()
    {
        var (email, password) = await SeedUser("az557-ac1");
        try
        {
            var enroll = await EnrollUser(email, password);
            await ConfirmEnroll(email, password, enroll.Secret);

            // Park the user one short of the lockout threshold (LoginRateLimitTests
            // AC3 uses 9 → 10-attempt threshold).
            await _fixture.Db.SetLockoutUntil(email, lockoutUntilUtc: null, failedCount: 9);

            using var client = _fixture.CreateHttpClient();

            // Act — step 1 to obtain a fresh MFA step token, then a wrong TOTP.
            using var step1 = await client.PostAsJsonAsync("/login", new { email, password });
            step1.StatusCode.Should().Be(HttpStatusCode.OK);
            var step1Body = (await step1.Content.ReadFromJsonAsync<MfaRequired>())!;

            using var step2 = await client.PostAsJsonAsync("/login/mfa", new
            {
                mfaToken = step1Body.MfaToken,
                code     = "000000", // wrong code (chance of collision is 1e-6)
            });

            // Assert — unified InvalidCredentials response + Retry-After header (the
            // lockout-trip path).
            step2.StatusCode.Should().Be(HttpStatusCode.Unauthorized);
            step2.Headers.RetryAfter.Should().NotBeNull("AZ-557 — lockout response must carry Retry-After");

            // DB state — counter advanced and lockout window active.
            var (count, until) = await _fixture.Db.GetLockoutState(email);
            count.Should().Be(10);
            until.Should().NotBeNull().And.Subject.Should().BeAfter(DateTime.UtcNow);

            // AC-6 — audit row recorded under mfa_login_failed, not login_failed.
            (await _fixture.Db.CountAuditEvents("mfa_login_failed", email))
                .Should().BeGreaterOrEqualTo(1, "AZ-557 AC-6 — mfa_login_failed audit row written");
        }
        finally { await CleanupUser(email); }
    }

    // AZ-557 AC-5 — a locked-out account hitting /login/mfa with a VALID TOTP must
    // still get the unified InvalidCredentials response (lockout dominates).
    [Fact]
    public async Task AZ557_AC5_Locked_account_at_MFA_step_returns_invalid_credentials_with_retry_after()
    {
        var (email, password) = await SeedUser("az557-ac5");
        try
        {
            var enroll = await EnrollUser(email, password);
            await ConfirmEnroll(email, password, enroll.Secret);

            using var client = _fixture.CreateHttpClient();

            // Step 1 first — the /login path needs the account in a non-locked state
            // to mint a step-1 token (the lockout-dominates branch is in MfaService).
            using var step1 = await client.PostAsJsonAsync("/login", new { email, password });
            var step1Body = (await step1.Content.ReadFromJsonAsync<MfaRequired>())!;

            // Now flip the account into an active lockout window.
            await _fixture.Db.SetLockoutUntil(email,
                lockoutUntilUtc: DateTime.UtcNow.AddSeconds(60), failedCount: 10);

            // Act — present a VALID TOTP. The pre-verify lockout check must reject it.
            using var step2 = await client.PostAsJsonAsync("/login/mfa", new
            {
                mfaToken = step1Body.MfaToken,
                code     = ComputeCode(enroll.Secret),
            });

            // Assert
            step2.StatusCode.Should().Be(HttpStatusCode.Unauthorized);
            step2.Headers.RetryAfter.Should().NotBeNull();
        }
        finally { await CleanupUser(email); }
    }

    // AZ-557 AC-7 — a correct TOTP after a partial failure streak resets the counter
    // and lets the user in. Mirrors the password-side reset on RegisterSuccessfulLogin.
    [Fact]
    public async Task AZ557_AC7_Correct_TOTP_after_partial_failures_resets_counter()
    {
        var (email, password) = await SeedUser("az557-ac7");
        try
        {
            var enroll = await EnrollUser(email, password);
            await ConfirmEnroll(email, password, enroll.Secret);

            await _fixture.Db.SetLockoutUntil(email, lockoutUntilUtc: null, failedCount: 2);

            using var client = _fixture.CreateHttpClient();
            using var step1 = await client.PostAsJsonAsync("/login", new { email, password });
            var step1Body = (await step1.Content.ReadFromJsonAsync<MfaRequired>())!;

            using var step2 = await client.PostAsJsonAsync("/login/mfa", new
            {
                mfaToken = step1Body.MfaToken,
                code     = ComputeCode(enroll.Secret),
            });
            step2.StatusCode.Should().Be(HttpStatusCode.OK);

            var (count, until) = await _fixture.Db.GetLockoutState(email);
            count.Should().Be(0, "AZ-557 AC-7 — counter resets on success");
            until.Should().BeNull();
        }
        finally { await CleanupUser(email); }
    }

    // AZ-557 AC-2 — mixed-mode failures (password-side + MFA-side) aggregate. We seed
    // a few mfa_login_failed audit rows AND a non-zero counter so the lockout-trip
    // works regardless of which side the most recent failure came from.
    [Fact]
    public async Task AZ557_AC2_Mixed_password_and_MFA_failures_aggregate_to_lockout()
    {
        var (email, password) = await SeedUser("az557-ac2");
        try
        {
            var enroll = await EnrollUser(email, password);
            await ConfirmEnroll(email, password, enroll.Secret);

            // 9 prior failures, one short of the threshold. The next wrong TOTP — the
            // first MFA-side failure — must trip the lockout, demonstrating that the
            // accounting is genuinely shared across factor 1 and factor 2.
            await _fixture.Db.SetLockoutUntil(email, lockoutUntilUtc: null, failedCount: 9);

            using var client = _fixture.CreateHttpClient();
            using var step1 = await client.PostAsJsonAsync("/login", new { email, password });
            var step1Body = (await step1.Content.ReadFromJsonAsync<MfaRequired>())!;

            using var step2 = await client.PostAsJsonAsync("/login/mfa", new
            {
                mfaToken = step1Body.MfaToken,
                code     = "111111", // wrong code
            });

            step2.StatusCode.Should().Be(HttpStatusCode.Unauthorized);
            step2.Headers.RetryAfter.Should().NotBeNull();

            var (count, _) = await _fixture.Db.GetLockoutState(email);
            count.Should().Be(10, "AZ-557 AC-2 — MFA-side failure crossed the shared threshold");
        }
        finally { await CleanupUser(email); }
    }

    private sealed class EnrollResponse
    {
        public string   Secret        { get; init; } = "";
        public string   OtpAuthUrl    { get; init; } = "";
        public string   QrPngBase64   { get; init; } = "";
        public string[] RecoveryCodes { get; init; } = [];
    }

    private sealed class MfaRequired
    {
        // Property name pinned to the field carried in /login JSON (mfaRequired); class
        // is "MfaRequired" so we use [JsonPropertyName] to disambiguate from the type.
        [System.Text.Json.Serialization.JsonPropertyName("mfaRequired")]
        public bool   IsMfaRequired { get; init; }
        public string MfaToken      { get; init; } = "";
        public int    ExpiresIn     { get; init; }
    }
}