Refactor AcceptOrgUserCommand to enforce two-factor authentication policy based on feature flag; update validation logic and tests accordingly.

src/Core/AdminConsole/OrganizationFeatures/OrganizationUsers/AcceptOrgUserCommand.cs

@@ -202,9 +202,21 @@ public class AcceptOrgUserCommand : IAcceptOrgUserCommand
         }
 
         // Enforce Two Factor Authentication Policy of organization user is trying to join
-        if (!await _twoFactorIsEnabledQuery.TwoFactorIsEnabledAsync(user))
+        if (_featureService.IsEnabled(FeatureFlagKeys.PolicyRequirements))
         {
-            await ValidateTwoFactorAuthenticationPolicy(user, orgUser.OrganizationId);
+            await ValidateTwoFactorAuthenticationPolicyAsync(user, orgUser.OrganizationId);
+        }
+        else
+        {
+            if (!await _twoFactorIsEnabledQuery.TwoFactorIsEnabledAsync(user))
+            {
+                var invitedTwoFactorPolicies = await _policyService.GetPoliciesApplicableToUserAsync(user.Id,
+                    PolicyType.TwoFactorAuthentication, OrganizationUserStatusType.Invited);
+                if (invitedTwoFactorPolicies.Any(p => p.OrganizationId == orgUser.OrganizationId))
+                {
+                    throw new BadRequestException("You cannot join this organization until you enable two-step login on your user account.");
+                }
+            }
         }
 
         orgUser.Status = OrganizationUserStatusType.Accepted;
@@ -225,24 +237,28 @@ public class AcceptOrgUserCommand : IAcceptOrgUserCommand
         return orgUser;
     }
 
-    private async Task ValidateTwoFactorAuthenticationPolicy(User user, Guid organizationId)
+    /// <summary>
+    /// Validates the two-factor authentication policy for the organization user.
+    /// If the user has two-step login enabled, the policy is not enforced.
+    /// </summary>
+    /// <param name="user">The user to validate the policy for.</param>
+    /// <param name="organizationId">The ID of the organization to validate the policy for.</param>
+    /// <exception cref="BadRequestException">Thrown if the user does not have two-step login enabled.</exception>
+    private async Task ValidateTwoFactorAuthenticationPolicyAsync(User user, Guid organizationId)
     {
-        if (_featureService.IsEnabled(FeatureFlagKeys.PolicyRequirements))
+        var userTwoFactorEnabled = await _twoFactorIsEnabledQuery.TwoFactorIsEnabledAsync(user);
+        if (userTwoFactorEnabled)
         {
-            var requireTwoFactorPolicyRequirement = await _policyRequirementQuery.GetAsync<RequireTwoFactorPolicyRequirement>(user.Id);
+            // If the user has two-step login enabled, the policy is not enforced.
-            if (requireTwoFactorPolicyRequirement.RequireTwoFactor)
+            return;
-            {
-                throw new BadRequestException("You cannot join this organization until you enable two-step login on your user account.");
-            }
         }
-        else
+
+        var requirement = await _policyRequirementQuery.GetAsync<RequireTwoFactorPolicyRequirement>(user.Id);
+        var canAcceptInvitation = requirement.CanAcceptInvitation(userTwoFactorEnabled, organizationId);
+
+        if (!canAcceptInvitation)
         {
-            var invitedTwoFactorPolicies = await _policyService.GetPoliciesApplicableToUserAsync(user.Id,
+            throw new BadRequestException("You cannot join this organization until you enable two-step login on your user account.");
-                PolicyType.TwoFactorAuthentication, OrganizationUserStatusType.Invited);
-            if (invitedTwoFactorPolicies.Any(p => p.OrganizationId == organizationId))
-            {
-                throw new BadRequestException("You cannot join this organization until you enable two-step login on your user account.");
-            }
         }
     }
 }

test/Core.Test/AdminConsole/OrganizationFeatures/OrganizationUsers/AcceptOrgUserCommandTests.cs

@@ -1,5 +1,6 @@
 using Bit.Core.AdminConsole.Entities;
 using Bit.Core.AdminConsole.Enums;
+using Bit.Core.AdminConsole.Models.Data.Organizations.Policies;
 using Bit.Core.AdminConsole.OrganizationFeatures.Policies;
 using Bit.Core.AdminConsole.Services;
 using Bit.Core.Auth.Models.Business.Tokenables;
@@ -30,7 +31,6 @@ namespace Bit.Core.Test.OrganizationFeatures.OrganizationUsers;
 public class AcceptOrgUserCommandTests
 {
     private readonly IUserService _userService = Substitute.For<IUserService>();
-    private readonly ITwoFactorIsEnabledQuery _twoFactorIsEnabledQuery = Substitute.For<ITwoFactorIsEnabledQuery>();
     private readonly IOrgUserInviteTokenableFactory _orgUserInviteTokenableFactory = Substitute.For<IOrgUserInviteTokenableFactory>();
     private readonly IDataProtectorTokenFactory<OrgUserInviteTokenable> _orgUserInviteTokenDataFactory = new FakeDataProtectorTokenFactory<OrgUserInviteTokenable>();
 
@@ -168,7 +168,7 @@ public class AcceptOrgUserCommandTests
         SetupCommonAcceptOrgUserMocks(sutProvider, user, org, orgUser, adminUserDetails);
 
         // User doesn't have 2FA enabled
-        _twoFactorIsEnabledQuery.TwoFactorIsEnabledAsync(user).Returns(false);
+        sutProvider.GetDependency<ITwoFactorIsEnabledQuery>().TwoFactorIsEnabledAsync(user).Returns(false);
 
         // Organization they are trying to join requires 2FA
         var twoFactorPolicy = new OrganizationUserPolicyDetails { OrganizationId = orgUser.OrganizationId };
@@ -192,7 +192,6 @@ public class AcceptOrgUserCommandTests
         SutProvider<AcceptOrgUserCommand> sutProvider,
         User user, Organization org, OrganizationUser orgUser, OrganizationUserUserDetails adminUserDetails)
     {
-        // Arrange
         SetupCommonAcceptOrgUserMocks(sutProvider, user, org, orgUser, adminUserDetails);
 
         sutProvider.GetDependency<IFeatureService>()
@@ -200,14 +199,23 @@ public class AcceptOrgUserCommandTests
             .Returns(true);
 
         // User doesn't have 2FA enabled
-        _userService.TwoFactorIsEnabledAsync(user).Returns(false);
+        sutProvider.GetDependency<ITwoFactorIsEnabledQuery>()
+            .TwoFactorIsEnabledAsync(user)
+            .Returns(false);
 
         // Organization they are trying to join requires 2FA
         sutProvider.GetDependency<IPolicyRequirementQuery>()
             .GetAsync<RequireTwoFactorPolicyRequirement>(user.Id)
-            .Returns(new RequireTwoFactorPolicyRequirement { RequireTwoFactor = true });
+            .Returns(new RequireTwoFactorPolicyRequirement(
+            [
+                new PolicyDetails
+                {
+                    OrganizationId = orgUser.OrganizationId,
+                    OrganizationUserStatus = OrganizationUserStatusType.Invited,
+                    PolicyType = PolicyType.TwoFactorAuthentication,
+                }
+            ]));
 
-        // Act & Assert
         var exception = await Assert.ThrowsAsync<BadRequestException>(() =>
             sutProvider.Sut.AcceptOrgUserAsync(orgUser, user, _userService));
 
@@ -228,12 +236,20 @@ public class AcceptOrgUserCommandTests
             .Returns(true);
 
         // User has 2FA enabled
-        _userService.TwoFactorIsEnabledAsync(user).Returns(true);
+        sutProvider.GetDependency<ITwoFactorIsEnabledQuery>().TwoFactorIsEnabledAsync(Arg.Any<User>()).Returns(true);
 
         // Organization they are trying to join requires 2FA
         sutProvider.GetDependency<IPolicyRequirementQuery>()
             .GetAsync<RequireTwoFactorPolicyRequirement>(user.Id)
-            .Returns(new RequireTwoFactorPolicyRequirement { RequireTwoFactor = true });
+            .Returns(new RequireTwoFactorPolicyRequirement(
+            [
+                new PolicyDetails
+                {
+                    OrganizationId = orgUser.OrganizationId,
+                    OrganizationUserStatus = OrganizationUserStatusType.Invited,
+                    PolicyType = PolicyType.TwoFactorAuthentication,
+                }
+            ]));
 
         await sutProvider.Sut.AcceptOrgUserAsync(orgUser, user, _userService);
 
@@ -253,26 +269,34 @@ public class AcceptOrgUserCommandTests
         bool userTwoFactorEnabled, SutProvider<AcceptOrgUserCommand> sutProvider,
         User user, Organization org, OrganizationUser orgUser, OrganizationUserUserDetails adminUserDetails)
     {
-        // Arrange
         SetupCommonAcceptOrgUserMocks(sutProvider, user, org, orgUser, adminUserDetails);
 
         sutProvider.GetDependency<IFeatureService>()
             .IsEnabled(FeatureFlagKeys.PolicyRequirements)
             .Returns(true);
 
-        // User doesn't have 2FA enabled
+        // User 2FA status
-        _userService.TwoFactorIsEnabledAsync(user).Returns(userTwoFactorEnabled);
+        sutProvider.GetDependency<ITwoFactorIsEnabledQuery>()
+            .TwoFactorIsEnabledAsync(user)
+            .Returns(userTwoFactorEnabled);
 
         // Organization they are trying to join doesn't require 2FA
         sutProvider.GetDependency<IPolicyRequirementQuery>()
             .GetAsync<RequireTwoFactorPolicyRequirement>(user.Id)
-            .Returns(new RequireTwoFactorPolicyRequirement { RequireTwoFactor = false });
+            .Returns(new RequireTwoFactorPolicyRequirement(
+            [
+                new PolicyDetails
+                {
+                    OrganizationId = Guid.NewGuid(),
+                    OrganizationUserStatus = OrganizationUserStatusType.Invited,
+                    PolicyType = PolicyType.TwoFactorAuthentication,
+                }
+            ]));
 
-        // Act
         await sutProvider.Sut.AcceptOrgUserAsync(orgUser, user, _userService);
 
-        // Assert
+        await sutProvider.GetDependency<IOrganizationUserRepository>()
-        await sutProvider.GetDependency<IOrganizationUserRepository>().Received(1)
+            .Received(1)
             .ReplaceAsync(Arg.Is<OrganizationUser>(ou => ou.Status == OrganizationUserStatusType.Accepted));
     }
 
@@ -739,7 +763,9 @@ public class AcceptOrgUserCommandTests
             .Returns(false);
 
         // User doesn't have 2FA enabled
-        _twoFactorIsEnabledQuery.TwoFactorIsEnabledAsync(user).Returns(false);
+        sutProvider.GetDependency<ITwoFactorIsEnabledQuery>()
+            .TwoFactorIsEnabledAsync(user)
+            .Returns(false);
 
         // Org does not require 2FA
         sutProvider.GetDependency<IPolicyService>().GetPoliciesApplicableToUserAsync(user.Id,