[PM-6664] base request validator - Two Factor flows integration tests (#4643)

* initial commit added two factor tests

* initial commit

* updated two factor tests

* fixed formatting

src/Identity/IdentityServer/BaseRequestValidator.cs

@@ -350,9 +350,8 @@ public abstract class BaseRequestValidator<T> where T : class
                                  (await _userManager.GetValidTwoFactorProvidersAsync(user)).Count > 0;
 
         Organization firstEnabledOrg = null;
-        var orgs = (await CurrentContext.OrganizationMembershipAsync(_organizationUserRepository, user.Id))
-            .ToList();
-        if (orgs.Any())
+        var orgs = (await CurrentContext.OrganizationMembershipAsync(_organizationUserRepository, user.Id)).ToList();
+        if (orgs.Count > 0)
         {
             var orgAbilities = await _applicationCacheService.GetOrganizationAbilitiesAsync();
             var twoFactorOrgs = orgs.Where(o => OrgUsing2fa(orgAbilities, o.Id));

test/Identity.IntegrationTest/Endpoints/IdentityServerTwoFactorTests.cs

@@ -1,48 +1,53 @@
-using System.Text.Json;
+using System.Security.Claims;
+using System.Text.Json;
 using Bit.Core.AdminConsole.Entities;
+using Bit.Core.Auth.Entities;
 using Bit.Core.Auth.Enums;
+using Bit.Core.Auth.Models.Data;
+using Bit.Core.Auth.Repositories;
 using Bit.Core.Entities;
 using Bit.Core.Enums;
+using Bit.Core.Models.Data;
 using Bit.Core.Repositories;
 using Bit.Core.Services;
+using Bit.Core.Utilities;
 using Bit.Identity.Models.Request.Accounts;
 using Bit.IntegrationTestCommon.Factories;
 using Bit.Test.Common.AutoFixture.Attributes;
 using Bit.Test.Common.Helpers;
-using Microsoft.AspNetCore.TestHost;
+using Duende.IdentityServer.Models;
+using Duende.IdentityServer.Stores;
+using IdentityModel;
+using LinqToDB;
+using NSubstitute;
 using Xunit;
 
+// #nullable enable
+
 namespace Bit.Identity.IntegrationTest.Endpoints;
 
 public class IdentityServerTwoFactorTests : IClassFixture<IdentityApplicationFactory>
 {
+    const string _organizationTwoFactor = """{"6":{"Enabled":true,"MetaData":{"IKey":"DIEFB13LB49IEB3459N2","SKey":"0ZnsZHav0KcNPBZTS6EOUwqLPoB0sfMd5aJeWExQ","Host":"api-example.duosecurity.com"}}}""";
+    const string _testEmail = "test+2farequired@email.com";
+    const string _testPassword = "master_password_hash";
+    const string _userEmailTwoFactor = """{"1": { "Enabled": true, "MetaData": { "Email": "test+2farequired@email.com"}}}""";
+
     private readonly IdentityApplicationFactory _factory;
-    private readonly IUserRepository _userRepository;
-    private readonly IUserService _userService;
 
     public IdentityServerTwoFactorTests(IdentityApplicationFactory factory)
     {
         _factory = factory;
-        _userRepository = _factory.GetService<IUserRepository>();
-        _userService = _factory.GetService<IUserService>();
     }
 
-    [Theory, BitAutoData]
-    public async Task TokenEndpoint_UserTwoFactorRequired_NoTwoFactorProvided_Fails(string deviceId)
+    [Fact]
+    public async Task TokenEndpoint_GrantTypePassword_UserTwoFactorRequired_NoTwoFactorProvided_Fails()
     {
         // Arrange
-        var username = "test+2farequired@email.com";
-        var twoFactor = """{"1": { "Enabled": true, "MetaData": { "Email": "test+2farequired@email.com"}}}""";
-
-        await CreateUserAsync(_factory.Server, username, deviceId, async () =>
-        {
-            var user = await _userRepository.GetByEmailAsync(username);
-            user.TwoFactorProviders = twoFactor;
-            await _userService.UpdateTwoFactorProviderAsync(user, TwoFactorProviderType.Email);
-        });
+        await CreateUserAsync(_factory, _testEmail, _userEmailTwoFactor);
 
         // Act
-        var context = await PostLoginAsync(_factory.Server, username, deviceId);
+        var context = await _factory.ContextFromPasswordAsync(_testEmail, _testPassword);
 
         // Assert
         var body = await AssertHelper.AssertResponseTypeIs<JsonDocument>(context);
@@ -52,92 +57,437 @@ public class IdentityServerTwoFactorTests : IClassFixture<IdentityApplicationFac
         Assert.Equal("Two factor required.", error);
     }
 
-    [Theory, BitAutoData]
-    public async Task TokenEndpoint_OrgTwoFactorRequired_NoTwoFactorProvided_Fails(string deviceId)
+    [Fact]
+    public async Task TokenEndpoint_GrantTypePassword_UserTwoFactorRequired_TwoFactorProvided_Success()
     {
         // Arrange
-        var username = "test+org2farequired@email.com";
-        // use valid length keys so DuoWeb.SignRequest doesn't throw
-        // ikey: 20, skey: 40, akey: 40
-        var orgTwoFactor =
-            """{"6":{"Enabled":true,"MetaData":{"IKey":"DIEFB13LB49IEB3459N2","SKey":"0ZnsZHav0KcNPBZTS6EOUwqLPoB0sfMd5aJeWExQ","Host":"api-example.duosecurity.com"}}}""";
+        // we can't use the class factory here.
+        var factory = new IdentityApplicationFactory();
 
-        var server = _factory.WithWebHostBuilder(builder =>
+        string emailToken = null;
+        factory.SubstituteService<IMailService>(mailService =>
         {
-            builder.UseSetting("globalSettings:Duo:AKey", "WJHB374KM3N5hglO9hniwbkibg$789EfbhNyLpNq1");
-        }).Server;
-
-
-        await CreateUserAsync(server, username, deviceId, async () =>
-        {
-            var user = await _userRepository.GetByEmailAsync(username);
-
-            var organizationRepository = _factory.Services.GetService<IOrganizationRepository>();
-            var organization = await organizationRepository.CreateAsync(new Organization
-            {
-                Name = "Test Org",
-                Use2fa = true,
-                TwoFactorProviders = orgTwoFactor,
-                BillingEmail = "billing-email@example.com",
-                Plan = "Enterprise",
-            });
-
-            await _factory.Services.GetService<IOrganizationUserRepository>()
-                .CreateAsync(new OrganizationUser
-                {
-                    UserId = user.Id,
-                    OrganizationId = organization.Id,
-                    Status = OrganizationUserStatusType.Confirmed,
-                    Type = OrganizationUserType.User,
-                });
+            mailService.SendTwoFactorEmailAsync(Arg.Any<string>(), Arg.Do<string>(t => emailToken = t))
+                .Returns(Task.CompletedTask);
         });
 
+        // Create Test User
+        await CreateUserAsync(factory, _testEmail, _userEmailTwoFactor);
+
         // Act
-        var context = await PostLoginAsync(server, username, deviceId);
+        var failedTokenContext = await factory.ContextFromPasswordAsync(_testEmail, _testPassword);
+
+        Assert.Equal(StatusCodes.Status400BadRequest, failedTokenContext.Response.StatusCode);
+        Assert.NotNull(emailToken);
+
+        var twoFactorProvidedContext = await factory.ContextFromPasswordWithTwoFactorAsync(
+            _testEmail,
+            _testPassword,
+            twoFactorToken: emailToken);
+
+        // Assert
+        var body = await AssertHelper.AssertResponseTypeIs<JsonDocument>(twoFactorProvidedContext);
+        var root = body.RootElement;
+
+        var result = AssertHelper.AssertJsonProperty(root, "access_token", JsonValueKind.String).GetString();
+        Assert.NotNull(result);
+    }
+
+    [Fact]
+    public async Task TokenEndpoint_GrantTypePassword_InvalidTwoFactorToken_Fails()
+    {
+        // Arrange
+        await CreateUserAsync(_factory, _testEmail, _userEmailTwoFactor);
+
+        // Act
+        var context = await _factory.ContextFromPasswordWithTwoFactorAsync(
+                                _testEmail, _testPassword, twoFactorProviderType: "Email");
 
         // Assert
         var body = await AssertHelper.AssertResponseTypeIs<JsonDocument>(context);
         var root = body.RootElement;
 
+        var errorModel = AssertHelper.AssertJsonProperty(root, "ErrorModel", JsonValueKind.Object);
+        var errorMessage = AssertHelper.AssertJsonProperty(errorModel, "Message", JsonValueKind.String).GetString();
+        Assert.Equal("Two-step token is invalid. Try again.", errorMessage);
+
         var error = AssertHelper.AssertJsonProperty(root, "error_description", JsonValueKind.String).GetString();
-        Assert.Equal("Two factor required.", error);
+        Assert.Equal("invalid_username_or_password", error);
     }
 
-    private async Task CreateUserAsync(TestServer server, string username, string deviceId,
-        Func<Task> twoFactorSetup)
+    [Theory, BitAutoData]
+    public async Task TokenEndpoint_GrantTypePassword_OrgDuoTwoFactorRequired_NoTwoFactorProvided_Fails(string deviceId)
     {
-        // Register user
-        await _factory.RegisterAsync(new RegisterRequestModel
+        // Arrange
+        var challenge = new string('c', 50);
+        var ssoConfigData = new SsoConfigurationData
         {
-            Email = username,
-            MasterPasswordHash = "master_password_hash"
-        });
+            MemberDecryptionType = MemberDecryptionType.MasterPassword,
+        };
+        await CreateSsoOrganizationAndUserAsync(
+            _factory, ssoConfigData, challenge, _testEmail, orgTwoFactor: _organizationTwoFactor);
 
-        // Add two factor
-        if (twoFactorSetup != null)
-        {
-            await twoFactorSetup();
-        }
-    }
-
-    private async Task<HttpContext> PostLoginAsync(TestServer server, string username, string deviceId,
-        Action<HttpContext> extraConfiguration = null)
-    {
-        return await server.PostAsync("/connect/token", new FormUrlEncodedContent(new Dictionary<string, string>
+        // Act
+        var context = await _factory.Server.PostAsync("/connect/token", new FormUrlEncodedContent(new Dictionary<string, string>
         {
             { "scope", "api offline_access" },
             { "client_id", "web" },
-            { "deviceType", DeviceTypeAsString(DeviceType.FirefoxBrowser) },
+            { "deviceType", "12" },
             { "deviceIdentifier", deviceId },
-            { "deviceName", "firefox" },
+            { "deviceName", "edge" },
             { "grant_type", "password" },
-            { "username", username },
-            { "password", "master_password_hash" },
-        }), context => context.SetAuthEmail(username));
+            { "username", _testEmail },
+            { "password", _testPassword },
+        }), context => context.Request.Headers.Append("Auth-Email", CoreHelpers.Base64UrlEncodeString(_testEmail)));
+
+        // Assert 
+        using var responseBody = await AssertHelper.AssertResponseTypeIs<JsonDocument>(context);
+        var root = responseBody.RootElement;
+        var error = AssertHelper.AssertJsonProperty(root, "error_description", JsonValueKind.String).GetString();
+        Assert.Equal("Two factor required.", error);
     }
 
-    private static string DeviceTypeAsString(DeviceType deviceType)
+    [Fact]
+    public async Task TokenEndpoint_GrantTypePassword_RememberTwoFactorType_InvalidTwoFactorToken_Fails()
     {
-        return ((int)deviceType).ToString();
+        // Arrange
+        await CreateUserAsync(_factory, _testEmail, _userEmailTwoFactor);
+
+        // Act
+        var context = await _factory.ContextFromPasswordWithTwoFactorAsync(
+                                _testEmail, _testPassword, twoFactorProviderType: "Remember");
+
+        // Assert
+        var body = await AssertHelper.AssertResponseTypeIs<JsonDocument>(context);
+        var root = body.RootElement;
+
+        var error = AssertHelper.AssertJsonProperty(root, "error_description", JsonValueKind.String).GetString();
+        Assert.Equal("Two factor required.", error);
+    }
+
+    [Theory, BitAutoData]
+    public async Task TokenEndpoint_GrantTypeClientCredential_OrgTwoFactorRequired_Success(Organization organization, OrganizationApiKey organizationApiKey)
+    {
+        // Arrange        
+        organization.Enabled = true;
+        organization.UseApi = true;
+        organization.Use2fa = true;
+        organization.TwoFactorProviders = _organizationTwoFactor;
+
+        var orgRepo = _factory.Services.GetRequiredService<IOrganizationRepository>();
+        organization = await orgRepo.CreateAsync(organization);
+
+        organizationApiKey.OrganizationId = organization.Id;
+        organizationApiKey.Type = OrganizationApiKeyType.Default;
+
+        var orgApiKeyRepo = _factory.Services.GetRequiredService<IOrganizationApiKeyRepository>();
+        await orgApiKeyRepo.CreateAsync(organizationApiKey);
+
+        // Act
+        var context = await _factory.Server.PostAsync("/connect/token", new FormUrlEncodedContent(new Dictionary<string, string>
+        {
+            { "grant_type", "client_credentials" },
+            { "client_id", $"organization.{organization.Id}" },
+            { "client_secret", organizationApiKey.ApiKey },
+            { "scope", "api.organization" },
+        }));
+
+        // Assert
+        Assert.Equal(StatusCodes.Status200OK, context.Response.StatusCode);
+
+        var body = await AssertHelper.AssertResponseTypeIs<JsonDocument>(context);
+        var root = body.RootElement;
+        var token = AssertHelper.AssertJsonProperty(root, "access_token", JsonValueKind.String).GetString();
+        Assert.NotNull(token);
+    }
+
+    [Theory, BitAutoData]
+    public async Task TokenEndpoint_GrantTypeClientCredential_IndvTwoFactorRequired_Success(string deviceId)
+    {
+        // Arrange
+        await CreateUserAsync(_factory, _testEmail, _userEmailTwoFactor);
+
+        var database = _factory.GetDatabaseContext();
+        var user = await database.Users.FirstAsync(u => u.Email == _testEmail);
+
+        // Act
+        var context = await _factory.Server.PostAsync("/connect/token", new FormUrlEncodedContent(new Dictionary<string, string>
+        {
+            { "grant_type", "client_credentials" },
+            { "client_id", $"user.{user.Id}" },
+            { "client_secret", user.ApiKey },
+            { "scope", "api" },
+            { "DeviceIdentifier", deviceId },
+            { "DeviceType",  ((int)DeviceType.FirefoxBrowser).ToString() },
+            { "DeviceName", "firefox" },
+        }));
+
+        // Assert
+        Assert.Equal(StatusCodes.Status200OK, context.Response.StatusCode);
+
+        var body = await AssertHelper.AssertResponseTypeIs<JsonDocument>(context);
+        var root = body.RootElement;
+        var token = AssertHelper.AssertJsonProperty(root, "access_token", JsonValueKind.String).GetString();
+        Assert.NotNull(token);
+    }
+
+    [Theory, BitAutoData]
+    public async Task TokenEndpoint_GrantTypeAuthCode_OrgTwoFactorRequired_IndvTwoFactor_NoTwoFactorProvided_Fails(string deviceId)
+    {
+        // Arrange
+        var localFactory = new IdentityApplicationFactory();
+        var challenge = new string('c', 50);
+        var ssoConfigData = new SsoConfigurationData
+        {
+            MemberDecryptionType = MemberDecryptionType.MasterPassword,
+        };
+        await CreateSsoOrganizationAndUserAsync(
+            localFactory, ssoConfigData, challenge, _testEmail, userTwoFactor: _userEmailTwoFactor);
+
+        // Act
+        var context = await localFactory.Server.PostAsync("/connect/token", new FormUrlEncodedContent(new Dictionary<string, string>
+        {
+            { "scope", "api offline_access" },
+            { "client_id", "web" },
+            { "deviceType", "12" },
+            { "deviceIdentifier", deviceId },
+            { "deviceName", "edge" },
+            { "grant_type", "authorization_code" },
+            { "code", "test_code" },
+            { "code_verifier", challenge },
+            { "redirect_uri", "https://localhost:8080/sso-connector.html" }
+        }), context => context.Request.Headers.Append("Auth-Email", CoreHelpers.Base64UrlEncodeString(_testEmail)));
+
+        // Assert 
+        using var responseBody = await AssertHelper.AssertResponseTypeIs<JsonDocument>(context);
+        var root = responseBody.RootElement;
+        var error = AssertHelper.AssertJsonProperty(root, "error_description", JsonValueKind.String).GetString();
+        Assert.Equal("Two factor required.", error);
+    }
+
+    [Theory, BitAutoData]
+    public async Task TokenEndpoint_GrantTypeAuthCode_OrgTwoFactorRequired_IndvTwoFactor_TwoFactorProvided_Success(string deviceId)
+    {
+        // Arrange
+        var localFactory = new IdentityApplicationFactory();
+        string emailToken = null;
+        localFactory.SubstituteService<IMailService>(mailService =>
+        {
+            mailService.SendTwoFactorEmailAsync(Arg.Any<string>(), Arg.Do<string>(t => emailToken = t))
+                .Returns(Task.CompletedTask);
+        });
+
+        // Create Test User
+        var challenge = new string('c', 50);
+        var ssoConfigData = new SsoConfigurationData
+        {
+            MemberDecryptionType = MemberDecryptionType.MasterPassword,
+        };
+        await CreateSsoOrganizationAndUserAsync(
+            localFactory, ssoConfigData, challenge, _testEmail, userTwoFactor: _userEmailTwoFactor);
+
+        // Act
+        var failedTokenContext = await localFactory.Server.PostAsync("/connect/token", new FormUrlEncodedContent(new Dictionary<string, string>
+        {
+            { "scope", "api offline_access" },
+            { "client_id", "web" },
+            { "deviceType", "12" },
+            { "deviceIdentifier", deviceId },
+            { "deviceName", "edge" },
+            { "grant_type", "authorization_code" },
+            { "code", "test_code" },
+            { "code_verifier", challenge },
+            { "redirect_uri", "https://localhost:8080/sso-connector.html" }
+        }), context => context.Request.Headers.Append("Auth-Email", CoreHelpers.Base64UrlEncodeString(_testEmail)));
+
+        Assert.Equal(StatusCodes.Status400BadRequest, failedTokenContext.Response.StatusCode);
+        Assert.NotNull(emailToken);
+
+        var twoFactorProvidedContext = await localFactory.Server.PostAsync("/connect/token", new FormUrlEncodedContent(new Dictionary<string, string>
+        {
+            { "scope", "api offline_access" },
+            { "client_id", "web" },
+            { "deviceType", "12" },
+            { "deviceIdentifier", deviceId },
+            { "deviceName", "edge" },
+            { "twoFactorToken", emailToken},
+            { "twoFactorProvider", "1" },
+            { "twoFactorRemember", "0" },
+            { "grant_type", "authorization_code" },
+            { "code", "test_code" },
+            { "code_verifier", challenge },
+            { "redirect_uri", "https://localhost:8080/sso-connector.html" }
+        }), context => context.Request.Headers.Append("Auth-Email", CoreHelpers.Base64UrlEncodeString(_testEmail)));
+
+
+        // Assert 
+        var body = await AssertHelper.AssertResponseTypeIs<JsonDocument>(twoFactorProvidedContext);
+        var root = body.RootElement;
+
+        var result = AssertHelper.AssertJsonProperty(root, "access_token", JsonValueKind.String).GetString();
+        Assert.NotNull(result);
+    }
+
+    [Theory, BitAutoData]
+    public async Task TokenEndpoint_GrantTypeAuthCode_OrgTwoFactorRequired_OrgDuoTwoFactor_NoTwoFactorProvided_Fails(string deviceId)
+    {
+        // Arrange
+        var localFactory = new IdentityApplicationFactory();
+        var challenge = new string('c', 50);
+        var ssoConfigData = new SsoConfigurationData
+        {
+            MemberDecryptionType = MemberDecryptionType.MasterPassword,
+        };
+        await CreateSsoOrganizationAndUserAsync(
+            localFactory, ssoConfigData, challenge, _testEmail, orgTwoFactor: _organizationTwoFactor);
+
+        // Act
+        var context = await localFactory.Server.PostAsync("/connect/token", new FormUrlEncodedContent(new Dictionary<string, string>
+        {
+            { "scope", "api offline_access" },
+            { "client_id", "web" },
+            { "deviceType", "12" },
+            { "deviceIdentifier", deviceId },
+            { "deviceName", "edge" },
+            { "grant_type", "authorization_code" },
+            { "code", "test_code" },
+            { "code_verifier", challenge },
+            { "redirect_uri", "https://localhost:8080/sso-connector.html" }
+        }), context => context.Request.Headers.Append("Auth-Email", CoreHelpers.Base64UrlEncodeString(_testEmail)));
+
+        // Assert 
+        using var responseBody = await AssertHelper.AssertResponseTypeIs<JsonDocument>(context);
+        var root = responseBody.RootElement;
+        var error = AssertHelper.AssertJsonProperty(root, "error_description", JsonValueKind.String).GetString();
+        Assert.Equal("Two factor required.", error);
+    }
+
+    private async Task CreateUserAsync(
+        IdentityApplicationFactory factory,
+        string testEmail,
+        string userTwoFactor = null)
+    {
+        // Create Test User
+        await factory.RegisterAsync(new RegisterRequestModel
+        {
+            Email = testEmail,
+            MasterPasswordHash = _testPassword,
+        });
+
+        var userRepository = factory.Services.GetRequiredService<IUserRepository>();
+        var user = await userRepository.GetByEmailAsync(testEmail);
+        Assert.NotNull(user);
+
+        var userService = factory.GetService<IUserService>();
+        if (userTwoFactor != null)
+        {
+            user.TwoFactorProviders = userTwoFactor;
+            await userService.UpdateTwoFactorProviderAsync(user, TwoFactorProviderType.Email);
+            user = await userRepository.GetByEmailAsync(testEmail);
+            Assert.NotNull(user.TwoFactorProviders);
+        }
+    }
+
+    private async Task<IdentityApplicationFactory> CreateSsoOrganizationAndUserAsync(
+        IdentityApplicationFactory factory,
+        SsoConfigurationData ssoConfigurationData,
+        string challenge,
+        string testEmail,
+        string orgTwoFactor = null,
+        string userTwoFactor = null,
+        Permissions permissions = null)
+    {
+        var authorizationCode = new AuthorizationCode
+        {
+            ClientId = "web",
+            CreationTime = DateTime.UtcNow,
+            Lifetime = (int)TimeSpan.FromMinutes(5).TotalSeconds,
+            RedirectUri = "https://localhost:8080/sso-connector.html",
+            RequestedScopes = ["api", "offline_access"],
+            CodeChallenge = challenge.Sha256(),
+            CodeChallengeMethod = "plain",
+            Subject = null!, // Temporarily set it to null
+        };
+
+        factory.SubstituteService<IAuthorizationCodeStore>(service =>
+        {
+            service.GetAuthorizationCodeAsync("test_code")
+                .Returns(authorizationCode);
+        });
+
+        // Create Test User
+        var registerResponse = await factory.RegisterAsync(new RegisterRequestModel
+        {
+            Email = testEmail,
+            MasterPasswordHash = _testPassword,
+        });
+
+        var userRepository = factory.Services.GetRequiredService<IUserRepository>();
+        var user = await userRepository.GetByEmailAsync(testEmail);
+        Assert.NotNull(user);
+
+        var userService = factory.GetService<IUserService>();
+        if (userTwoFactor != null)
+        {
+            user.TwoFactorProviders = userTwoFactor;
+            await userService.UpdateTwoFactorProviderAsync(user, TwoFactorProviderType.Email);
+        }
+
+        // Create Organization
+        var organizationRepository = factory.Services.GetRequiredService<IOrganizationRepository>();
+        var organization = await organizationRepository.CreateAsync(new Organization
+        {
+            Name = "Test Org",
+            BillingEmail = "billing-email@example.com",
+            Plan = "Enterprise",
+            UsePolicies = true,
+            UseSso = true,
+            Use2fa = !string.IsNullOrEmpty(userTwoFactor) || !string.IsNullOrEmpty(orgTwoFactor),
+            TwoFactorProviders = orgTwoFactor,
+        });
+
+        if (orgTwoFactor != null)
+        {
+            factory.WithWebHostBuilder(builder =>
+            {
+                builder.UseSetting("globalSettings:Duo:AKey", "WJHB374KM3N5hglO9hniwbkibg$789EfbhNyLpNq1");
+            });
+        }
+
+        // Register User to Organization
+        var organizationUserRepository = factory.Services.GetRequiredService<IOrganizationUserRepository>();
+        var orgUserPermissions =
+            (permissions == null) ? null : JsonSerializer.Serialize(permissions, JsonHelpers.CamelCase);
+        var organizationUser = await organizationUserRepository.CreateAsync(new OrganizationUser
+        {
+            UserId = user.Id,
+            OrganizationId = organization.Id,
+            Status = OrganizationUserStatusType.Confirmed,
+            Type = OrganizationUserType.User,
+            Permissions = orgUserPermissions
+        });
+
+        // Configure SSO
+        var ssoConfigRepository = factory.Services.GetRequiredService<ISsoConfigRepository>();
+        await ssoConfigRepository.CreateAsync(new SsoConfig
+        {
+            OrganizationId = organization.Id,
+            Enabled = true,
+            Data = JsonSerializer.Serialize(ssoConfigurationData, JsonHelpers.CamelCase),
+        });
+
+        var subject = new ClaimsPrincipal(new ClaimsIdentity([
+            new Claim(JwtClaimTypes.Subject, user.Id.ToString()), // Get real user id
+            new Claim(JwtClaimTypes.Name, testEmail),
+            new Claim(JwtClaimTypes.IdentityProvider, "sso"),
+            new Claim("organizationId", organization.Id.ToString()),
+            new Claim(JwtClaimTypes.SessionId, "SOMETHING"),
+            new Claim(JwtClaimTypes.AuthenticationMethod, "external"),
+            new Claim(JwtClaimTypes.AuthenticationTime, DateTime.UtcNow.AddMinutes(-1).ToEpochTime().ToString())
+        ], "Duende.IdentityServer", JwtClaimTypes.Name, JwtClaimTypes.Role));
+
+        authorizationCode.Subject = subject;
+
+        return factory;
     }
 }

test/IntegrationTestCommon/Factories/IdentityApplicationFactory.cs

@@ -6,6 +6,7 @@ using Bit.Core.Utilities;
 using Bit.Identity;
 using Bit.Identity.Models.Request.Accounts;
 using Bit.Test.Common.Helpers;
+using HandlebarsDotNet;
 using Microsoft.AspNetCore.Http;
 
 namespace Bit.IntegrationTestCommon.Factories;
@@ -34,7 +35,25 @@ public class IdentityApplicationFactory : WebApplicationFactoryBase<Startup>
         return await Server.PostAsync("/accounts/register/verification-email-clicked", JsonContent.Create(model));
     }
 
-    public async Task<(string Token, string RefreshToken)> TokenFromPasswordAsync(string username,
+    public async Task<(string Token, string RefreshToken)> TokenFromPasswordAsync(
+        string username,
+        string password,
+        string deviceIdentifier = DefaultDeviceIdentifier,
+        string clientId = "web",
+        DeviceType deviceType = DeviceType.FirefoxBrowser,
+        string deviceName = "firefox")
+    {
+        var context = await ContextFromPasswordAsync(
+            username, password, deviceIdentifier, clientId, deviceType, deviceName);
+
+        using var body = await AssertHelper.AssertResponseTypeIs<JsonDocument>(context);
+        var root = body.RootElement;
+
+        return (root.GetProperty("access_token").GetString(), root.GetProperty("refresh_token").GetString());
+    }
+
+    public async Task<HttpContext> ContextFromPasswordAsync(
+        string username,
         string password,
         string deviceIdentifier = DefaultDeviceIdentifier,
         string clientId = "web",
@@ -53,14 +72,50 @@ public class IdentityApplicationFactory : WebApplicationFactoryBase<Startup>
             { "password", password },
         }), context => context.Request.Headers.Append("Auth-Email", CoreHelpers.Base64UrlEncodeString(username)));
 
-        using var body = await AssertHelper.AssertResponseTypeIs<JsonDocument>(context);
-        var root = body.RootElement;
+        return context;
+    }
 
-        return (root.GetProperty("access_token").GetString(), root.GetProperty("refresh_token").GetString());
+    public async Task<HttpContext> ContextFromPasswordWithTwoFactorAsync(
+        string username,
+        string password,
+        string deviceIdentifier = DefaultDeviceIdentifier,
+        string clientId = "web",
+        DeviceType deviceType = DeviceType.FirefoxBrowser,
+        string deviceName = "firefox",
+        string twoFactorProviderType = "Email",
+        string twoFactorToken = "two-factor-token")
+    {
+        var context = await Server.PostAsync("/connect/token", new FormUrlEncodedContent(new Dictionary<string, string>
+        {
+            { "scope", "api offline_access" },
+            { "client_id", clientId },
+            { "deviceType", ((int)deviceType).ToString() },
+            { "deviceIdentifier", deviceIdentifier },
+            { "deviceName", deviceName },
+            { "grant_type", "password" },
+            { "username", username },
+            { "password", password },
+            { "TwoFactorToken", twoFactorToken },
+            { "TwoFactorProvider", twoFactorProviderType },
+            { "TwoFactorRemember", "1" },
+        }), context => context.Request.Headers.Append("Auth-Email", CoreHelpers.Base64UrlEncodeString(username)));
+
+        return context;
     }
 
     public async Task<string> TokenFromAccessTokenAsync(Guid clientId, string clientSecret,
         DeviceType deviceType = DeviceType.SDK)
+    {
+        var context = await ContextFromAccessTokenAsync(clientId, clientSecret, deviceType);
+
+        using var body = await AssertHelper.AssertResponseTypeIs<JsonDocument>(context);
+        var root = body.RootElement;
+
+        return root.GetProperty("access_token").GetString();
+    }
+
+    public async Task<HttpContext> ContextFromAccessTokenAsync(Guid clientId, string clientSecret,
+        DeviceType deviceType = DeviceType.SDK)
     {
         var context = await Server.PostAsync("/connect/token",
             new FormUrlEncodedContent(new Dictionary<string, string>
@@ -72,13 +127,21 @@ public class IdentityApplicationFactory : WebApplicationFactoryBase<Startup>
                 { "deviceType", ((int)deviceType).ToString() }
             }));
 
+        return context;
+    }
+
+    public async Task<string> TokenFromOrganizationApiKeyAsync(string clientId, string clientSecret,
+        DeviceType deviceType = DeviceType.FirefoxBrowser)
+    {
+        var context = await ContextFromOrganizationApiKeyAsync(clientId, clientSecret, deviceType);
+
         using var body = await AssertHelper.AssertResponseTypeIs<JsonDocument>(context);
         var root = body.RootElement;
 
         return root.GetProperty("access_token").GetString();
     }
 
-    public async Task<string> TokenFromOrganizationApiKeyAsync(string clientId, string clientSecret,
+    public async Task<HttpContext> ContextFromOrganizationApiKeyAsync(string clientId, string clientSecret,
         DeviceType deviceType = DeviceType.FirefoxBrowser)
     {
         var context = await Server.PostAsync("/connect/token",
@@ -90,10 +153,6 @@ public class IdentityApplicationFactory : WebApplicationFactoryBase<Startup>
                 { "grant_type", "client_credentials" },
                 { "deviceType", ((int)deviceType).ToString() }
             }));
-
-        using var body = await AssertHelper.AssertResponseTypeIs<JsonDocument>(context);
-        var root = body.RootElement;
-
-        return root.GetProperty("access_token").GetString();
+        return context;
     }
 }