app-MAIL-temp/app/proton/proton_callback_handler.py
Adrià Casajús 025d4feba0
Sync on partner user creation + several fixes (#2214)
* Do not close session since it leads to orphan user object

* Redirect instead of render to avoid having to have a mailbox object

* On inital partner link/login trigger sync

* Update github action upload/artifact to v4

* Remove sys.exit used to test script locally

* Simplified script to update alias flags and note
2024-09-12 16:58:20 +02:00

122 lines
3.8 KiB
Python

from dataclasses import dataclass
from enum import Enum
from flask import url_for
from typing import Optional
import arrow
from app import config
from app.errors import LinkException
from app.models import User, Partner, Job
from app.proton.proton_client import ProtonClient, ProtonUser
from app.account_linking import (
process_login_case,
process_link_case,
PartnerLinkRequest,
)
class Action(Enum):
Login = 1
Link = 2
@dataclass
class ProtonCallbackResult:
redirect_to_login: bool
flash_message: Optional[str]
flash_category: Optional[str]
redirect: Optional[str]
user: Optional[User]
def generate_account_not_allowed_to_log_in() -> ProtonCallbackResult:
return ProtonCallbackResult(
redirect_to_login=True,
flash_message="This account is not allowed to log in with Proton. Please convert your account to a full Proton account",
flash_category="error",
redirect=None,
user=None,
)
class ProtonCallbackHandler:
def __init__(self, proton_client: ProtonClient):
self.proton_client = proton_client
def _initial_alias_sync(self, user: User):
Job.create(
name=config.JOB_SEND_ALIAS_CREATION_EVENTS,
payload={"user_id": user.id},
run_at=arrow.now(),
commit=True,
)
def handle_login(self, partner: Partner) -> ProtonCallbackResult:
try:
user = self.__get_partner_user()
if user is None:
return generate_account_not_allowed_to_log_in()
res = process_login_case(user, partner)
self._initial_alias_sync(res.user)
return ProtonCallbackResult(
redirect_to_login=False,
flash_message=None,
flash_category=None,
redirect=None,
user=res.user,
)
except LinkException as e:
return ProtonCallbackResult(
redirect_to_login=True,
flash_message=e.message,
flash_category="error",
redirect=None,
user=None,
)
def handle_link(
self,
current_user: Optional[User],
partner: Partner,
) -> ProtonCallbackResult:
if current_user is None:
raise Exception("Cannot link account with current_user being None")
try:
user = self.__get_partner_user()
if user is None:
return generate_account_not_allowed_to_log_in()
res = process_link_case(user, current_user, partner)
self._initial_alias_sync(res.user)
return ProtonCallbackResult(
redirect_to_login=False,
flash_message="Account successfully linked",
flash_category="success",
redirect=url_for("dashboard.setting"),
user=res.user,
)
except LinkException as e:
return ProtonCallbackResult(
redirect_to_login=False,
flash_message=e.message,
flash_category="error",
redirect=None,
user=None,
)
def __get_partner_user(self) -> Optional[PartnerLinkRequest]:
proton_user = self.__get_proton_user()
if proton_user is None:
return None
return PartnerLinkRequest(
email=proton_user.email,
external_user_id=proton_user.id,
name=proton_user.name,
plan=proton_user.plan,
from_partner=False, # The user has started this flow, so we don't mark it as created by a partner
)
def __get_proton_user(self) -> Optional[ProtonUser]:
user = self.proton_client.get_user()
if user is None:
return None
return ProtonUser(email=user.email, plan=user.plan, name=user.name, id=user.id)