Add different magic token intents for email address verification

- Use intent: register when testing magic link endpoints

- Update magic link tests

- Allow any email to make an account in testing mode

Author: ml-evs

pydatalab/src/pydatalab/routes/v0_1/auth.py

@@ -14,11 +14,11 @@
 
 import jwt
 from bson import ObjectId
-from flask import Blueprint, Response, g, jsonify, redirect, request
+from flask import Blueprint, g, jsonify, redirect, request
 from flask_dance.consumer import OAuth2ConsumerBlueprint, oauth_authorized
 from flask_login import current_user, login_user
 from flask_login.utils import LocalProxy
-from werkzeug.exceptions import BadRequest
+from werkzeug.exceptions import BadRequest, Forbidden
 
 from pydatalab.config import CONFIG
 from pydatalab.errors import UserRegistrationForbidden
@@ -253,6 +253,9 @@ def _check_email_domain(email: str, allow_list: list[str] | None) -> bool:
         Whether the email address is allowed to register an account.
 
     """
+    if CONFIG.TESTING:
+        return True
+
     domain = email.split("@")[-1]
     if isinstance(allow_list, list) and not allow_list:
         return False
@@ -420,48 +423,36 @@ def attach_identity_to_user(
         wrapped_login_user(get_by_id(str(user.immutable_id)))
 
 
-def _validate_magic_link_request(email: str, referrer: str) -> tuple[Response | None, int]:
+def _validate_magic_link_request(email: str, referrer: str) -> None:
     if not email:
-        return jsonify({"status": "error", "detail": "No email provided."}), 400
+        raise BadRequest("No email provided")
 
     if not re.match(r"^\S+@\S+.\S+$", email):
-        return jsonify({"status": "error", "detail": "Invalid email provided."}), 400
+        raise BadRequest("Invalid email provided.")
 
     if not referrer:
-        LOGGER.warning("No referrer provided for magic link request")
-        return (
-            jsonify(
-                {
-                    "status": "error",
-                    "detail": "Referrer address not provided, please contact the datalab administrator.",
-                }
-            ),
-            400,
-        )
-
-    return None, 200
+        raise BadRequest("Referrer address not provided, please contact the datalab administrator")
 
 
-def _generate_and_store_token(email: str, is_test: bool = False) -> str:
+def _generate_and_store_token(email: str, intent: str = "register") -> str:
     """Generate a JWT for the user with a short expiration and store it in the session.
 
     The session itself persists beyond the JWT expiration. The `exp` key is a standard
     part of JWT that PyJWT treats as an expiration time and will correctly encode the datetime.
 
     Args:
         email: The user's email address to include in the token.
-        is_test: If True, generates a token for testing purposes that may have different
-                 expiration or validation rules. Defaults to False.
+        intent: The intent of the magic link, e.g., "register" "verify", or "login".
 
     Returns:
         The generated JWT token string.
+
     """
     payload = {
         "exp": datetime.datetime.now(datetime.timezone.utc) + LINK_EXPIRATION,
         "email": email,
+        "intent": intent,
     }
-    if is_test:
-        payload["is_test"] = True
 
     token = jwt.encode(
         payload,
@@ -474,45 +465,49 @@ def _generate_and_store_token(email: str, is_test: bool = False) -> str:
     return token
 
 
-def _check_user_registration_allowed(email: str) -> tuple[Response | None, int]:
+def _check_user_registration_allowed(email: str) -> None:
     user = find_user_with_identity(email, IdentityType.EMAIL, verify=False)
 
     if not user:
         allowed = _check_email_domain(email, CONFIG.EMAIL_DOMAIN_ALLOW_LIST)
         if not allowed:
-            LOGGER.info("Did not allow %s to register an account", email)
-            return (
-                jsonify(
-                    {
-                        "status": "error",
-                        "detail": f"Email address {email} is not allowed to register an account. Please contact the administrator if you believe this is an error.",
-                    }
-                ),
-                403,
+            LOGGER.warning("Did not allow %s to register an account", email)
+            raise Forbidden(
+                f"Email address {email} is not allowed to register an account. Please contact the administrator if you believe this is an error."
             )
 
-    return None, 200
 
+def _send_magic_link_email(
+    email: str, token: str, referrer: str | None, purpose: str = "authorize"
+) -> None:
+    if not referrer:
+        referrer = "https://example.com"
 
-def _send_magic_link_email(email: str, token: str, referrer: str) -> tuple[Response | None, int]:
     link = f"{referrer}?token={token}"
     instance_url = referrer.replace("https://", "")
-    user = find_user_with_identity(email, IdentityType.EMAIL, verify=False)
 
-    if user is not None:
-        subject = "Datalab Sign-in Magic Link"
-        body = f"Click the link below to sign-in to the datalab instance at {instance_url}:\n\n{link}\n\nThis link is single-use and will expire in 1 hour."
+    if purpose == "authorize":
+        user = find_user_with_identity(email, IdentityType.EMAIL, verify=False)
+        if user is not None:
+            subject = "datalab Sign-in Magic Link"
+            body = f"Click the link below to sign-in to the datalab instance at {instance_url}:\n\n{link}\n\nThis link is single-use and will expire in 1 hour."
+        else:
+            subject = "datalab Registration Magic Link"
+            body = f"Click the link below to register for the datalab instance at {instance_url}:\n\n{link}\n\nThis link is single-use and will expire in 1 hour."
+
+    elif purpose == "verify":
+        subject = "datalab Email Address Verification"
+        body = f"Click the link below to verify your email address for the datalab instance at {instance_url}:\n\n{link}\n\nThis link is single-use and will expire in 1 hour."
+
     else:
-        subject = "Datalab Registration Magic Link"
-        body = f"Click the link below to register for the datalab instance at {instance_url}:\n\n{link}\n\nThis link is single-use and will expire in 1 hour."
+        LOGGER.critical("Unknown purpose %s for magic link email", purpose)
+        raise RuntimeError("Unknown error occurred")
 
     try:
         send_mail(email, subject, body)
     except Exception as exc:
         LOGGER.warning("Failed to send email to %s: %s", email, exc)
-        return jsonify({"status": "error", "detail": "Email not sent successfully."}), 400
-
-    return None, 200
+        raise RuntimeError("Email not sent successfully.")
 
 
 @EMAIL_BLUEPRINT.route("/magic-link", methods=["POST"])
@@ -526,21 +521,12 @@ def generate_and_share_magic_link():
     email = request_json.get("email")
     referrer = request_json.get("referrer")
 
-    error_response, status_code = _validate_magic_link_request(email, referrer)
-    if error_response:
-        return error_response, status_code
-
-    error_response, status_code = _check_user_registration_allowed(email)
-    if error_response:
-        return error_response, status_code
+    _validate_magic_link_request(email, referrer)
+    _check_user_registration_allowed(email)
+    token = _generate_and_store_token(email, intent="register")
+    _send_magic_link_email(email, token, referrer)
 
-    token = _generate_and_store_token(email)
-
-    error_response, status_code = _send_magic_link_email(email, token, referrer)
-    if error_response:
-        return error_response, status_code
-
-    return jsonify({"status": "success", "detail": "Email sent successfully."}), 200
+    return jsonify({"status": "success", "message": "Email sent successfully."}), 200
 
 
 @EMAIL_BLUEPRINT.route("/email")
@@ -580,21 +566,22 @@ def email_logged_in():
     # If the email domain list is explicitly configured to None, this allows any
     # email address to make an active account, otherwise the email domain must match
     # the list of allowed domains and the admin must verify the user
-    is_test = data.get("is_test", False)
 
-    if not is_test:
+    if data.get("intent") == "register":
         allowed = _check_email_domain(email, CONFIG.EMAIL_DOMAIN_ALLOW_LIST)
         if not allowed:
             raise UserRegistrationForbidden
 
-    create_account = AccountStatus.UNVERIFIED
-    if (
-        CONFIG.EMAIL_DOMAIN_ALLOW_LIST is None
-        or CONFIG.EMAIL_AUTO_ACTIVATE_ACCOUNTS
-        or CONFIG.AUTO_ACTIVATE_ACCOUNTS
-        or is_test
-    ):
-        create_account = AccountStatus.ACTIVE
+        create_account = AccountStatus.UNVERIFIED
+        if (
+            CONFIG.EMAIL_DOMAIN_ALLOW_LIST is None
+            or CONFIG.EMAIL_AUTO_ACTIVATE_ACCOUNTS
+            or CONFIG.AUTO_ACTIVATE_ACCOUNTS
+        ):
+            create_account = AccountStatus.ACTIVE
+
+    else:
+        create_account = False
 
     find_create_or_modify_user(
         email,
@@ -751,10 +738,8 @@ def create_test_magic_link():
     email = request_json.get("email")
     referrer = request_json.get("referrer", "http://localhost:8080")
 
-    error_response, status_code = _validate_magic_link_request(email, referrer)
-    if error_response:
-        return error_response, status_code
+    _validate_magic_link_request(email, referrer)
 
-    token = _generate_and_store_token(email, is_test=True)
+    token = _generate_and_store_token(email, intent="register")
 
     return jsonify({"status": "success", "token": token}), 200

pydatalab/src/pydatalab/routes/v0_1/users.py

@@ -1,11 +1,14 @@
 from bson import ObjectId
 from flask import Blueprint, jsonify, request
 from flask_login import current_user
+from werkzeug.exceptions import BadRequest, Forbidden, Unauthorized
 
 from pydatalab.config import CONFIG
-from pydatalab.models.people import DisplayName, EmailStr
+from pydatalab.logger import LOGGER
+from pydatalab.models.people import AccountStatus, DisplayName, EmailStr
 from pydatalab.mongo import flask_mongo
 from pydatalab.permissions import active_users_or_get_only
+from pydatalab.routes.v0_1.auth import _generate_and_store_token, _send_magic_link_email
 
 USERS = Blueprint("users", __name__)
 
@@ -29,51 +32,98 @@ def save_user(user_id):
         account_status = request_json.get("account_status", None)
 
     if not current_user.is_authenticated and not CONFIG.TESTING:
-        return (jsonify({"status": "error", "message": "No user authenticated."}), 401)
+        raise Unauthorized("No user authenticated.")
 
     if not CONFIG.TESTING and current_user.id != user_id and current_user.role != "admin":
-        return (
-            jsonify({"status": "error", "message": "User not allowed to edit this profile."}),
-            403,
-        )
+        raise Forbidden("Current user not allowed to edit this profile.")
 
     update = {}
 
     try:
         if display_name:
             update["display_name"] = DisplayName(display_name)
 
+    except ValueError:
+        raise BadRequest(f"Invalid display name {display_name!r} was passed")
+
+    try:
         if contact_email or contact_email in (None, ""):
             if contact_email in ("", None):
                 update["contact_email"] = None
             else:
                 update["contact_email"] = EmailStr(contact_email)
 
-        if account_status:
-            update["account_status"] = account_status
+    except ValueError:
+        raise BadRequest(f"Invalid email address {contact_email!r} was passed")
+
+    trigger_email_verification = False
+    if update.get("contact_email"):
+        # Check if this email identity already exists for this user
+        existing_email_identity = False
+        if update.get("contact_email") is not None:
+            existing_email_identity = flask_mongo.db.users.find_one(
+                {
+                    "_id": ObjectId(user_id),
+                    "identities": {"$elemMatch": {"type": "email", "identifier": contact_email}},
+                }
+            )
+            if not existing_email_identity:
+                # If not, push it as a new unverified identity
+                flask_mongo.db.users.update_one(
+                    {"_id": ObjectId(user_id)},
+                    {
+                        "$push": {
+                            "identities": {
+                                "identity_type": "email",
+                                "identifier": contact_email,
+                                "name": contact_email,
+                                "verified": False,
+                            }
+                        }
+                    },
+                )
+                trigger_email_verification = True
+
+        if existing_email_identity and not existing_email_identity.get("verified"):
+            # If this did exist, but is not yet verified, also trigger an email
+            trigger_email_verification = True
+
+        if trigger_email_verification:
+            token = _generate_and_store_token(
+                email=contact_email,
+                intent="verify",
+            )
+            try:
+                _send_magic_link_email(
+                    email=contact_email,
+                    token=token,
+                    referrer=CONFIG.APP_URL,
+                    purpose="verify",
+                )
+            except RuntimeError as e:
+                trigger_email_verification = False
+                LOGGER.critical(f"Unable to send verification email on this deployment: {e}")
 
-    except ValueError as e:
-        return jsonify(
-            {"status": "error", "message": f"Invalid display name or email was passed: {str(e)}"}
-        ), 400
+    try:
+        if account_status:
+            update["account_status"] = AccountStatus(account_status)
+    except ValueError:
+        raise BadRequest(f"Invalid account status {account_status!r} was passed")
 
     if not update:
-        return jsonify({"status": "success", "message": "No update was performed."}), 200
+        return jsonify({"status": "success", "message": "No update to perform."}), 200
 
     update_result = flask_mongo.db.users.update_one({"_id": ObjectId(user_id)}, {"$set": update})
 
     if update_result.matched_count != 1:
-        return (jsonify({"status": "error", "message": "Unable to update user."}), 400)
+        raise BadRequest("Unable to update user.")
 
-    if update_result.modified_count != 1:
+    if trigger_email_verification:
         return (
             jsonify(
-                {
-                    "status": "success",
-                    "message": "No update was performed",
-                }
+                {"message": f"Verification email sent to {contact_email}", "status": "success"}
             ),
             200,
         )
 
-    return (jsonify({"status": "success"}), 200)
+    return (jsonify({"message": "User updated successfully", "status": "success"}), 200)

pydatalab/tests/server/test_auth.py

@@ -19,7 +19,7 @@ def test_magic_link_account_creation(unauthenticated_client, app, database):
             "/login/magic-link",
             json={"email": "[email protected]", "referrer": "datalab.example.org"},
         )
-        assert response.json["detail"] == "Email sent successfully."
+        assert response.json["message"] == "Email sent successfully."
         assert response.status_code == 200
         assert len(outbox) == 1
 
@@ -46,7 +46,7 @@ def test_magic_links_expected_failures(unauthenticated_client, app):
         )
         assert response.status_code == 400
         assert len(outbox) == 0
-        assert response.json["detail"] == "Invalid email provided."
+        assert response.json["message"] == "Invalid email provided."
 
         response = unauthenticated_client.post(
             "/login/magic-link",

pydatalab/tests/server/test_users.py

@@ -56,6 +56,8 @@ def test_user_update(client, real_mongo_client, user_id, admin_user_id):
     assert resp.status_code == 200
     user = real_mongo_client.get_database().users.find_one({"_id": user_id})
     assert user["contact_email"] == "[email protected]"
+    assert user["identities"][-1]["identifier"] == "[email protected]"
+    assert not user["identities"][-1]["verified"]
 
     # Test that display name -> None does not remove display name
     user_request = {"display_name": None}