[PM-18238] Add RequireTwoFactorPolicyRequirement (#5840)

* Add RequireTwoFactorPolicyRequirement and its factory with unit tests

* Implemented RequireTwoFactorPolicyRequirement to enforce two-factor authentication policies.
* Created RequireTwoFactorPolicyRequirementFactory to generate policy requirements based on user status.
* Added unit tests for the factory to validate behavior with various user statuses and policy details.

* Enhance AcceptOrgUserCommand to use IPolicyRequirementQuery for two-factor authentication validation

* Update ConfirmOrganizationUserCommand to use RequireTwoFactorPolicyRequirement to check for 2FA requirement

* Implement CanAcceptInvitation and CanBeConfirmed methods in RequireTwoFactorPolicyRequirement; update tests to reflect new logic for two-factor authentication policy handling.

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

* Enhance ConfirmOrganizationUserCommand to validate two-factor authentication policy based on feature flag; refactor validation logic and update related tests for improved policy handling.

* Remove unused method and its dependencies from OrganizationService.

* Implement CanBeRestored method in RequireTwoFactorPolicyRequirement to determine user restoration eligibility based on two-factor authentication status; add corresponding unit tests for various scenarios.

* Update RestoreOrganizationUserCommand to use IPolicyRequirementQuery for two-factor authentication policies checks

* Remove redundant vNext tests

* Add TwoFactorPoliciesForActiveMemberships property to RequireTwoFactorPolicyRequirement and corresponding unit tests for policy retrieval based on user status

* Refactor UserService to integrate IPolicyRequirementQuery for two-factor authentication policy checks

* Add XML documentation for TwoFactorPoliciesForActiveMemberships property in RequireTwoFactorPolicyRequirement to clarify its purpose and return value.

* Add exception documentation for ValidateTwoFactorAuthenticationPolicyAsync method in ConfirmOrganizationUserCommand to clarify error handling for users without two-step login enabled.

* Update comments in AcceptOrgUserCommand and ConfirmOrganizationUserCommand to clarify handling of two-step login and 2FA policy checks.

* Add RequireTwoFactorPolicyRequirementFactory to PolicyServiceCollectionExtensions

* Refactor two-factor authentication policy checks in AcceptOrgUserCommand and ConfirmOrganizationUserCommand to streamline validation logic and improve clarity. Update RequireTwoFactorPolicyRequirement to provide a method for checking if two-factor authentication is required for an organization. Adjust related unit tests accordingly.

* Add PolicyRequirements namespace

* Update comments in AcceptOrgUserCommand and ConfirmOrganizationUserCommand to clarify two-factor authentication policy requirements and exception handling.

* Refactor RequireTwoFactorPolicyRequirement to return tuples of (OrganizationId, OrganizationUserId) for active memberships requiring two-factor authentication. Update UserService and related tests to reflect this change.

* Refactor AcceptOrgUserCommand: delegate feature flag check to the ValidateTwoFactorAuthenticationPolicyAsync method

* Skip policy check if two-step login is enabled for the user

* Refactor ConfirmOrganizationUserCommand to streamline two-factor authentication policy validation logic

* Refactor AcceptOrgUserCommand to simplify two-factor authentication check by removing intermediate variable

* Update documentation in RequireTwoFactorPolicyRequirement to clarify the purpose of the IsTwoFactorRequiredForOrganization

* Refactor AcceptOrgUserCommandTests to remove redundant two-factor authentication checks and simplify test setup

* Refactor AcceptOrgUserCommand and ConfirmOrganizationUserCommand to streamline two-factor authentication checks by removing redundant conditions and simplifying logic flow.

* Rename removeOrgUserTasks variable in UserService

* Refactor RestoreOrganizationUserCommand to simplify two-factor authentication compliance checks by consolidating logic into a new method, IsTwoFactorRequiredForOrganizationAsync.

* Remove outdated two-factor authentication validation documentation from AcceptOrgUserCommand

* Invert two-factor compliance check in RestoreOrganizationUserCommand to ensure correct validation of organization user policies.

* Refactor UserService to enhance two-factor compliance checks by optimizing organization retrieval and logging when no organizations require two-factor authentication.

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

@@ -1,4 +1,6 @@
 using Bit.Core.AdminConsole.Enums;
+using Bit.Core.AdminConsole.OrganizationFeatures.Policies;
+using Bit.Core.AdminConsole.OrganizationFeatures.Policies.PolicyRequirements;
 using Bit.Core.AdminConsole.Services;
 using Bit.Core.Auth.Models.Business.Tokenables;
 using Bit.Core.Auth.UserFeatures.TwoFactorAuth.Interfaces;
@@ -27,6 +29,8 @@ public class AcceptOrgUserCommand : IAcceptOrgUserCommand
     private readonly IUserRepository _userRepository;
     private readonly ITwoFactorIsEnabledQuery _twoFactorIsEnabledQuery;
     private readonly IDataProtectorTokenFactory<OrgUserInviteTokenable> _orgUserInviteTokenDataFactory;
+    private readonly IFeatureService _featureService;
+    private readonly IPolicyRequirementQuery _policyRequirementQuery;
 
     public AcceptOrgUserCommand(
         IDataProtectionProvider dataProtectionProvider,
@@ -37,9 +41,10 @@ public class AcceptOrgUserCommand : IAcceptOrgUserCommand
         IMailService mailService,
         IUserRepository userRepository,
         ITwoFactorIsEnabledQuery twoFactorIsEnabledQuery,
-        IDataProtectorTokenFactory<OrgUserInviteTokenable> orgUserInviteTokenDataFactory)
+        IDataProtectorTokenFactory<OrgUserInviteTokenable> orgUserInviteTokenDataFactory,
+        IFeatureService featureService,
+        IPolicyRequirementQuery policyRequirementQuery)
     {
-
         // TODO: remove data protector when old token validation removed
         _dataProtector = dataProtectionProvider.CreateProtector(OrgUserInviteTokenable.DataProtectorPurpose);
         _globalSettings = globalSettings;
@@ -50,6 +55,8 @@ public class AcceptOrgUserCommand : IAcceptOrgUserCommand
         _userRepository = userRepository;
         _twoFactorIsEnabledQuery = twoFactorIsEnabledQuery;
         _orgUserInviteTokenDataFactory = orgUserInviteTokenDataFactory;
+        _featureService = featureService;
+        _policyRequirementQuery = policyRequirementQuery;
     }
 
     public async Task<OrganizationUser> AcceptOrgUserByEmailTokenAsync(Guid organizationUserId, User user, string emailToken,
@@ -196,15 +203,7 @@ public class AcceptOrgUserCommand : IAcceptOrgUserCommand
         }
 
         // Enforce Two Factor Authentication Policy of organization user is trying to join
-        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.");
-            }
-        }
+        await ValidateTwoFactorAuthenticationPolicyAsync(user, orgUser.OrganizationId);
 
         orgUser.Status = OrganizationUserStatusType.Accepted;
         orgUser.UserId = user.Id;
@@ -224,4 +223,33 @@ public class AcceptOrgUserCommand : IAcceptOrgUserCommand
         return orgUser;
     }
 
+    private async Task ValidateTwoFactorAuthenticationPolicyAsync(User user, Guid organizationId)
+    {
+        if (_featureService.IsEnabled(FeatureFlagKeys.PolicyRequirements))
+        {
+            if (await _twoFactorIsEnabledQuery.TwoFactorIsEnabledAsync(user))
+            {
+                // If the user has two-step login enabled, we skip checking the 2FA policy
+                return;
+            }
+
+            var twoFactorPolicyRequirement = await _policyRequirementQuery.GetAsync<RequireTwoFactorPolicyRequirement>(user.Id);
+            if (twoFactorPolicyRequirement.IsTwoFactorRequiredForOrganization(organizationId))
+            {
+                throw new BadRequestException("You cannot join this organization until you enable two-step login on your user account.");
+            }
+
+            return;
+        }
+
+        if (!await _twoFactorIsEnabledQuery.TwoFactorIsEnabledAsync(user))
+        {
+            var invitedTwoFactorPolicies = await _policyService.GetPoliciesApplicableToUserAsync(user.Id,
+                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.");
+            }
+        }
+    }
 }

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

@@ -1,5 +1,7 @@
 using Bit.Core.AdminConsole.Enums;
 using Bit.Core.AdminConsole.OrganizationFeatures.OrganizationUsers.Interfaces;
+using Bit.Core.AdminConsole.OrganizationFeatures.Policies;
+using Bit.Core.AdminConsole.OrganizationFeatures.Policies.PolicyRequirements;
 using Bit.Core.AdminConsole.Services;
 using Bit.Core.Auth.UserFeatures.TwoFactorAuth.Interfaces;
 using Bit.Core.Billing.Enums;
@@ -24,6 +26,8 @@ public class ConfirmOrganizationUserCommand : IConfirmOrganizationUserCommand
     private readonly IPushRegistrationService _pushRegistrationService;
     private readonly IPolicyService _policyService;
     private readonly IDeviceRepository _deviceRepository;
+    private readonly IPolicyRequirementQuery _policyRequirementQuery;
+    private readonly IFeatureService _featureService;
 
     public ConfirmOrganizationUserCommand(
         IOrganizationRepository organizationRepository,
@@ -35,7 +39,9 @@ public class ConfirmOrganizationUserCommand : IConfirmOrganizationUserCommand
         IPushNotificationService pushNotificationService,
         IPushRegistrationService pushRegistrationService,
         IPolicyService policyService,
-        IDeviceRepository deviceRepository)
+        IDeviceRepository deviceRepository,
+        IPolicyRequirementQuery policyRequirementQuery,
+        IFeatureService featureService)
     {
         _organizationRepository = organizationRepository;
         _organizationUserRepository = organizationUserRepository;
@@ -47,6 +53,8 @@ public class ConfirmOrganizationUserCommand : IConfirmOrganizationUserCommand
         _pushRegistrationService = pushRegistrationService;
         _policyService = policyService;
         _deviceRepository = deviceRepository;
+        _policyRequirementQuery = policyRequirementQuery;
+        _featureService = featureService;
     }
 
     public async Task<OrganizationUser> ConfirmUserAsync(Guid organizationId, Guid organizationUserId, string key,
@@ -118,8 +126,8 @@ public class ConfirmOrganizationUserCommand : IConfirmOrganizationUserCommand
                     }
                 }
 
-                var twoFactorEnabled = usersTwoFactorEnabled.FirstOrDefault(tuple => tuple.userId == user.Id).twoFactorIsEnabled;
-                await CheckPoliciesAsync(organizationId, user, orgUsers, twoFactorEnabled);
+                var userTwoFactorEnabled = usersTwoFactorEnabled.FirstOrDefault(tuple => tuple.userId == user.Id).twoFactorIsEnabled;
+                await CheckPoliciesAsync(organizationId, user, orgUsers, userTwoFactorEnabled);
                 orgUser.Status = OrganizationUserStatusType.Confirmed;
                 orgUser.Key = keys[orgUser.Id];
                 orgUser.Email = null;
@@ -142,15 +150,10 @@ public class ConfirmOrganizationUserCommand : IConfirmOrganizationUserCommand
     }
 
     private async Task CheckPoliciesAsync(Guid organizationId, User user,
-        ICollection<OrganizationUser> userOrgs, bool twoFactorEnabled)
+        ICollection<OrganizationUser> userOrgs, bool userTwoFactorEnabled)
     {
         // Enforce Two Factor Authentication Policy for this organization
-        var orgRequiresTwoFactor = (await _policyService.GetPoliciesApplicableToUserAsync(user.Id, PolicyType.TwoFactorAuthentication))
-            .Any(p => p.OrganizationId == organizationId);
-        if (orgRequiresTwoFactor && !twoFactorEnabled)
-        {
-            throw new BadRequestException("User does not have two-step login enabled.");
-        }
+        await ValidateTwoFactorAuthenticationPolicyAsync(user, organizationId, userTwoFactorEnabled);
 
         var hasOtherOrgs = userOrgs.Any(ou => ou.OrganizationId != organizationId);
         var singleOrgPolicies = await _policyService.GetPoliciesApplicableToUserAsync(user.Id, PolicyType.SingleOrg);
@@ -168,6 +171,33 @@ public class ConfirmOrganizationUserCommand : IConfirmOrganizationUserCommand
         }
     }
 
+    private async Task ValidateTwoFactorAuthenticationPolicyAsync(User user, Guid organizationId, bool userTwoFactorEnabled)
+    {
+        if (_featureService.IsEnabled(FeatureFlagKeys.PolicyRequirements))
+        {
+            if (userTwoFactorEnabled)
+            {
+                // If the user has two-step login enabled, we skip checking the 2FA policy
+                return;
+            }
+
+            var twoFactorPolicyRequirement = await _policyRequirementQuery.GetAsync<RequireTwoFactorPolicyRequirement>(user.Id);
+            if (twoFactorPolicyRequirement.IsTwoFactorRequiredForOrganization(organizationId))
+            {
+                throw new BadRequestException("User does not have two-step login enabled.");
+            }
+
+            return;
+        }
+
+        var orgRequiresTwoFactor = (await _policyService.GetPoliciesApplicableToUserAsync(user.Id, PolicyType.TwoFactorAuthentication))
+            .Any(p => p.OrganizationId == organizationId);
+        if (orgRequiresTwoFactor && !userTwoFactorEnabled)
+        {
+            throw new BadRequestException("User does not have two-step login enabled.");
+        }
+    }
+
     private async Task DeleteAndPushUserRegistrationAsync(Guid organizationId, Guid userId)
     {
         var devices = await GetUserDeviceIdsAsync(userId);

src/Core/AdminConsole/OrganizationFeatures/OrganizationUsers/RestoreUser/v1/RestoreOrganizationUserCommand.cs

@@ -1,5 +1,7 @@
 using Bit.Core.AdminConsole.Entities;
 using Bit.Core.AdminConsole.Enums;
+using Bit.Core.AdminConsole.OrganizationFeatures.Policies;
+using Bit.Core.AdminConsole.OrganizationFeatures.Policies.PolicyRequirements;
 using Bit.Core.AdminConsole.Services;
 using Bit.Core.Auth.UserFeatures.TwoFactorAuth.Interfaces;
 using Bit.Core.Billing.Enums;
@@ -22,7 +24,9 @@ public class RestoreOrganizationUserCommand(
     ITwoFactorIsEnabledQuery twoFactorIsEnabledQuery,
     IPolicyService policyService,
     IUserRepository userRepository,
-    IOrganizationService organizationService) : IRestoreOrganizationUserCommand
+    IOrganizationService organizationService,
+    IFeatureService featureService,
+    IPolicyRequirementQuery policyRequirementQuery) : IRestoreOrganizationUserCommand
 {
     public async Task RestoreUserAsync(OrganizationUser organizationUser, Guid? restoringUserId)
     {
@@ -270,12 +274,7 @@ public class RestoreOrganizationUserCommand(
         // Enforce 2FA Policy of organization user is trying to join
         if (!userHasTwoFactorEnabled)
         {
-            var invitedTwoFactorPolicies = await policyService.GetPoliciesApplicableToUserAsync(userId,
-                PolicyType.TwoFactorAuthentication, OrganizationUserStatusType.Revoked);
-            if (invitedTwoFactorPolicies.Any(p => p.OrganizationId == orgUser.OrganizationId))
-            {
-                twoFactorCompliant = false;
-            }
+            twoFactorCompliant = !await IsTwoFactorRequiredForOrganizationAsync(userId, orgUser.OrganizationId);
         }
 
         var user = await userRepository.GetByIdAsync(userId);
@@ -299,4 +298,17 @@ public class RestoreOrganizationUserCommand(
             throw new BadRequestException(user.Email + " is not compliant with the two-step login policy");
         }
     }
+
+    private async Task<bool> IsTwoFactorRequiredForOrganizationAsync(Guid userId, Guid organizationId)
+    {
+        if (featureService.IsEnabled(FeatureFlagKeys.PolicyRequirements))
+        {
+            var requirement = await policyRequirementQuery.GetAsync<RequireTwoFactorPolicyRequirement>(userId);
+            return requirement.IsTwoFactorRequiredForOrganization(organizationId);
+        }
+
+        var invitedTwoFactorPolicies = await policyService.GetPoliciesApplicableToUserAsync(userId,
+            PolicyType.TwoFactorAuthentication, OrganizationUserStatusType.Revoked);
+        return invitedTwoFactorPolicies.Any(p => p.OrganizationId == organizationId);
+    }
 }

src/Core/AdminConsole/OrganizationFeatures/Policies/PolicyRequirements/RequireTwoFactorPolicyRequirement.cs

@@ -0,0 +1,52 @@
+using Bit.Core.AdminConsole.Enums;
+using Bit.Core.AdminConsole.Models.Data.Organizations.Policies;
+using Bit.Core.Enums;
+
+namespace Bit.Core.AdminConsole.OrganizationFeatures.Policies.PolicyRequirements;
+
+/// <summary>
+/// Policy requirements for the Require Two-Factor Authentication policy.
+/// </summary>
+public class RequireTwoFactorPolicyRequirement : IPolicyRequirement
+{
+    private readonly IEnumerable<PolicyDetails> _policyDetails;
+
+    public RequireTwoFactorPolicyRequirement(IEnumerable<PolicyDetails> policyDetails)
+    {
+        _policyDetails = policyDetails;
+    }
+
+    /// <summary>
+    /// Checks if two-factor authentication is required for the organization due to an active policy.
+    /// </summary>
+    /// <param name="organizationId">The ID of the organization to check.</param>
+    /// <returns>True if two-factor authentication is required for the organization, false otherwise.</returns>
+    /// <remarks>
+    /// This should be used to check whether the member needs to have 2FA enabled before being
+    /// accepted, confirmed, or restored to the organization.
+    /// </remarks>
+    public bool IsTwoFactorRequiredForOrganization(Guid organizationId) =>
+        _policyDetails.Any(p => p.OrganizationId == organizationId);
+
+    /// <summary>
+    /// Returns tuples of (OrganizationId, OrganizationUserId) for active memberships where two-factor authentication is required.
+    /// Users should be revoked from these organizations if they disable all 2FA methods.
+    /// </summary>
+    public IEnumerable<(Guid OrganizationId, Guid OrganizationUserId)> OrganizationsRequiringTwoFactor =>
+        _policyDetails
+            .Where(p => p.OrganizationUserStatus is
+                OrganizationUserStatusType.Accepted or
+                OrganizationUserStatusType.Confirmed)
+            .Select(p => (p.OrganizationId, p.OrganizationUserId));
+}
+
+public class RequireTwoFactorPolicyRequirementFactory : BasePolicyRequirementFactory<RequireTwoFactorPolicyRequirement>
+{
+    public override PolicyType PolicyType => PolicyType.TwoFactorAuthentication;
+    protected override IEnumerable<OrganizationUserStatusType> ExemptStatuses => [];
+
+    public override RequireTwoFactorPolicyRequirement Create(IEnumerable<PolicyDetails> policyDetails)
+    {
+        return new RequireTwoFactorPolicyRequirement(policyDetails);
+    }
+}

src/Core/AdminConsole/OrganizationFeatures/Policies/PolicyServiceCollectionExtensions.cs

@@ -36,5 +36,6 @@ public static class PolicyServiceCollectionExtensions
         services.AddScoped<IPolicyRequirementFactory<IPolicyRequirement>, ResetPasswordPolicyRequirementFactory>();
         services.AddScoped<IPolicyRequirementFactory<IPolicyRequirement>, PersonalOwnershipPolicyRequirementFactory>();
         services.AddScoped<IPolicyRequirementFactory<IPolicyRequirement>, RequireSsoPolicyRequirementFactory>();
+        services.AddScoped<IPolicyRequirementFactory<IPolicyRequirement>, RequireTwoFactorPolicyRequirementFactory>();
     }
 }