server/src/Billing/Services/Implementations/SubscriptionUpdatedHandler.cs

using Bit.Billing.Constants;
using Bit.Billing.Jobs;
using Bit.Core.AdminConsole.OrganizationFeatures.Organizations.Interfaces;
using Bit.Core.AdminConsole.Repositories;
using Bit.Core.AdminConsole.Services;
using Bit.Core.Billing.Constants;
using Bit.Core.Billing.Extensions;
using Bit.Core.Billing.Pricing;
using Bit.Core.OrganizationFeatures.OrganizationSponsorships.FamiliesForEnterprise.Interfaces;
using Bit.Core.Repositories;
using Bit.Core.Services;
using Quartz;
using Stripe;
using Stripe.TestHelpers;
using Event = Stripe.Event;

namespace Bit.Billing.Services.Implementations;

public class SubscriptionUpdatedHandler : ISubscriptionUpdatedHandler
{
    private readonly IStripeEventService _stripeEventService;
    private readonly IStripeEventUtilityService _stripeEventUtilityService;
    private readonly IOrganizationService _organizationService;
    private readonly IStripeFacade _stripeFacade;
    private readonly IOrganizationSponsorshipRenewCommand _organizationSponsorshipRenewCommand;
    private readonly IUserService _userService;
    private readonly IOrganizationRepository _organizationRepository;
    private readonly ISchedulerFactory _schedulerFactory;
    private readonly IOrganizationEnableCommand _organizationEnableCommand;
    private readonly IOrganizationDisableCommand _organizationDisableCommand;
    private readonly IPricingClient _pricingClient;
    private readonly IFeatureService _featureService;
    private readonly IProviderRepository _providerRepository;
    private readonly IProviderService _providerService;
    private readonly ILogger<SubscriptionUpdatedHandler> _logger;
    private readonly IPushNotificationAdapter _pushNotificationAdapter;

    public SubscriptionUpdatedHandler(
        IStripeEventService stripeEventService,
        IStripeEventUtilityService stripeEventUtilityService,
        IOrganizationService organizationService,
        IStripeFacade stripeFacade,
        IOrganizationSponsorshipRenewCommand organizationSponsorshipRenewCommand,
        IUserService userService,
        IOrganizationRepository organizationRepository,
        ISchedulerFactory schedulerFactory,
        IOrganizationEnableCommand organizationEnableCommand,
        IOrganizationDisableCommand organizationDisableCommand,
        IPricingClient pricingClient,
        IFeatureService featureService,
        IProviderRepository providerRepository,
        IProviderService providerService,
        ILogger<SubscriptionUpdatedHandler> logger,
        IPushNotificationAdapter pushNotificationAdapter)
    {
        _stripeEventService = stripeEventService;
        _stripeEventUtilityService = stripeEventUtilityService;
        _organizationService = organizationService;
        _providerService = providerService;
        _stripeFacade = stripeFacade;
        _organizationSponsorshipRenewCommand = organizationSponsorshipRenewCommand;
        _userService = userService;
        _organizationRepository = organizationRepository;
        _providerRepository = providerRepository;
        _schedulerFactory = schedulerFactory;
        _organizationEnableCommand = organizationEnableCommand;
        _organizationDisableCommand = organizationDisableCommand;
        _pricingClient = pricingClient;
        _featureService = featureService;
        _providerRepository = providerRepository;
        _providerService = providerService;
        _logger = logger;
        _pushNotificationAdapter = pushNotificationAdapter;
    }

    /// <summary>
    /// Handles the <see cref="HandledStripeWebhook.SubscriptionUpdated"/> event type from Stripe.
    /// </summary>
    /// <param name="parsedEvent"></param>
    public async Task HandleAsync(Event parsedEvent)
    {
        var subscription = await _stripeEventService.GetSubscription(parsedEvent, true, ["customer", "discounts", "latest_invoice", "test_clock"]);
        var (organizationId, userId, providerId) = _stripeEventUtilityService.GetIdsFromMetadata(subscription.Metadata);

        var currentPeriodEnd = subscription.GetCurrentPeriodEnd();

        switch (subscription.Status)
        {
            case StripeSubscriptionStatus.Unpaid or StripeSubscriptionStatus.IncompleteExpired
                when organizationId.HasValue:
                {
                    await _organizationDisableCommand.DisableAsync(organizationId.Value, currentPeriodEnd);
                    if (subscription.Status == StripeSubscriptionStatus.Unpaid &&
                        subscription.LatestInvoice is { BillingReason: "subscription_cycle" or "subscription_create" })
                    {
                        await ScheduleCancellationJobAsync(subscription.Id, organizationId.Value);
                    }
                    break;
                }
            case StripeSubscriptionStatus.Unpaid or StripeSubscriptionStatus.IncompleteExpired when providerId.HasValue:
                {
                    await HandleUnpaidProviderSubscriptionAsync(providerId.Value, parsedEvent, subscription);
                    break;
                }
            case StripeSubscriptionStatus.Unpaid or StripeSubscriptionStatus.IncompleteExpired:
                {
                    if (!userId.HasValue)
                    {
                        break;
                    }

                    if (await IsPremiumSubscriptionAsync(subscription))
                    {
                        await CancelSubscription(subscription.Id);
                        await VoidOpenInvoices(subscription.Id);
                    }

                    await _userService.DisablePremiumAsync(userId.Value, currentPeriodEnd);

                    break;
                }
            case StripeSubscriptionStatus.Incomplete when userId.HasValue:
                {
                    // Handle Incomplete subscriptions for Premium users that have open invoices from failed payments
                    // This prevents duplicate subscriptions when users retry the subscription flow
                    if (await IsPremiumSubscriptionAsync(subscription) &&
                        subscription.LatestInvoice is { Status: StripeInvoiceStatus.Open })
                    {
                        await CancelSubscription(subscription.Id);
                        await VoidOpenInvoices(subscription.Id);
                        await _userService.DisablePremiumAsync(userId.Value, currentPeriodEnd);
                    }

                    break;
                }
            case StripeSubscriptionStatus.Active when organizationId.HasValue:
                {
                    await _organizationEnableCommand.EnableAsync(organizationId.Value);
                    var organization = await _organizationRepository.GetByIdAsync(organizationId.Value);
                    if (organization != null)
                    {
                        await _pushNotificationAdapter.NotifyEnabledChangedAsync(organization);
                    }
                    break;
                }
            case StripeSubscriptionStatus.Active when providerId.HasValue:
                {
                    var provider = await _providerRepository.GetByIdAsync(providerId.Value);
                    if (provider != null)
                    {
                        provider.Enabled = true;
                        await _providerService.UpdateAsync(provider);

                        if (IsProviderSubscriptionNowActive(parsedEvent, subscription))
                        {
                            // Update the CancelAtPeriodEnd subscription option to prevent the now active provider subscription from being cancelled
                            var subscriptionUpdateOptions = new SubscriptionUpdateOptions { CancelAtPeriodEnd = false };
                            await _stripeFacade.UpdateSubscription(subscription.Id, subscriptionUpdateOptions);
                        }
                    }
                    break;
                }
            case StripeSubscriptionStatus.Active:
                {
                    if (userId.HasValue)
                    {
                        await _userService.EnablePremiumAsync(userId.Value, currentPeriodEnd);
                    }
                    break;
                }
        }

        if (organizationId.HasValue)
        {
            await _organizationService.UpdateExpirationDateAsync(organizationId.Value, currentPeriodEnd);
            if (_stripeEventUtilityService.IsSponsoredSubscription(subscription) && currentPeriodEnd.HasValue)
            {
                await _organizationSponsorshipRenewCommand.UpdateExpirationDateAsync(organizationId.Value, currentPeriodEnd.Value);
            }

            await RemovePasswordManagerCouponIfRemovingSecretsManagerTrialAsync(parsedEvent, subscription);
        }
        else if (userId.HasValue)
        {
            await _userService.UpdatePremiumExpirationAsync(userId.Value, currentPeriodEnd);
        }
    }

    private async Task CancelSubscription(string subscriptionId) =>
        await _stripeFacade.CancelSubscription(subscriptionId, new SubscriptionCancelOptions());

    private async Task VoidOpenInvoices(string subscriptionId)
    {
        var options = new InvoiceListOptions
        {
            Status = StripeInvoiceStatus.Open,
            Subscription = subscriptionId
        };
        var invoices = await _stripeFacade.ListInvoices(options);
        foreach (var invoice in invoices)
        {
            await _stripeFacade.VoidInvoice(invoice.Id);
        }
    }

    private async Task<bool> IsPremiumSubscriptionAsync(Subscription subscription)
    {
        var premiumPlans = await _pricingClient.ListPremiumPlans();
        var premiumPriceIds = premiumPlans.SelectMany(p => new[] { p.Seat.StripePriceId, p.Storage.StripePriceId }).ToHashSet();
        return subscription.Items.Any(i => premiumPriceIds.Contains(i.Price.Id));
    }

    /// <summary>
    /// Checks if the provider subscription status has changed from a non-active to an active status type
    /// If the previous status is already active(active,past-due,trialing),canceled,or null, then this will return false.
    /// </summary>
    /// <param name="parsedEvent">The event containing the previous subscription status</param>
    /// <param name="subscription">The current subscription status</param>
    /// <returns>A boolean that represents whether the event status has changed from a non-active status to an active status</returns>
    private static bool IsProviderSubscriptionNowActive(Event parsedEvent, Subscription subscription)
    {
        if (parsedEvent.Data.PreviousAttributes == null)
        {
            return false;
        }

        var previousSubscription = parsedEvent
            .Data
            .PreviousAttributes
            .ToObject<Subscription>() as Subscription;

        return previousSubscription?.Status switch
        {
            StripeSubscriptionStatus.IncompleteExpired
                or StripeSubscriptionStatus.Paused
                or StripeSubscriptionStatus.Incomplete
                or StripeSubscriptionStatus.Unpaid
                when subscription.Status == StripeSubscriptionStatus.Active => true,
            _ => false
        };
    }

    /// <summary>
    /// Removes the Password Manager coupon if the organization is removing the Secrets Manager trial.
    /// Only applies to organizations that have a subscription from the Secrets Manager trial.
    /// </summary>
    /// <param name="parsedEvent"></param>
    /// <param name="subscription"></param>
    private async Task RemovePasswordManagerCouponIfRemovingSecretsManagerTrialAsync(
        Event parsedEvent,
        Subscription subscription)
    {
        if (parsedEvent.Data.PreviousAttributes?.items is null)
        {
            return;
        }

        var organization = subscription.Metadata.TryGetValue("organizationId", out var organizationId)
            ? await _organizationRepository.GetByIdAsync(Guid.Parse(organizationId))
            : null;

        if (organization == null)
        {
            return;
        }

        var plan = await _pricingClient.GetPlanOrThrow(organization.PlanType);

        if (!plan.SupportsSecretsManager)
        {
            return;
        }

        var previousSubscription = parsedEvent.Data
            .PreviousAttributes
            .ToObject<Subscription>() as Subscription;

        // Get all plan IDs that include Secrets Manager support to check if the organization has secret manager in the
        // previous and/or current subscriptions.
        var planIdsOfPlansWithSecretManager = (await _pricingClient.ListPlans())
            .Where(orgPlan => orgPlan.SupportsSecretsManager && orgPlan.SecretsManager.StripeSeatPlanId != null)
            .Select(orgPlan => orgPlan.SecretsManager.StripeSeatPlanId)
            .ToHashSet();

        // This being false doesn't necessarily mean that the organization doesn't subscribe to Secrets Manager.
        // If there are changes to any subscription item, Stripe sends every item in the subscription, both
        // changed and unchanged.
        var previousSubscriptionHasSecretsManager =
            previousSubscription?.Items is not null &&
            previousSubscription.Items.Any(
                previousSubscriptionItem => planIdsOfPlansWithSecretManager.Contains(previousSubscriptionItem.Plan.Id));

        var currentSubscriptionHasSecretsManager =
            subscription.Items.Any(
                currentSubscriptionItem => planIdsOfPlansWithSecretManager.Contains(currentSubscriptionItem.Plan.Id));

        if (!previousSubscriptionHasSecretsManager || currentSubscriptionHasSecretsManager)
        {
            return;
        }

        var customerHasSecretsManagerTrial = subscription.Customer
            ?.Discount
            ?.Coupon
            ?.Id == "sm-standalone";

        var subscriptionHasSecretsManagerTrial = subscription.Discounts.Select(discount => discount.Coupon.Id)
            .Contains(StripeConstants.CouponIDs.SecretsManagerStandalone);

        if (customerHasSecretsManagerTrial)
        {
            await _stripeFacade.DeleteCustomerDiscount(subscription.CustomerId);
        }

        if (subscriptionHasSecretsManagerTrial)
        {
            await _stripeFacade.DeleteSubscriptionDiscount(subscription.Id);
        }
    }

    private async Task ScheduleCancellationJobAsync(string subscriptionId, Guid organizationId)
    {
        var scheduler = await _schedulerFactory.GetScheduler();

        var job = JobBuilder.Create<SubscriptionCancellationJob>()
            .WithIdentity($"cancel-sub-{subscriptionId}", "subscription-cancellations")
            .UsingJobData("subscriptionId", subscriptionId)
            .UsingJobData("organizationId", organizationId.ToString())
            .Build();

        var trigger = TriggerBuilder.Create()
            .WithIdentity($"cancel-trigger-{subscriptionId}", "subscription-cancellations")
            .StartAt(DateTimeOffset.UtcNow.AddDays(7))
            .Build();

        await scheduler.ScheduleJob(job, trigger);
    }

    private async Task HandleUnpaidProviderSubscriptionAsync(
        Guid providerId,
        Event parsedEvent,
        Subscription currentSubscription)
    {
        var provider = await _providerRepository.GetByIdAsync(providerId);
        if (provider == null)
        {
            return;
        }

        try
        {
            provider.Enabled = false;
            await _providerService.UpdateAsync(provider);

            if (parsedEvent.Data.PreviousAttributes != null)
            {
                var previousSubscription = parsedEvent.Data.PreviousAttributes.ToObject<Subscription>() as Subscription;

                if (previousSubscription is
                    {
                        Status:
                        StripeSubscriptionStatus.Trialing or
                        StripeSubscriptionStatus.Active or
                        StripeSubscriptionStatus.PastDue
                    } && currentSubscription is
                    {
                        Status: StripeSubscriptionStatus.Unpaid,
                        LatestInvoice.BillingReason: "subscription_cycle" or "subscription_create"
                    })
                {
                    if (currentSubscription.TestClock != null)
                    {
                        await WaitForTestClockToAdvanceAsync(currentSubscription.TestClock);
                    }

                    var now = currentSubscription.TestClock?.FrozenTime ?? DateTime.UtcNow;

                    var subscriptionUpdateOptions = new SubscriptionUpdateOptions { CancelAt = now.AddDays(7) };

                    await _stripeFacade.UpdateSubscription(currentSubscription.Id, subscriptionUpdateOptions);
                }
            }
        }
        catch (Exception exception)
        {
            _logger.LogError(exception, "An error occurred while trying to disable and schedule subscription cancellation for provider ({ProviderID})", providerId);
        }
    }

    private async Task WaitForTestClockToAdvanceAsync(TestClock testClock)
    {
        while (testClock.Status != "ready")
        {
            await Task.Delay(TimeSpan.FromSeconds(2));
            testClock = await _stripeFacade.GetTestClock(testClock.Id);
            if (testClock.Status == "internal_failure")
            {
                throw new Exception("Stripe Test Clock encountered an internal failure");
            }
        }
    }
}