Merges pull request #3

Off of Amazon!!!

3 files changed, 184 insertions, 0 deletions

diff --git a/jb/api/__init__.py b/jb/api/__init__.py
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/jb/api/__init__.py
diff --git a/jb/api/auth.py b/jb/api/auth.py
new file mode 100644
index 0000000..1542e70
--- /dev/null
+++ b/jb/api/auth.py
@@ -0,0 +1,88 @@
+from datetime import datetime, timedelta, timezone
+from typing import Annotated
+from uuid import uuid4
+
+import jwt
+from fastapi import Depends, HTTPException, Request, status
+from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
+
+from jb.config import settings
+from jb.dependencies import get_gr_api_manager
+from jb.managers.gr_api import GRApiManager
+from jb.models.auth import User
+
+bearer = HTTPBearer(auto_error=False)
+
+SESSION_COOKIE_NAME = "jb_session"
+JWT_ISSUER = "jamesbillings67"
+JWT_AUDIENCE = "jamesbillings67"
+
+
+def get_authenticated_user(
+    request: Request,
+    credentials: Annotated[HTTPAuthorizationCredentials | None, Depends(bearer)],
+    gr_api: Annotated[GRApiManager, Depends(get_gr_api_manager)],
+) -> User:
+    """FastAPI dependency for endpoints requiring a valid session."""
+    if settings.session_jwt_secret is None:
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail="Account session signing key is not configured",
+        )
+
+    token = request.cookies.get(SESSION_COOKIE_NAME)
+    if credentials is not None and credentials.scheme.lower() == "bearer":
+        token = credentials.credentials
+
+    if token is None:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Missing session token",
+            headers={"WWW-Authenticate": "Bearer"},
+        )
+
+    try:
+        claims = jwt.decode(
+            token,
+            settings.session_jwt_secret.get_secret_value(),
+            algorithms=["HS256"],
+            issuer=JWT_ISSUER,
+            audience=JWT_AUDIENCE,
+            options={"require": ["sub", "iat", "exp", "jti", "iss", "aud", "type"]},
+        )
+    except jwt.PyJWTError:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Invalid or expired session token",
+            headers={"WWW-Authenticate": "Bearer"},
+        )
+
+    if claims["type"] != "session" or not isinstance(claims["sub"], str):
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Invalid or expired session token",
+            headers={"WWW-Authenticate": "Bearer"},
+        )
+    product_user_id = claims.get("sub")
+
+    user = gr_api.get_user(product_user_id=product_user_id)
+
+    return user
+
+
+def create_session(product_user_id: str) -> str:
+    now = datetime.now(timezone.utc)
+    assert settings.session_jwt_secret
+    return jwt.encode(
+        {
+            "sub": product_user_id,
+            "iat": now,
+            "exp": now + timedelta(seconds=settings.session_token_ttl_seconds),
+            "jti": uuid4().hex,
+            "iss": JWT_ISSUER,
+            "aud": JWT_AUDIENCE,
+            "type": "session",
+        },
+        settings.session_jwt_secret.get_secret_value(),
+        algorithm="HS256",
+    )
diff --git a/jb/api/magic_token.py b/jb/api/magic_token.py
new file mode 100644
index 0000000..b60d5e9
--- /dev/null
+++ b/jb/api/magic_token.py
@@ -0,0 +1,96 @@
+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 = 10 * 60  # 10 minutes, in seconds
+
+
+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
+) -> 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=MAGIC_TOKEN_TTL,
+    )
+    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)