This commit is contained in:
MrMeeb 2024-10-24 12:00:05 +01:00
parent bc48198bb1
commit 44fda2d94e
12 changed files with 180 additions and 59 deletions

@ -112,6 +112,7 @@ def ensure_partner_user_exists_for_user(
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}"

@ -16,6 +16,8 @@ from flask_admin.contrib import sqla
from flask_login import current_user
from app.db import Session
from app.events.event_dispatcher import EventDispatcher
from app.events.generated.event_pb2 import EventContent, UserPlanChanged
from app.models import (
User,
ManualSubscription,
@ -39,6 +41,7 @@ from app.models import (
UserAuditLog,
)
from app.newsletter_utils import send_newsletter_to_user, send_newsletter_to_address
from app.user_audit_log_utils import emit_user_audit_log, UserAuditLogAction
def _admin_action_formatter(view, context, model, name):
@ -351,17 +354,42 @@ def manual_upgrade(way: str, ids: [int], is_giveaway: bool):
manual_sub.end_at = manual_sub.end_at.shift(years=1)
else:
manual_sub.end_at = arrow.now().shift(years=1, days=1)
emit_user_audit_log(
user=user,
action=UserAuditLogAction.Upgrade,
message=f"Admin {current_user.email} extended manual subscription to user {user.email}",
)
EventDispatcher.send_event(
user=user,
content=EventContent(
user_plan_change=UserPlanChanged(
plan_end_time=manual_sub.end_at.timestamp
)
),
)
flash(f"Subscription extended to {manual_sub.end_at.humanize()}", "success")
continue
else:
emit_user_audit_log(
user=user,
action=UserAuditLogAction.Upgrade,
message=f"Admin {current_user.email} created manual subscription to user {user.email}",
)
manual_sub = ManualSubscription.create(
user_id=user.id,
end_at=arrow.now().shift(years=1, days=1),
comment=way,
is_giveaway=is_giveaway,
)
EventDispatcher.send_event(
user=user,
content=EventContent(
user_plan_change=UserPlanChanged(
plan_end_time=manual_sub.end_at.timestamp
)
),
)
ManualSubscription.create(
user_id=user.id,
end_at=arrow.now().shift(years=1, days=1),
comment=way,
is_giveaway=is_giveaway,
)
flash(f"New {way} manual subscription for {user} is created", "success")
flash(f"New {way} manual subscription for {user} is created", "success")
Session.commit()
@ -453,14 +481,7 @@ class ManualSubscriptionAdmin(SLModelView):
"Extend 1 year more?",
)
def extend_1y(self, ids):
for ms in ManualSubscription.filter(ManualSubscription.id.in_(ids)):
ms.end_at = ms.end_at.shift(years=1)
flash(f"Extend subscription for 1 year for {ms.user}", "success")
AdminAuditLog.extend_subscription(
current_user.id, ms.user.id, ms.end_at, "1 year"
)
Session.commit()
self.__extend_manual_subscription(ids, msg="1 year", years=1)
@action(
"extend_1m",
@ -468,11 +489,26 @@ class ManualSubscriptionAdmin(SLModelView):
"Extend 1 month more?",
)
def extend_1m(self, ids):
self.__extend_manual_subscription(ids, msg="1 month", months=1)
def __extend_manual_subscription(self, ids: List[int], msg: str, **kwargs):
for ms in ManualSubscription.filter(ManualSubscription.id.in_(ids)):
ms.end_at = ms.end_at.shift(months=1)
flash(f"Extend subscription for 1 month for {ms.user}", "success")
sub: ManualSubscription = ms
sub.end_at = sub.end_at.shift(**kwargs)
flash(f"Extend subscription for {msg} for {sub.user}", "success")
emit_user_audit_log(
user=sub.user,
action=UserAuditLogAction.Upgrade,
message=f"Admin {current_user.email} extended manual subscription for {msg} for {sub.user}",
)
AdminAuditLog.extend_subscription(
current_user.id, ms.user.id, ms.end_at, "1 month"
current_user.id, sub.user.id, sub.end_at, msg
)
EventDispatcher.send_event(
user=sub.user,
content=EventContent(
user_plan_change=UserPlanChanged(plan_end_time=sub.end_at.timestamp)
),
)
Session.commit()

@ -38,7 +38,11 @@ def create_mailbox():
the new mailbox dict
"""
user = g.user
mailbox_email = sanitize_email(request.get_json().get("email"))
email = request.get_json().get("email")
if not email:
return jsonify(error="Invalid email"), 400
mailbox_email = sanitize_email(email)
try:
new_mailbox = mailbox_utils.create_mailbox(user, mailbox_email).mailbox

@ -121,10 +121,16 @@ def mailbox_route():
@login_required
def mailbox_verify():
mailbox_id = request.args.get("mailbox_id")
if not mailbox_id:
LOG.i("Missing mailbox_id")
flash("You followed an invalid link", "error")
return redirect(url_for("dashboard.mailbox_route"))
code = request.args.get("code")
if not code:
# Old way
return verify_with_signed_secret(mailbox_id)
try:
mailbox = mailbox_utils.verify_mailbox_code(current_user, mailbox_id, code)
except mailbox_utils.MailboxError as e:

@ -171,17 +171,17 @@ def verify_mailbox_code(user: User, mailbox_id: int, code: str) -> Mailbox:
f"User {user} failed to verify mailbox {mailbox_id} because it does not exist"
)
raise MailboxError("Invalid mailbox")
if mailbox.user_id != user.id:
LOG.i(
f"User {user} failed to verify mailbox {mailbox_id} because it's owned by another user"
)
raise MailboxError("Invalid mailbox")
if mailbox.verified:
LOG.i(
f"User {user} failed to verify mailbox {mailbox_id} because it's already verified"
)
clear_activation_codes_for_mailbox(mailbox)
return mailbox
if mailbox.user_id != user.id:
LOG.i(
f"User {user} failed to verify mailbox {mailbox_id} because it's owned by another user"
)
raise MailboxError("Invalid mailbox")
activation = (
MailboxActivation.filter(MailboxActivation.mailbox_id == mailbox_id)

@ -24,6 +24,7 @@ from sqlalchemy import text, desc, CheckConstraint, Index, Column
from sqlalchemy.dialects.postgresql import TSVECTOR
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import deferred
from sqlalchemy.orm.exc import ObjectDeletedError
from sqlalchemy.sql import and_
from sqlalchemy_utils import ArrowType
@ -3781,15 +3782,18 @@ class SyncEvent(Base, ModelMixin):
)
def mark_as_taken(self, allow_taken_older_than: Optional[Arrow] = None) -> bool:
taken_condition = ["taken_time IS NULL"]
args = {"taken_time": arrow.now().datetime, "sync_event_id": self.id}
if allow_taken_older_than:
taken_condition.append("taken_time < :taken_older_than")
args["taken_older_than"] = allow_taken_older_than.datetime
sql_taken_condition = "({})".format(" OR ".join(taken_condition))
sql = f"UPDATE sync_event SET taken_time = :taken_time WHERE id = :sync_event_id AND {sql_taken_condition}"
res = Session.execute(sql, args)
Session.commit()
try:
taken_condition = ["taken_time IS NULL"]
args = {"taken_time": arrow.now().datetime, "sync_event_id": self.id}
if allow_taken_older_than:
taken_condition.append("taken_time < :taken_older_than")
args["taken_older_than"] = allow_taken_older_than.datetime
sql_taken_condition = "({})".format(" OR ".join(taken_condition))
sql = f"UPDATE sync_event SET taken_time = :taken_time WHERE id = :sync_event_id AND {sql_taken_condition}"
res = Session.execute(sql, args)
Session.commit()
except ObjectDeletedError:
return False
return res.rowcount > 0

@ -1,8 +1,10 @@
from typing import Optional
import arrow
from arrow import Arrow
from app.models import PartnerUser, PartnerSubscription, User
from app import config
from app.models import PartnerUser, PartnerSubscription, User, Job
from app.user_audit_log_utils import emit_user_audit_log, UserAuditLogAction
@ -15,6 +17,11 @@ def create_partner_user(
partner_email=partner_email,
external_user_id=external_user_id,
)
Job.create(
name=config.JOB_SEND_ALIAS_CREATION_EVENTS,
payload={"user_id": user.id},
run_at=arrow.now(),
)
emit_user_audit_log(
user=user,
action=UserAuditLogAction.LinkAccount,

@ -2,11 +2,9 @@ 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.models import User, Partner
from app.proton.proton_client import ProtonClient, ProtonUser
from app.account_linking import (
process_login_case,
@ -43,21 +41,12 @@ 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,
@ -86,7 +75,6 @@ class ProtonCallbackHandler:
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",

@ -8,6 +8,7 @@
<tr>
<th scope="col">User ID</th>
<th scope="col">Email</th>
<th scope="col">Verified</th>
<th scope="col">Status</th>
<th scope="col">Paid</th>
<th>Subscription</th>
@ -20,8 +21,12 @@
<tr>
<td>{{ user.id }}</td>
<td><a href="?email={{ user.email }}">{{ user.email }}</a></td>
{% if user.activated %}
<td class="text-success">Activated</td>
{% else %}
<td class="text-warning">Pending</td>
{% endif %}
{% if user.disabled %}
<td class="text-danger">Disabled</td>
{% else %}
<td class="text-success">Enabled</td>

@ -5,7 +5,7 @@ from app.models import Mailbox
from tests.utils import login
def test_create_mailbox(flask_client):
def test_create_mailbox_valid(flask_client):
login(flask_client)
r = flask_client.post(
@ -21,10 +21,34 @@ def test_create_mailbox(flask_client):
assert r.json["default"] is False
assert r.json["nb_alias"] == 0
# invalid email address
def test_create_mailbox_invalid_email(flask_client):
login(flask_client)
r = flask_client.post(
"/api/mailboxes",
json={"email": "gmail.com"},
json={"email": "gmail.com"}, # not an email address
)
assert r.status_code == 400
assert r.json == {"error": "Invalid email"}
def test_create_mailbox_empty_payload(flask_client):
login(flask_client)
r = flask_client.post(
"/api/mailboxes",
json={},
)
assert r.status_code == 400
assert r.json == {"error": "Invalid email"}
def test_create_mailbox_empty_email(flask_client):
login(flask_client)
r = flask_client.post(
"/api/mailboxes",
json={"email": ""},
)
assert r.status_code == 400

@ -25,15 +25,17 @@ class MockProtonClient(ProtonClient):
return self.user
def check_initial_sync_job(user: User):
def check_initial_sync_job(user: User, expected: bool):
found = False
for job in Job.yield_per_query(10).filter_by(
name=config.JOB_SEND_ALIAS_CREATION_EVENTS,
state=JobState.ready.value,
):
if job.payload.get("user_id") == user.id:
found = True
Job.delete(job.id)
return
assert False
break
assert expected == found
def test_proton_callback_handler_unexistant_sl_user():
@ -69,10 +71,9 @@ def test_proton_callback_handler_unexistant_sl_user():
)
assert partner_user is not None
assert partner_user.external_user_id == external_id
check_initial_sync_job(res.user)
def test_proton_callback_handler_existant_sl_user():
def test_proton_callback_handler_existing_sl_user():
email = random_email()
sl_user = User.create(email, commit=True)
@ -98,7 +99,43 @@ def test_proton_callback_handler_existant_sl_user():
sa = PartnerUser.get_by(user_id=sl_user.id, partner_id=get_proton_partner().id)
assert sa is not None
assert sa.partner_email == user.email
check_initial_sync_job(res.user)
check_initial_sync_job(res.user, True)
def test_proton_callback_handler_linked_sl_user():
email = random_email()
external_id = random_string()
sl_user = User.create(email, commit=True)
PartnerUser.create(
user_id=sl_user.id,
partner_id=get_proton_partner().id,
external_user_id=external_id,
partner_email=email,
commit=True,
)
user = UserInformation(
email=email,
name=random_string(),
id=external_id,
plan=SLPlan(type=SLPlanType.Premium, expiration=Arrow.utcnow().shift(hours=2)),
)
handler = ProtonCallbackHandler(MockProtonClient(user=user))
res = handler.handle_login(get_proton_partner())
assert res.user is not None
assert res.user.id == sl_user.id
# Ensure the user is not marked as created from partner
assert User.FLAG_CREATED_FROM_PARTNER != (
res.user.flags & User.FLAG_CREATED_FROM_PARTNER
)
assert res.user.notification is True
assert res.user.trial_end is not None
sa = PartnerUser.get_by(user_id=sl_user.id, partner_id=get_proton_partner().id)
assert sa is not None
assert sa.partner_email == user.email
check_initial_sync_job(res.user, False)
def test_proton_callback_handler_none_user_login():

@ -286,6 +286,15 @@ def test_verify_other_users_mailbox():
mailbox_utils.verify_mailbox_code(user, mailbox.id, "9999999")
def test_verify_other_users_already_verified_mailbox():
other = create_new_user()
mailbox = Mailbox.create(
user_id=other.id, email=random_email(), verified=True, commit=True
)
with pytest.raises(mailbox_utils.MailboxError):
mailbox_utils.verify_mailbox_code(user, mailbox.id, "9999999")
@mail_sender.store_emails_test_decorator
def test_verify_fail():
output = mailbox_utils.create_mailbox(user, random_email())