Aether/_deprecated_py_src/services/auth/oauth/state.py

from __future__ import annotations

import json
import secrets
import time
from collections.abc import Awaitable
from dataclasses import dataclass
from typing import Any, cast

from redis.asyncio import Redis

OAUTH_STATE_TTL_SECONDS = 600
OAUTH_STATE_KEY_PREFIX = "oauth_state:"

# OAuth bind token: 用于安全地在浏览器跳转时传递用户身份
# 短期有效（5分钟），一次性使用
OAUTH_BIND_TOKEN_TTL_SECONDS = 300
OAUTH_BIND_TOKEN_KEY_PREFIX = "oauth_bind_token:"


CONSUME_STATE_SCRIPT = r"""
local value = redis.call("GET", KEYS[1])
if value then
    redis.call("DEL", KEYS[1])
end
return value
"""


@dataclass(frozen=True)
class OAuthStateData:
    nonce: str
    provider_type: str
    action: str  # "login" | "bind"
    user_id: str | None
    client_device_id: str | None
    created_at: int

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> OAuthStateData:
        return cls(
            nonce=str(data.get("nonce") or ""),
            provider_type=str(data.get("provider_type") or ""),
            action=str(data.get("action") or ""),
            user_id=data.get("user_id"),
            client_device_id=data.get("client_device_id"),
            created_at=int(data.get("created_at") or 0),
        )


def _state_key(nonce: str) -> str:
    return f"{OAUTH_STATE_KEY_PREFIX}{nonce}"


async def create_oauth_state(
    redis: Redis,
    *,
    provider_type: str,
    action: str,
    user_id: str | None = None,
    client_device_id: str | None = None,
) -> str:
    nonce = secrets.token_urlsafe(24)
    data = {
        "nonce": nonce,
        "provider_type": provider_type,
        "action": action,
        "user_id": user_id,
        "client_device_id": client_device_id,
        "created_at": int(time.time()),
    }
    await redis.setex(_state_key(nonce), OAUTH_STATE_TTL_SECONDS, json.dumps(data))
    return nonce


async def consume_oauth_state(redis: Redis, nonce: str) -> OAuthStateData | None:
    if not nonce:
        return None

    key = _state_key(nonce)
    # redis-py 的类型标注在 sync/async 之间会出现 Union；这里明确按 async 处理。
    raw = await cast(Awaitable[str | None], redis.eval(CONSUME_STATE_SCRIPT, 1, key))
    if not raw:
        return None

    try:
        parsed = json.loads(raw)
    except json.JSONDecodeError:
        return None

    return OAuthStateData.from_dict(parsed)


@dataclass(frozen=True)
class OAuthBindTokenData:
    """OAuth 绑定临时令牌数据，用于浏览器跳转场景的安全认证"""

    token: str
    user_id: str
    provider_type: str
    created_at: int

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> OAuthBindTokenData:
        return cls(
            token=str(data.get("token") or ""),
            user_id=str(data.get("user_id") or ""),
            provider_type=str(data.get("provider_type") or ""),
            created_at=int(data.get("created_at") or 0),
        )


def _bind_token_key(token: str) -> str:
    return f"{OAUTH_BIND_TOKEN_KEY_PREFIX}{token}"


async def create_oauth_bind_token(redis: Redis, *, user_id: str, provider_type: str) -> str:
    """创建一次性 OAuth 绑定令牌，用于浏览器跳转场景"""
    token = secrets.token_urlsafe(32)
    data = {
        "token": token,
        "user_id": user_id,
        "provider_type": provider_type,
        "created_at": int(time.time()),
    }
    await redis.setex(_bind_token_key(token), OAUTH_BIND_TOKEN_TTL_SECONDS, json.dumps(data))
    return token


async def consume_oauth_bind_token(redis: Redis, token: str) -> OAuthBindTokenData | None:
    """消费（验证并删除）OAuth 绑定令牌，返回令牌数据或 None"""
    if not token:
        return None

    key = _bind_token_key(token)
    raw = await cast(Awaitable[str | None], redis.eval(CONSUME_STATE_SCRIPT, 1, key))
    if not raw:
        return None

    try:
        parsed = json.loads(raw)
    except json.JSONDecodeError:
        return None

    return OAuthBindTokenData.from_dict(parsed)