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); }
    }

    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; }
    }
}