From c852575a9eaf67789eb0bbb0c80b10ff7de018b8 Mon Sep 17 00:00:00 2001
From: Alex Morask <144709477+amorask-bitwarden@users.noreply.github.com>
Date: Thu, 12 Dec 2024 07:08:17 -0500
Subject: [PATCH] [PM-14984] Use provider subscription for MSP managed
 enterprise license (#5102)

* Use provider subscription when creating license for MSP managed enterprise organization

* Run dotnet format
---
 .../Cloud/CloudGetOrganizationLicenseQuery.cs | 20 +++++++++--
 .../CloudGetOrganizationLicenseQueryTests.cs  | 33 +++++++++++++++++++
 2 files changed, 51 insertions(+), 2 deletions(-)

diff --git a/src/Core/OrganizationFeatures/OrganizationLicenses/Cloud/CloudGetOrganizationLicenseQuery.cs b/src/Core/OrganizationFeatures/OrganizationLicenses/Cloud/CloudGetOrganizationLicenseQuery.cs
index a4b08736c2..d7782fcd98 100644
--- a/src/Core/OrganizationFeatures/OrganizationLicenses/Cloud/CloudGetOrganizationLicenseQuery.cs
+++ b/src/Core/OrganizationFeatures/OrganizationLicenses/Cloud/CloudGetOrganizationLicenseQuery.cs
@@ -1,4 +1,6 @@
 using Bit.Core.AdminConsole.Entities;
+using Bit.Core.AdminConsole.Repositories;
+using Bit.Core.Enums;
 using Bit.Core.Exceptions;
 using Bit.Core.Models.Business;
 using Bit.Core.OrganizationFeatures.OrganizationLicenses.Interfaces;
@@ -12,15 +14,18 @@ public class CloudGetOrganizationLicenseQuery : ICloudGetOrganizationLicenseQuer
     private readonly IInstallationRepository _installationRepository;
     private readonly IPaymentService _paymentService;
     private readonly ILicensingService _licensingService;
+    private readonly IProviderRepository _providerRepository;
 
     public CloudGetOrganizationLicenseQuery(
         IInstallationRepository installationRepository,
         IPaymentService paymentService,
-        ILicensingService licensingService)
+        ILicensingService licensingService,
+        IProviderRepository providerRepository)
     {
         _installationRepository = installationRepository;
         _paymentService = paymentService;
         _licensingService = licensingService;
+        _providerRepository = providerRepository;
     }
 
     public async Task<OrganizationLicense> GetLicenseAsync(Organization organization, Guid installationId,
@@ -32,11 +37,22 @@ public class CloudGetOrganizationLicenseQuery : ICloudGetOrganizationLicenseQuer
             throw new BadRequestException("Invalid installation id");
         }
 
-        var subscriptionInfo = await _paymentService.GetSubscriptionAsync(organization);
+        var subscriptionInfo = await GetSubscriptionAsync(organization);
 
         return new OrganizationLicense(organization, subscriptionInfo, installationId, _licensingService, version)
         {
             Token = await _licensingService.CreateOrganizationTokenAsync(organization, installationId, subscriptionInfo)
         };
     }
+
+    private async Task<SubscriptionInfo> GetSubscriptionAsync(Organization organization)
+    {
+        if (organization is not { Status: OrganizationStatusType.Managed })
+        {
+            return await _paymentService.GetSubscriptionAsync(organization);
+        }
+
+        var provider = await _providerRepository.GetByOrganizationIdAsync(organization.Id);
+        return await _paymentService.GetSubscriptionAsync(provider);
+    }
 }
diff --git a/test/Core.Test/OrganizationFeatures/OrganizationLicenses/CloudGetOrganizationLicenseQueryTests.cs b/test/Core.Test/OrganizationFeatures/OrganizationLicenses/CloudGetOrganizationLicenseQueryTests.cs
index 00a4b12b2e..52bee7068f 100644
--- a/test/Core.Test/OrganizationFeatures/OrganizationLicenses/CloudGetOrganizationLicenseQueryTests.cs
+++ b/test/Core.Test/OrganizationFeatures/OrganizationLicenses/CloudGetOrganizationLicenseQueryTests.cs
@@ -1,4 +1,6 @@
 using Bit.Core.AdminConsole.Entities;
+using Bit.Core.AdminConsole.Entities.Provider;
+using Bit.Core.AdminConsole.Repositories;
 using Bit.Core.Entities;
 using Bit.Core.Enums;
 using Bit.Core.Exceptions;
@@ -11,6 +13,7 @@ using Bit.Test.Common.AutoFixture;
 using Bit.Test.Common.AutoFixture.Attributes;
 using NSubstitute;
 using NSubstitute.ReturnsExtensions;
+using Stripe;
 using Xunit;
 
 namespace Bit.Core.Test.OrganizationFeatures.OrganizationLicenses;
@@ -62,4 +65,34 @@ public class CloudGetOrganizationLicenseQueryTests
         Assert.Equal(installationId, result.InstallationId);
         Assert.Equal(licenseSignature, result.SignatureBytes);
     }
+
+    [Theory]
+    [BitAutoData]
+    public async Task GetLicenseAsync_MSPManagedOrganization_UsesProviderSubscription(SutProvider<CloudGetOrganizationLicenseQuery> sutProvider,
+        Organization organization, Guid installationId, Installation installation, SubscriptionInfo subInfo,
+        byte[] licenseSignature, Provider provider)
+    {
+        organization.Status = OrganizationStatusType.Managed;
+        organization.ExpirationDate = null;
+
+        subInfo.Subscription = new SubscriptionInfo.BillingSubscription(new Subscription
+        {
+            CurrentPeriodStart = DateTime.UtcNow,
+            CurrentPeriodEnd = DateTime.UtcNow.AddMonths(1)
+        });
+
+        installation.Enabled = true;
+        sutProvider.GetDependency<IInstallationRepository>().GetByIdAsync(installationId).Returns(installation);
+        sutProvider.GetDependency<IProviderRepository>().GetByOrganizationIdAsync(organization.Id).Returns(provider);
+        sutProvider.GetDependency<IPaymentService>().GetSubscriptionAsync(provider).Returns(subInfo);
+        sutProvider.GetDependency<ILicensingService>().SignLicense(Arg.Any<ILicense>()).Returns(licenseSignature);
+
+        var result = await sutProvider.Sut.GetLicenseAsync(organization, installationId);
+
+        Assert.Equal(LicenseType.Organization, result.LicenseType);
+        Assert.Equal(organization.Id, result.Id);
+        Assert.Equal(installationId, result.InstallationId);
+        Assert.Equal(licenseSignature, result.SignatureBytes);
+        Assert.Equal(DateTime.UtcNow.AddYears(1).Date, result.Expires!.Value.Date);
+    }
 }