From 5c7d47402d2424f32bcfe86434065531db743975 Mon Sep 17 00:00:00 2001
From: Thomas Rittson <trittson@bitwarden.com>
Date: Mon, 7 Oct 2024 12:55:23 +1000
Subject: [PATCH] Fix missing validation

---
 .../RequireSsoPolicyDefinition.cs             | 32 +++++++++++++++++++
 .../ResetPasswordPolicyDefinition.cs          | 19 +++++++----
 .../SingleOrgPolicyDefinition.cs              | 14 ++++++--
 3 files changed, 56 insertions(+), 9 deletions(-)

diff --git a/src/Core/AdminConsole/OrganizationFeatures/Policies/Implementations/RequireSsoPolicyDefinition.cs b/src/Core/AdminConsole/OrganizationFeatures/Policies/Implementations/RequireSsoPolicyDefinition.cs
index 02bfd7d4cc..8bf74947ad 100644
--- a/src/Core/AdminConsole/OrganizationFeatures/Policies/Implementations/RequireSsoPolicyDefinition.cs
+++ b/src/Core/AdminConsole/OrganizationFeatures/Policies/Implementations/RequireSsoPolicyDefinition.cs
@@ -1,11 +1,43 @@
 #nullable enable
 
+using Bit.Core.AdminConsole.Entities;
 using Bit.Core.AdminConsole.Enums;
+using Bit.Core.Auth.Enums;
+using Bit.Core.Auth.Repositories;
 
 namespace Bit.Core.AdminConsole.OrganizationFeatures.Policies.Implementations;
 
 public class RequireSsoPolicyDefinition : IPolicyDefinition
 {
+    private readonly ISsoConfigRepository _ssoConfigRepository;
+
+    public RequireSsoPolicyDefinition(ISsoConfigRepository ssoConfigRepository)
+    {
+        _ssoConfigRepository = ssoConfigRepository;
+    }
+
     public PolicyType Type => PolicyType.RequireSso;
     public IEnumerable<PolicyType> RequiredPolicies => [PolicyType.SingleOrg];
+
+    public async Task<string?> ValidateAsync(Policy? currentPolicy, Policy modifiedPolicy)
+    {
+        if (modifiedPolicy is not { Enabled: true })
+        {
+            return await ValidateDisableAsync(modifiedPolicy);
+        }
+
+        return null;
+    }
+
+    private async Task<string?> ValidateDisableAsync(Policy policy)
+    {
+        // Do not allow this policy to be disabled if Key Connector or TDE are being used
+        var ssoConfig = await _ssoConfigRepository.GetByOrganizationIdAsync(policy.OrganizationId);
+        return ssoConfig?.GetData().MemberDecryptionType switch
+        {
+            MemberDecryptionType.KeyConnector => "Key Connector is enabled.",
+            MemberDecryptionType.TrustedDeviceEncryption => "Trusted device encryption is on and requires this policy.",
+            _ => null
+        };
+    }
 }
diff --git a/src/Core/AdminConsole/OrganizationFeatures/Policies/Implementations/ResetPasswordPolicyDefinition.cs b/src/Core/AdminConsole/OrganizationFeatures/Policies/Implementations/ResetPasswordPolicyDefinition.cs
index 7f900e6380..9a350351b4 100644
--- a/src/Core/AdminConsole/OrganizationFeatures/Policies/Implementations/ResetPasswordPolicyDefinition.cs
+++ b/src/Core/AdminConsole/OrganizationFeatures/Policies/Implementations/ResetPasswordPolicyDefinition.cs
@@ -22,13 +22,20 @@ public class ResetPasswordPolicyDefinition : IPolicyDefinition
     public async Task<string?> ValidateAsync(Policy? currentPolicy, Policy modifiedPolicy)
     {
         if (modifiedPolicy is not { Enabled:true } ||
-            modifiedPolicy.GetDataModel<ResetPasswordDataModel>()?.AutoEnrollEnabled == false)
+            modifiedPolicy.GetDataModel<ResetPasswordDataModel>().AutoEnrollEnabled == false)
         {
-            var ssoConfig = await _ssoConfigRepository.GetByOrganizationIdAsync(modifiedPolicy.OrganizationId);
-            if (ssoConfig?.GetData()?.MemberDecryptionType == MemberDecryptionType.TrustedDeviceEncryption)
-            {
-                return "Trusted device encryption is on and requires this policy.";
-            }
+            return await ValidateDisableAsync(modifiedPolicy);
+        }
+
+        return null;
+    }
+
+    private async Task<string?> ValidateDisableAsync(Policy policy)
+    {
+        var ssoConfig = await _ssoConfigRepository.GetByOrganizationIdAsync(policy.OrganizationId);
+        if (ssoConfig?.GetData().MemberDecryptionType == MemberDecryptionType.TrustedDeviceEncryption)
+        {
+            return "Trusted device encryption is on and requires this policy.";
         }
 
         return null;
diff --git a/src/Core/AdminConsole/OrganizationFeatures/Policies/Implementations/SingleOrgPolicyDefinition.cs b/src/Core/AdminConsole/OrganizationFeatures/Policies/Implementations/SingleOrgPolicyDefinition.cs
index 2896388756..aaa6f43a2e 100644
--- a/src/Core/AdminConsole/OrganizationFeatures/Policies/Implementations/SingleOrgPolicyDefinition.cs
+++ b/src/Core/AdminConsole/OrganizationFeatures/Policies/Implementations/SingleOrgPolicyDefinition.cs
@@ -83,11 +83,19 @@ public class SingleOrgPolicyDefinition : IPolicyDefinition
 
     public async Task<string?> ValidateAsync(Policy? currentPolicy, Policy modifiedPolicy)
     {
-        var organizationId = modifiedPolicy.OrganizationId;
+        if (modifiedPolicy is not { Enabled: true })
+        {
+            return await ValidateDisableAsync(modifiedPolicy);
+        }
 
+        return null;
+    }
+
+    private async Task<string?> ValidateDisableAsync(Policy policy)
+    {
         // Do not allow this policy to be disabled if Key Connector is being used
-        var ssoConfig = await _ssoConfigRepository.GetByOrganizationIdAsync(organizationId);
-        if (ssoConfig?.GetData()?.MemberDecryptionType == MemberDecryptionType.KeyConnector)
+        var ssoConfig = await _ssoConfigRepository.GetByOrganizationIdAsync(policy.OrganizationId);
+        if (ssoConfig?.GetData().MemberDecryptionType == MemberDecryptionType.KeyConnector)
         {
             return "Key Connector is enabled.";
         }