app-MAIL-temp/app/account_linking.py

from abc import ABC, abstractmethod
from arrow import Arrow
from dataclasses import dataclass
from enum import Enum

from flask import url_for
from newrelic import agent
from typing import Optional

from app import config
from app.db import Session
from app.email_utils import send_email_at_most_times, send_welcome_email, render
from app.errors import AccountAlreadyLinkedToAnotherPartnerException
from app.log import LOG
from app.models import (
    PartnerSubscription,
    Partner,
    PartnerUser,
    User,
    AppleSubscription,
    Subscription,
)
from app.proton.utils import is_proton_partner
from app.utils import random_string


class SLPlanType(Enum):
    Free = 1
    Premium = 2


@dataclass
class SLPlan:
    type: SLPlanType
    expiration: Optional[Arrow]


@dataclass
class PartnerLinkRequest:
    name: str
    email: str
    external_user_id: str
    plan: SLPlan
    from_partner: bool


@dataclass
class LinkResult:
    user: User
    strategy: str


def send_double_subscription_if_needed(partner_user: PartnerUser):
    sub = partner_user.user.get_active_subscription()
    if isinstance(sub, AppleSubscription):
        channel = "Apple"
    elif isinstance(sub, Subscription):
        channel = "Paddle"
    else:
        return
    send_email_at_most_times(
        partner_user.user,
        config.ALERT_DUAL_SUBSCRIPTION_WITH_PARTNER,
        partner_user.user.email,
        f"You have two subscriptions in SimpleLogin",
        render(
            "transactional/double-subscription-partner.txt.jinja2",
            is_proton=is_proton_partner(partner_user.partner),
            partner=partner_user.partner,
            subscription_channel=channel,
            cancel_link=url_for("dashboard.billing"),
        ),
        render(
            "transactional/double-subscription-partner.html",
            is_proton=is_proton_partner(partner_user.partner),
            partner=partner_user.partner,
            subscription_channel=channel,
            cancel_link=url_for("dashboard.billing"),
        ),
    )


def set_plan_for_partner_user(partner_user: PartnerUser, plan: SLPlan):
    sub = PartnerSubscription.get_by(partner_user_id=partner_user.id)
    if plan.type == SLPlanType.Free:
        if sub is not None:
            LOG.i(
                f"Deleting partner_subscription [user_id={partner_user.user_id}] [partner_id={partner_user.partner_id}]"
            )
            PartnerSubscription.delete(sub.id)
            agent.record_custom_event("PlanChange", {"plan": "free"})
    else:
        if sub is None:
            LOG.i(
                f"Creating partner_subscription [user_id={partner_user.user_id}] [partner_id={partner_user.partner_id}]"
            )
            PartnerSubscription.create(
                partner_user_id=partner_user.id,
                end_at=plan.expiration,
            )
            agent.record_custom_event("PlanChange", {"plan": "premium", "type": "new"})
        else:
            if sub.end_at != plan.expiration:
                LOG.i(
                    f"Updating partner_subscription [user_id={partner_user.user_id}] [partner_id={partner_user.partner_id}]"
                )
                agent.record_custom_event(
                    "PlanChange", {"plan": "premium", "type": "extension"}
                )
                sub.end_at = plan.expiration
        send_double_subscription_if_needed(partner_user)
    Session.commit()


def set_plan_for_user(user: User, plan: SLPlan, partner: Partner):
    partner_user = PartnerUser.get_by(partner_id=partner.id, user_id=user.id)
    if partner_user is None:
        return
    return set_plan_for_partner_user(partner_user, plan)


def ensure_partner_user_exists_for_user(
    link_request: PartnerLinkRequest, sl_user: User, partner: Partner
) -> PartnerUser:
    # Find partner_user by user_id
    res = PartnerUser.get_by(user_id=sl_user.id)
    if res and res.partner_id != partner.id:
        raise AccountAlreadyLinkedToAnotherPartnerException()
    if not res:
        res = PartnerUser.create(
            user_id=sl_user.id,
            partner_id=partner.id,
            partner_email=link_request.email,
            external_user_id=link_request.external_user_id,
        )
        Session.commit()
        LOG.i(
            f"Created new partner_user for partner:{partner.id} user:{sl_user.id} external_user_id:{link_request.external_user_id}. PartnerUser.id is {res.id}"
        )
    return res


class ClientMergeStrategy(ABC):
    def __init__(
        self,
        link_request: PartnerLinkRequest,
        user: Optional[User],
        partner: Partner,
    ):
        if self.__class__ == ClientMergeStrategy:
            raise RuntimeError("Cannot directly instantiate a ClientMergeStrategy")
        self.link_request = link_request
        self.user = user
        self.partner = partner

    @abstractmethod
    def process(self) -> LinkResult:
        pass


class NewUserStrategy(ClientMergeStrategy):
    def process(self) -> LinkResult:
        # Will create a new SL User with a random password
        new_user = User.create(
            email=self.link_request.email,
            name=self.link_request.name,
            password=random_string(20),
            from_partner=self.link_request.from_partner,
        )
        partner_user = PartnerUser.create(
            user_id=new_user.id,
            partner_id=self.partner.id,
            external_user_id=self.link_request.external_user_id,
            partner_email=self.link_request.email,
        )
        LOG.i(
            f"Created new user for login request for partner:{self.partner.id} external_user_id:{self.link_request.external_user_id}. New user {new_user.id} partner_user:{partner_user.id}"
        )
        set_plan_for_partner_user(
            partner_user,
            self.link_request.plan,
        )
        Session.commit()

        if not new_user.created_by_partner:
            send_welcome_email(new_user)

        agent.record_custom_event("PartnerUserCreation", {"partner": self.partner.name})

        return LinkResult(
            user=new_user,
            strategy=self.__class__.__name__,
        )


class ExistingUnlinkedUserStrategy(ClientMergeStrategy):
    def process(self) -> LinkResult:

        partner_user = ensure_partner_user_exists_for_user(
            self.link_request, self.user, self.partner
        )
        set_plan_for_partner_user(partner_user, self.link_request.plan)

        return LinkResult(
            user=self.user,
            strategy=self.__class__.__name__,
        )


class LinkedWithAnotherPartnerUserStrategy(ClientMergeStrategy):
    def process(self) -> LinkResult:
        raise AccountAlreadyLinkedToAnotherPartnerException()


def get_login_strategy(
    link_request: PartnerLinkRequest, user: Optional[User], partner: Partner
) -> ClientMergeStrategy:
    if user is None:
        # We couldn't find any SimpleLogin user with the requested e-mail
        return NewUserStrategy(link_request, user, partner)
    # Check if user is already linked with another partner_user
    other_partner_user = PartnerUser.get_by(partner_id=partner.id, user_id=user.id)
    if other_partner_user is not None:
        return LinkedWithAnotherPartnerUserStrategy(link_request, user, partner)
    # There is a SimpleLogin user with the partner_user's e-mail
    return ExistingUnlinkedUserStrategy(link_request, user, partner)


def process_login_case(
    link_request: PartnerLinkRequest, partner: Partner
) -> LinkResult:
    # Try to find a SimpleLogin user registered with that partner user id
    partner_user = PartnerUser.get_by(
        partner_id=partner.id, external_user_id=link_request.external_user_id
    )
    if partner_user is None:
        # We didn't find any SimpleLogin user registered with that partner user id
        # Try to find it using the partner's e-mail address
        user = User.get_by(email=link_request.email)
        return get_login_strategy(link_request, user, partner).process()
    else:
        # We found the SL user registered with that partner user id
        # We're done
        set_plan_for_partner_user(partner_user, link_request.plan)
        # It's the same user. No need to do anything
        return LinkResult(
            user=partner_user.user,
            strategy="Link",
        )


def link_user(
    link_request: PartnerLinkRequest, current_user: User, partner: Partner
) -> LinkResult:
    partner_user = ensure_partner_user_exists_for_user(
        link_request, current_user, partner
    )
    set_plan_for_partner_user(partner_user, link_request.plan)

    agent.record_custom_event("AccountLinked", {"partner": partner.name})
    Session.commit()
    return LinkResult(
        user=current_user,
        strategy="Link",
    )


def switch_already_linked_user(
    link_request: PartnerLinkRequest, partner_user: PartnerUser, current_user: User
):
    # Find if the user has another link and unlink it
    other_partner_user = PartnerUser.get_by(
        user_id=current_user.id,
        partner_id=partner_user.partner_id,
    )
    if other_partner_user is not None:
        LOG.i(
            f"Deleting previous partner_user:{other_partner_user.id} from user:{current_user.id}"
        )
        PartnerUser.delete(other_partner_user.id)
    LOG.i(f"Linking partner_user:{partner_user.id} to user:{current_user.id}")
    # Link this partner_user to the current user
    partner_user.user_id = current_user.id
    # Set plan
    set_plan_for_partner_user(partner_user, link_request.plan)
    Session.commit()
    return LinkResult(
        user=current_user,
        strategy="Link",
    )


def process_link_case(
    link_request: PartnerLinkRequest,
    current_user: User,
    partner: Partner,
) -> LinkResult:
    # Try to find a SimpleLogin user linked with this Partner account
    partner_user = PartnerUser.get_by(
        partner_id=partner.id, external_user_id=link_request.external_user_id
    )
    if partner_user is None:
        # There is no SL user linked with the partner. Proceed with linking
        return link_user(link_request, current_user, partner)

    # There is a SL user registered with the partner. Check if is the current one
    if partner_user.id == current_user.id:
        # Update plan
        set_plan_for_partner_user(partner_user, link_request.plan)
        # It's the same user. No need to do anything
        return LinkResult(
            user=current_user,
            strategy="Link",
        )
    else:

        return switch_already_linked_user(link_request, partner_user, current_user)