From 0017f7a025e58313defa1061a83becf67b82194b Mon Sep 17 00:00:00 2001 From: Zixuan James Li Date: Tue, 28 Jun 2022 15:52:24 -0400 Subject: [PATCH] confirmation: Support more models as ConfirmationObjT. Signed-off-by: Zixuan James Li --- confirmation/models.py | 2 +- zerver/views/auth.py | 7 +++++-- zerver/views/realm.py | 1 + zerver/views/registration.py | 7 +++++-- zerver/views/unsubscribe.py | 1 + 5 files changed, 13 insertions(+), 5 deletions(-) diff --git a/confirmation/models.py b/confirmation/models.py index bf1492bd4a..87bf2522bb 100644 --- a/confirmation/models.py +++ b/confirmation/models.py @@ -54,7 +54,7 @@ def generate_key() -> str: return b32encode(secrets.token_bytes(15)).decode().lower() -ConfirmationObjT = Union[MultiuseInvite, PreregistrationUser, EmailChangeStatus] +ConfirmationObjT = Union[MultiuseInvite, PreregistrationUser, EmailChangeStatus, UserProfile, Realm] def get_object_from_key( diff --git a/zerver/views/auth.py b/zerver/views/auth.py index 9b341cb76b..7660497d33 100644 --- a/zerver/views/auth.py +++ b/zerver/views/auth.py @@ -190,11 +190,14 @@ def maybe_send_to_registration( if multiuse_object_key: from_multiuse_invite = True try: - multiuse_obj = get_object_from_key(multiuse_object_key, [Confirmation.MULTIUSE_INVITE]) + confirmation_obj = get_object_from_key( + multiuse_object_key, [Confirmation.MULTIUSE_INVITE] + ) except ConfirmationKeyException as exception: return render_confirmation_key_error(request, exception) - assert multiuse_obj is not None + assert isinstance(confirmation_obj, MultiuseInvite) + multiuse_obj = confirmation_obj if realm != multiuse_obj.realm: return render(request, "confirmation/link_does_not_exist.html", status=404) diff --git a/zerver/views/realm.py b/zerver/views/realm.py index bd30f8f7cd..acfee2e3e3 100644 --- a/zerver/views/realm.py +++ b/zerver/views/realm.py @@ -330,6 +330,7 @@ def realm_reactivation(request: HttpRequest, confirmation_key: str) -> HttpRespo realm = get_object_from_key(confirmation_key, [Confirmation.REALM_REACTIVATION]) except ConfirmationKeyException: return render(request, "zerver/realm_reactivation_link_error.html") + assert isinstance(realm, Realm) do_reactivate_realm(realm) context = {"realm": realm} return render(request, "zerver/realm_reactivation.html", context) diff --git a/zerver/views/registration.py b/zerver/views/registration.py index ce051bd612..4eb3dec543 100644 --- a/zerver/views/registration.py +++ b/zerver/views/registration.py @@ -137,6 +137,7 @@ def check_prereg_key(request: HttpRequest, confirmation_key: str) -> Preregistra ] prereg_user = get_object_from_key(confirmation_key, confirmation_types, activate_object=False) + assert isinstance(prereg_user, PreregistrationUser) if prereg_user.status == confirmation_settings.STATUS_REVOKED: raise ConfirmationKeyException(ConfirmationKeyException.EXPIRED) @@ -741,9 +742,11 @@ def accounts_home( def accounts_home_from_multiuse_invite(request: HttpRequest, confirmation_key: str) -> HttpResponse: realm = get_realm_from_request(request) - multiuse_object = None + multiuse_object: Optional[MultiuseInvite] = None try: - multiuse_object = get_object_from_key(confirmation_key, [Confirmation.MULTIUSE_INVITE]) + confirmation_obj = get_object_from_key(confirmation_key, [Confirmation.MULTIUSE_INVITE]) + assert isinstance(confirmation_obj, MultiuseInvite) + multiuse_object = confirmation_obj if realm != multiuse_object.realm: return render(request, "confirmation/link_does_not_exist.html", status=404) # Required for OAuth 2 diff --git a/zerver/views/unsubscribe.py b/zerver/views/unsubscribe.py index e1b0c31a76..3ede06d94c 100644 --- a/zerver/views/unsubscribe.py +++ b/zerver/views/unsubscribe.py @@ -22,6 +22,7 @@ def process_unsubscribe( except ConfirmationKeyException: return render(request, "zerver/unsubscribe_link_error.html") + assert isinstance(user_profile, UserProfile) unsubscribe_function(user_profile) context = common_context(user_profile) context.update(subscription_type=subscription_type)