hello-agents/Co-creation-projects/2248652135-Trip-Planner-Pro/backend/app/redis_service.py

"""Redis服务封装"""
import json
from datetime import timedelta
import redis as redis_module
import os

REDIS_HOST = os.getenv("REDIS_HOST") 
REDIS_PORT = int(os.getenv("REDIS_PORT"))
REDIS_PASSWORD = os.getenv("REDIS_PASSWORD")
REDIS_DB = int(os.getenv("REDIS_DB"))

# Refresh Token 过期时间（与JWT refresh token一致）
REFRESH_TOKEN_TTL = timedelta(days=7)

# Key 前缀
PREFIX_REFRESH = "refresh_token:"   # refresh_token:{jti} -> user_id
PREFIX_USER_TOKENS = "user_tokens:"  # user_tokens:{user_id} -> set of jti


def get_redis() -> redis_module.Redis:
    """获取Redis连接"""
    return redis_module.Redis(
        host=REDIS_HOST,
        port=REDIS_PORT,
        password=REDIS_PASSWORD,
        db=REDIS_DB,
        decode_responses=True,
        socket_connect_timeout=3,
    )


def store_refresh_token(user_id: int, jti: str, ttl_seconds: int = None, user_agent: str = "") -> bool:
    """
    将 Refresh Token 的 jti 存入 Redis
    - refresh_token:{jti} -> json({"user_id": ..., "user_agent": ...}) （正向查：token -> user + 设备）
    - user_tokens:{user_id} -> set of jti （反向查：user -> tokens，用于踢下线）
    """
    if ttl_seconds is None:
        ttl_seconds = int(REFRESH_TOKEN_TTL.total_seconds())

    r = get_redis()
    try:
        pipe = r.pipeline()
        data = json.dumps({"user_id": user_id, "user_agent": user_agent})
        pipe.setex(f"{PREFIX_REFRESH}{jti}", ttl_seconds, data)
        pipe.sadd(f"{PREFIX_USER_TOKENS}{user_id}", jti)
        pipe.expire(f"{PREFIX_USER_TOKENS}{user_id}", ttl_seconds)
        pipe.execute()
        return True
    finally:
        r.close()


def validate_refresh_token(jti: str) -> dict | None:
    """
    验证 Refresh Token 是否在 Redis 中有效
    返回值: {"user_id": int, "user_agent": str} 或 None
    """
    r = get_redis()
    try:
        data = r.get(f"{PREFIX_REFRESH}{jti}")
        if data is None:
            return None
        parsed = json.loads(data)
        return {
            "user_id": int(parsed["user_id"]),
            "user_agent": parsed.get("user_agent", ""),
        }
    finally:
        r.close()


def revoke_refresh_token(jti: str) -> bool:
    """吊销单个 Refresh Token"""
    r = get_redis()
    try:
        data = r.get(f"{PREFIX_REFRESH}{jti}")
        if data is None:
            return False
        parsed = json.loads(data)
        uid = int(parsed["user_id"])
        pipe = r.pipeline()
        pipe.delete(f"{PREFIX_REFRESH}{jti}")
        pipe.srem(f"{PREFIX_USER_TOKENS}{uid}", jti)
        pipe.execute()
        return True
    finally:
        r.close()


def revoke_all_user_tokens(user_id: int) -> int:
    """吊销用户的所有 Refresh Token（全部踢下线）"""
    r = get_redis()
    try:
        jtis = r.smembers(f"{PREFIX_USER_TOKENS}{user_id}")
        if not jtis:
            return 0
        pipe = r.pipeline()
        for jti in jtis:
            pipe.delete(f"{PREFIX_REFRESH}{jti}")
        pipe.delete(f"{PREFIX_USER_TOKENS}{user_id}")
        pipe.execute()
        return len(jtis)
    finally:
        r.close()