1 files changed, 32 insertions, 2 deletions
diff --git a/jb/api/magic_token.py b/jb/api/magic_token.py
index 136e1b4..e0a1cca 100644
--- a/jb/api/magic_token.py
+++ b/jb/api/magic_token.py
@@ -4,15 +4,17 @@ import secrets
 from fastapi import HTTPException, status
 
 from jb.decorators import REDIS
+from jb.models.auth import AmtAccountLink, User
 
 MAGIC_TOKEN_PREFIX = "auth:magic:"
+AMT_ACCOUNT_LINK_TOKEN_PREFIX = "auth:amt-account-link:"
 MAGIC_TOKEN_TTL: int = 5 * 60  # 5 minutes, in seconds
 
 
-def redis_token_key(token: str) -> str:
+def redis_token_key(token: str, prefix: str = MAGIC_TOKEN_PREFIX) -> str:
     # Redis never contains a usable credential, even if its keys are exposed.
     digest = hashlib.sha256(token.encode("utf-8")).hexdigest()
-    return f"{MAGIC_TOKEN_PREFIX}{digest}"
+    return f"{prefix}{digest}"
 
 
 def create_magic_token(user_email: str) -> str:
@@ -39,3 +41,31 @@ def consume_magic_token(token: str) -> str:
             detail="Invalid or expired magic token",
         )
     return user_email
+
+
+def create_amt_account_link_token(user: User, amt_worker_id: str) -> str:
+    """Bind an email and AMT worker ID to an opaque, short-lived token."""
+    data = AmtAccountLink(
+        email=user.email,
+        amt_worker_id=amt_worker_id,
+    )
+    token = secrets.token_urlsafe(32)
+    REDIS.set(
+        redis_token_key(token, AMT_ACCOUNT_LINK_TOKEN_PREFIX),
+        data.model_dump_json(),
+        ex=MAGIC_TOKEN_TTL,
+    )
+    return token
+
+
+def consume_amt_account_link_token(token: str) -> AmtAccountLink:
+    """Atomically consume and validate an AMT account-link token."""
+    raw_data = REDIS.getdel(
+        redis_token_key(token, AMT_ACCOUNT_LINK_TOKEN_PREFIX)
+    )
+    if raw_data is None:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Invalid or expired account-link token",
+        )
+    return AmtAccountLink.model_validate_json(raw_data)