1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
|
import hashlib
import secrets
from fastapi import HTTPException, status
from generalresearch.redis_helper import RedisConfig
from jb.decorators import get_redis_config
from jb.models.auth import AmtAccountLink
MAGIC_TOKEN_PREFIX = "auth:magic:"
AMT_ACCOUNT_LINK_TOKEN_PREFIX = "auth:amt-account-link:"
MAGIC_TOKEN_TTL: int = 3 * 60 * 60 # 3hrs, in seconds
INVITE_AMT_LINK_TOKEN_TTL: int = 45 * 24 * 60 * 60 # 45 days
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"{prefix}{digest}"
def create_magic_token(user_email: str, redis_config: RedisConfig | None = None) -> str:
"""Create a short-lived, single-use token for a user.
The raw token can then be sent by email.
"""
if not user_email or not user_email.strip():
raise ValueError("user_email must not be empty")
token = secrets.token_urlsafe(32)
if redis_config is None:
redis_config = get_redis_config()
redis_client = redis_config.create_redis_client()
redis_client.set(
redis_token_key(token),
user_email,
ex=MAGIC_TOKEN_TTL,
)
return token
def consume_magic_token(token: str, redis_config: RedisConfig | None = None) -> str:
if redis_config is None:
redis_config = get_redis_config()
redis_client = redis_config.create_redis_client()
user_email = redis_client.getdel(redis_token_key(token))
if user_email is None:
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="Invalid or expired magic token",
)
return str(user_email)
def create_amt_account_link_token(
email: str,
amt_worker_id: str,
redis_config: RedisConfig | None = None,
ttl_seconds: int = MAGIC_TOKEN_TTL,
) -> str:
"""Bind an email and AMT worker ID to an opaque, short-lived token."""
if redis_config is None:
redis_config = get_redis_config()
redis_client = redis_config.create_redis_client()
data = AmtAccountLink(
email=email,
amt_worker_id=amt_worker_id,
)
token = secrets.token_urlsafe(32)
redis_client.set(
redis_token_key(token, AMT_ACCOUNT_LINK_TOKEN_PREFIX),
data.model_dump_json(),
ex=ttl_seconds,
)
return token
def consume_amt_account_link_token(
token: str, redis_config: RedisConfig | None = None
) -> AmtAccountLink:
"""Atomically consume and validate an AMT account-link token."""
if redis_config is None:
redis_config = get_redis_config()
redis_client = redis_config.create_redis_client()
raw_data = redis_client.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)
|
