wecom_it_smart_desk/backend/app/services/token_service.py

# =============================================================================
# 企微IT智能服务台 — 统一 Token 服务
# =============================================================================
# 说明：统一 Token 管理，支持以下功能：
#   1. 创建统一格式的 Token（包含角色信息）
#   2. 验证 Token 并获取用户信息
#   3. 切换当前角色
#   4. 兼容旧格式 Token（employee:token 和 agent:token）
# =============================================================================

import json
import logging
import secrets
from datetime import datetime
from typing import Dict, List, Optional

import redis.asyncio as aioredis

from app.config import settings

logger = logging.getLogger(__name__)

# Token TTL（8小时）
TOKEN_TTL_SECONDS = 8 * 60 * 60

# 统一 Token Key 前缀
UNIFIED_TOKEN_PREFIX = "user:token:"

# 旧格式 Token Key 前缀（兼容）
EMPLOYEE_TOKEN_PREFIX = "employee:token:"
AGENT_TOKEN_PREFIX = "agent:token:"


class TokenService:
    """统一 Token 服务。

    管理用户 Token 的创建、验证、角色切换等操作。
    """

    def __init__(self, redis_client: aioredis.Redis):
        """初始化 Token 服务。

        Args:
            redis_client: Redis 异步客户端
        """
        self.redis = redis_client

    async def create_token(
        self,
        employee_id: str,
        name: str,
        roles: List[str],
        department: Optional[str] = None,
        avatar: Optional[str] = None,
        login_source: str = "portal",
    ) -> str:
        """创建统一格式的 Token。

        Args:
            employee_id: 企微 UserID
            name: 用户姓名
            roles: 角色列表（如 ["user", "agent"]）
            department: 部门（可选）
            avatar: 头像URL（可选）
            login_source: 登录来源（portal/agent/h5）

        Returns:
            str: Token 字符串
        """
        # 生成 Token
        token = secrets.token_urlsafe(32)

        # 构建 Token 数据
        token_data = {
            "employee_id": employee_id,
            "name": name,
            "department": department or "",
            "avatar": avatar or "",
            "roles": roles,
            "current_role": self._get_default_role(roles),
            "login_source": login_source,
            "created_at": datetime.now().isoformat(),
            "last_active": datetime.now().isoformat(),
        }

        # 存入 Redis（统一格式）
        await self.redis.setex(
            f"{UNIFIED_TOKEN_PREFIX}{token}",
            TOKEN_TTL_SECONDS,
            json.dumps(token_data, ensure_ascii=False),
        )

        # 同时存入旧格式（兼容性）
        if login_source == "agent":
            await self.redis.setex(
                f"{AGENT_TOKEN_PREFIX}{token}",
                TOKEN_TTL_SECONDS,
                employee_id,
            )
        else:
            await self.redis.setex(
                f"{EMPLOYEE_TOKEN_PREFIX}{token}",
                TOKEN_TTL_SECONDS,
                employee_id,
            )

        logger.info(f"创建 Token: employee_id={employee_id}, roles={roles}, source={login_source}")

        return token

    async def get_user_info(self, token: str) -> Optional[Dict]:
        """获取 Token 对应的用户信息。

        支持新旧两种 Token 格式。

        Args:
            token: Token 字符串

        Returns:
            Optional[Dict]: 用户信息字典，如果 Token 无效返回 None
        """
        # 1. 尝试统一格式
        data = await self.redis.get(f"{UNIFIED_TOKEN_PREFIX}{token}")
        if data:
            try:
                user_info = json.loads(data)
                # 更新最后活跃时间
                user_info["last_active"] = datetime.now().isoformat()
                await self.redis.setex(
                    f"{UNIFIED_TOKEN_PREFIX}{token}",
                    TOKEN_TTL_SECONDS,
                    json.dumps(user_info, ensure_ascii=False),
                )
                return user_info
            except json.JSONDecodeError:
                logger.error(f"Token 数据解析失败: {token[:10]}...")
                return None

        # 2. 尝试旧格式（employee:token）
        employee_id = await self.redis.get(f"{EMPLOYEE_TOKEN_PREFIX}{token}")
        if employee_id:
            employee_id = employee_id.decode("utf-8") if isinstance(employee_id, bytes) else employee_id
            # 获取员工信息缓存
            info_data = await self.redis.get(f"employee:info:{employee_id}")
            if info_data:
                try:
                    info = json.loads(info_data)
                    return {
                        "employee_id": employee_id,
                        "name": info.get("employee_name", ""),
                        "department": info.get("department", ""),
                        "avatar": info.get("avatar", ""),
                        "roles": ["user"],
                        "current_role": "user",
                        "login_source": "h5",
                        "created_at": datetime.now().isoformat(),
                        "last_active": datetime.now().isoformat(),
                    }
                except json.JSONDecodeError:
                    pass

            # 降级：只有 employee_id
            return {
                "employee_id": employee_id,
                "name": "",
                "department": "",
                "avatar": "",
                "roles": ["user"],
                "current_role": "user",
                "login_source": "h5",
                "created_at": datetime.now().isoformat(),
                "last_active": datetime.now().isoformat(),
            }

        # 3. 尝试旧格式（agent:token）
        agent_id = await self.redis.get(f"{AGENT_TOKEN_PREFIX}{token}")
        if agent_id:
            agent_id = agent_id.decode("utf-8") if isinstance(agent_id, bytes) else agent_id
            return {
                "employee_id": agent_id,
                "name": "",
                "department": "",
                "avatar": "",
                "roles": ["user", "agent"],
                "current_role": "agent",
                "login_source": "agent",
                "created_at": datetime.now().isoformat(),
                "last_active": datetime.now().isoformat(),
            }

        return None

    async def switch_role(self, token: str, new_role: str) -> bool:
        """切换当前角色。

        Args:
            token: Token 字符串
            new_role: 目标角色标识

        Returns:
            bool: 是否切换成功
        """
        # 获取当前用户信息
        user_info = await self.get_user_info(token)
        if not user_info:
            return False

        # 验证用户是否有目标角色
        if new_role not in user_info.get("roles", []):
            logger.warning(f"用户 {user_info['employee_id']} 没有 {new_role} 角色")
            return False

        # 更新当前角色
        user_info["current_role"] = new_role
        user_info["last_active"] = datetime.now().isoformat()

        # 保存到 Redis（统一格式）
        await self.redis.setex(
            f"{UNIFIED_TOKEN_PREFIX}{token}",
            TOKEN_TTL_SECONDS,
            json.dumps(user_info, ensure_ascii=False),
        )

        # 同时更新旧格式（如果存在）
        # 注意：旧格式只存储 employee_id，不需要更新

        logger.info(f"用户 {user_info['employee_id']} 切换角色到 {new_role}")
        return True

    async def invalidate_token(self, token: str) -> None:
        """使 Token 失效。

        删除统一格式和旧格式的 Token。

        Args:
            token: Token 字符串
        """
        # 删除统一格式
        await self.redis.delete(f"{UNIFIED_TOKEN_PREFIX}{token}")

        # 删除旧格式
        await self.redis.delete(f"{EMPLOYEE_TOKEN_PREFIX}{token}")
        await self.redis.delete(f"{AGENT_TOKEN_PREFIX}{token}")

        logger.info(f"Token 已失效: {token[:10]}...")

    async def record_token_ip(self, token: str, ip_address: str) -> None:
        """记录Token使用的IP地址。

        用于后续的异常检测（如同一Token多IP使用）。

        Args:
            token: Token 字符串
            ip_address: 客户端IP地址
        """
        import hashlib
        import uuid

        # 计算token hash，避免存储明文token
        token_hash = hashlib.sha256(token.encode()).hexdigest()[:16]

        # Redis Key: token_ip:token_hash -> 逗号分隔的IP列表
        key = f"token_ip:{token_hash}"
        timestamp = datetime.now().isoformat()

        # 追加新IP和访问时间
        # 格式: "ip1@2026-07-14T10:00:00,ip2@2026-07-14T10:05:00"
        existing = await self.redis.get(key)
        if existing:
            existing_str = existing.decode("utf-8") if isinstance(existing, bytes) else existing
            # 只保留最近1小时的记录
            entries = existing_str.split(",")
            from datetime import timedelta
            one_hour_ago = datetime.now() - timedelta(hours=1)
            filtered = []
            for entry in entries:
                if "@" in entry:
                    ip, ts_str = entry.rsplit("@", 1)
                    try:
                        ts = datetime.fromisoformat(ts_str)
                        if ts > one_hour_ago:
                            filtered.append(entry)
                    except ValueError:
                        pass
            filtered.append(f"{ip_address}@{timestamp}")
            new_value = ",".join(filtered[-50:])  # 最多保留50条
        else:
            new_value = f"{ip_address}@{timestamp}"

        # 保留1小时
        await self.redis.setex(key, 3600, new_value)
        logger.debug(f"记录Token IP: token_hash={token_hash}, ip={ip_address}")

    async def get_token_ips(self, token: str) -> List[Dict]:
        """获取Token使用的IP列表。

        Args:
            token: Token 字符串

        Returns:
            List[Dict]: [{"ip": "x.x.x.x", "timestamp": "..."}]
        """
        import hashlib

        token_hash = hashlib.sha256(token.encode()).hexdigest()[:16]
        key = f"token_ip:{token_hash}"

        data = await self.redis.get(key)
        if not data:
            return []

        data_str = data.decode("utf-8") if isinstance(data, bytes) else data
        result = []
        for entry in data_str.split(","):
            if "@" in entry:
                ip, ts_str = entry.rsplit("@", 1)
                result.append({"ip": ip, "timestamp": ts_str})

        return result

    def _get_default_role(self, roles: List[str]) -> str:
        """获取默认角色。

        优先级：admin > agent > user

        Args:
            roles: 角色列表

        Returns:
            str: 默认角色标识
        """
        if "admin" in roles:
            return "admin"
        elif "agent" in roles:
            return "agent"
        else:
            return "user"