wecom_it_smart_desk/backend/app/api/auth_wecom_sso.py

# =============================================================================
# 企微IT智能服务台 — 企微入口 SSO（v0.7.1 新增）
# =============================================================================
# 说明: 解决 v0.7.0 hotfix1 用户报告的"企微工作台进入应用也要扫码"问题。
#
# 流程:
#   1. 前端 PortalSelect.vue 加载时检测 navigator.userAgent
#   2. 如果是 MicroMessenger / wxwork / DingTalk 等企微内置浏览器
#      → 调 /api/auth_wecom/sso/init?next=/itdesk/
#   3. 后端生成企微 OAuth2 授权 URL,302 跳转用户去企微授权
#   4. 企微回调 /api/auth_wecom/sso/callback?code=...&state=...
#   5. 用 code 换 userid,查 role (user/agent/admin),生成 token
#   6. 302 跳转到 next 路径 + token query param
#   7. 前端用 token 调 get_current_user 拉身份信息
#
# 配置要求:
#   - 企微管理后台 → 应用 → 网页授权及 JS-SDK → 可信域名: itsupport.servyou.com.cn
#   - 企微管理后台 → 应用 → 网页授权及 JS-SDK → 回调域: itsupport.servyou.com.cn
#   - 环境变量 WECOM_SSO_ENABLED=true 启用(默认 false,避免老用户被打扰)
# =============================================================================

import logging
import re
import secrets
import urllib.parse
from datetime import datetime, timedelta
from typing import Optional

from fastapi import APIRouter, Depends, Query, Request
from fastapi.responses import RedirectResponse
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from app.config import settings
from app.database import get_db
from app.models.role import Role
from app.models.user_role import UserRole
from app.services.wecom_service import WecomService
from app.services.audit_log_service import record_audit_log
from app.utils.response import AppException
from app.dependencies import get_redis

logger = logging.getLogger(__name__)

router = APIRouter(prefix="/auth_wecom", tags=["企微 SSO"])

# OAuth state 在 Redis 的 TTL (5 分钟,够用户授权 + 回调)
OAUTH_STATE_TTL = 300
# SSO token 长度
SSO_TOKEN_BYTES = 32
# Token TTL 常量（8小时）
TOKEN_TTL_SECONDS = 8 * 60 * 60  # 8小时
# 企微 API 超时设置（秒）
WECOM_API_TIMEOUT = 10

# --------------------------------------------------------------------------
# 企微环境检测（使用统一工具模块）
# --------------------------------------------------------------------------
from app.utils.wecom_auth import require_wecom_ua as _require_wework_ua


def _sso_enabled() -> bool:
    """检查是否启用企微 SSO。"""
    import os
    if os.getenv("WECOM_SSO_ENABLED", "false").lower() == "true":
        return True
    if getattr(settings, "wecom_sso_enabled", False):
        return True
    return False


def _get_oauth_callback_url(request: Request) -> str:
    """拼接 OAuth 回调 URL (绝对地址)。

    企微要求 redirect_uri 必须用可信域名(itsupport.servyou.com.cn)。
    不读 request.base_url 因为它可能是 127.0.0.1:8000(开发环境)。
    """
    # 优先用 settings 里的配置
    base = getattr(settings, "wecom_sso_callback_base", None)
    if not base:
        # 兜底: 读环境变量,默认生产域名
        import os
        base = os.getenv("WECOM_SSO_CALLBACK_BASE", "https://itsupport.servyou.com.cn")
    return f"{base.rstrip('/')}/api/auth_wecom/sso/callback"


def _build_oauth_url(state: str, callback_url: str) -> str:
    """拼企微 OAuth2 授权 URL。

    文档: https://developer.work.weixin.qq.com/document/path/91022
    """
    params = {
        "appid": settings.wecom_corp_id,
        "redirect_uri": callback_url,
        "response_type": "code",
        "scope": "snsapi_base",  # 静默授权
        "state": state,
        "agentid": settings.wecom_agent_id,
    }
    query = urllib.parse.urlencode(params)
    # 企业微信 OAuth2 地址（注意是 open.work.weixin.qq.com）
    return f"https://open.work.weixin.qq.com/connect/oauth2/authorize?{query}#wechat_redirect"


@router.get("/sso/init")
async def sso_init(
    request: Request,
    next: str = Query("/itdesk/", description="登录后跳转路径"),
    redis_client = Depends(get_redis),
):
    """初始化 SSO: 生成 state,302 跳转到企微 OAuth2 授权页。

    Args:
        next: 登录成功后跳转路径,如 /itdesk/ /itagent/ /itadmin/
    """
    # 后端第二道防线：非企微环境拒绝授权
    _require_wework_ua(request)

    if not _sso_enabled():
        raise AppException(1001, "企微 SSO 未启用, 请用扫码登录")

    # 1. 生成 state(防 CSRF + 携带 next 路径)
    state = secrets.token_urlsafe(24)
    state_payload = {
        "next": next,
        "created_at": datetime.now().isoformat(),
    }
    await redis_client.setex(
        f"wecom_sso:state:{state}",
        OAUTH_STATE_TTL,
        str(state_payload).encode("utf-8"),
    )

    # 2. 拼企微 OAuth URL（回调URL中包含next参数，用于state失效时仍能知道目标路径）
    callback_url = _get_oauth_callback_url(request)
    # 在回调URL中添加next参数
    separator = "&" if "?" in callback_url else "?"
    callback_url_with_next = f"{callback_url}{separator}next={urllib.parse.quote(next)}"
    oauth_url = _build_oauth_url(state, callback_url_with_next)

    logger.info(f"SSO init: state={state[:8]}..., next={next}")
    return RedirectResponse(url=oauth_url, status_code=302)


def _get_error_redirect_url(error_code: str, error_msg: str, next_path: str = "/itdesk/") -> str:
    """生成 OAuth 错误重定向 URL。

    异常时重定向到前端错误页面（ErrorPage），而不是返回 JSON 错误。
    ErrorPage 读取 code 和 message 参数显示友好错误提示。

    Args:
        error_code: 错误码
        error_msg: 错误信息
        next_path: 原始请求的目标路径（保留但不再用于决定重定向）
    """
    import os
    base = getattr(settings, "wecom_sso_callback_base", None)
    if not base:
        base = os.getenv("WECOM_SSO_CALLBACK_BASE", "https://itsupport.servyou.com.cn")

    # 重定向到 Portal 的 ErrorPage，带错误参数
    # ErrorPage 期望格式：?code=xxx&message=yyy
    return f"{base.rstrip('/')}/itportal/error?code={error_code}&message={urllib.parse.quote(error_msg)}"


@router.get("/sso/callback")
async def sso_callback(
    request: Request,
    code: Optional[str] = Query(None, description="企微 OAuth2 授权 code"),
    state: Optional[str] = Query(None, description="防 CSRF state"),
    errcode: Optional[int] = Query(None, description="企微 OAuth 错误码"),
    errmsg: Optional[str] = Query(None, description="企微 OAuth 错误信息"),
    next: Optional[str] = Query(None, description="原始请求的目标路径（可选，用于错误时重定向）"),
    redis_client = Depends(get_redis),
    db: AsyncSession = Depends(get_db),
):
    """企微 OAuth 回调: 用 code 换 userid → 查 role → 生成 token → 跳 next。

    异常时重定向到前端错误页面，避免白屏。
    所有未处理的异常都会记录详细日志（包含 traceback）。

    Args:
        next: 原始请求的目标路径，用于错误重定向。如果 state 验证失败，使用此参数决定重定向位置。
    """
    import traceback

    # 默认 next 路径
    next_path = next or "/itdesk/"

    try:
        # 后端第二道防线：非企微环境拒绝回调
        _require_wework_ua(request)

        # 0. 处理企微返回的错误（用户拒绝授权等）
        if errcode is not None:
            msg = errmsg or "用户取消授权或授权失败"
            logger.warning(f"SSO callback 企微返回错误: errcode={errcode}, errmsg={errmsg}")
            return RedirectResponse(url=_get_error_redirect_url(f"wecom_{errcode}", msg, next_path), status_code=302)

        # 1. 校验必要参数
        if not code or not state:
            logger.warning(f"SSO callback 缺少必要参数: code={bool(code)}, state={bool(state)}")
            return RedirectResponse(url=_get_error_redirect_url("missing_params", "授权参数不完整，请重试", next_path), status_code=302)

        # 2. 校验 state(防 CSRF)
        state_key = f"wecom_sso:state:{state}"
        try:
            state_raw = await redis_client.get(state_key)
        except Exception as e:
            logger.error(f"SSO callback Redis 获取 state 失败: {e}")
            return RedirectResponse(url=_get_error_redirect_url("redis_error", "服务暂不可用，请稍后重试", next_path), status_code=302)

        if not state_raw:
            logger.warning(f"SSO callback state 过期: state={state[:8]}...")
            return RedirectResponse(url=_get_error_redirect_url("state_expired", "授权已过期，请重新进入", next_path), status_code=302)

        # 删除 state(一次性)
        try:
            await redis_client.delete(state_key)
        except Exception as e:
            logger.warning(f"SSO callback 删除 state 失败: {e}")  # 不阻塞流程

        # 解析 state 数据（添加异常处理）
        import ast
        import json
        try:
            state_data = json.loads(state_raw.decode("utf-8"))
        except (json.JSONDecodeError, AttributeError) as e:
            # 兼容旧格式（使用 ast.literal_eval）
            try:
                state_data = ast.literal_eval(state_raw.decode("utf-8"))
            except (ValueError, SyntaxError) as e2:
                logger.error(f"SSO callback state 解析失败: {e2}")
                return RedirectResponse(url=_get_error_redirect_url("state_invalid", "授权信息无效，请重新进入", next_path), status_code=302)

        # 优先使用 state 中存储的 next，fallback 到 URL 参数
        next_path = state_data.get("next", next_path)

        # 3. 用 code 换 userid
        wecom = WecomService(redis_client)
        try:
            oauth_info = await wecom.get_oauth_user_info(code)
            user_id = oauth_info.get("userid", "")
            if not user_id:
                logger.warning("SSO callback 企微返回 userid 为空")
                return RedirectResponse(url=_get_error_redirect_url("empty_userid", "无法获取您的企业微信身份，请重试", next_path), status_code=302)

            user_info = await wecom.get_user_info(user_id)
            name = user_info.get("name", user_id)
        except Exception as e:
            logger.error(f"SSO callback 调企微 API 失败: code={code[:8]}..., error={e}")
            return RedirectResponse(url=_get_error_redirect_url("api_failed", f"企业微信服务异常: {str(e)}"), status_code=302)
        finally:
            try:
                await wecom.close()
            except Exception:
                pass

        # 3. 查 role (user/agent/admin)
        try:
            role_stmt = (
                select(Role)
                .join(UserRole, Role.id == UserRole.role_id)
                .where(UserRole.employee_id == user_id)
            )
            role_result = await db.execute(role_stmt)
            roles = role_result.scalars().all()
        except Exception as e:
            logger.error(f"SSO callback 查询角色失败: {e}")
            return RedirectResponse(url=_get_error_redirect_url("db_error", "服务暂不可用，请稍后重试", next_path), status_code=302)

        if not roles:
            # 没有绑定角色: 跳"无权限"页
            logger.warning(f"SSO: user_id={user_id} 没绑定任何角色")
            return RedirectResponse(url=f"/itdesk/no-role?user_id={user_id}", status_code=302)

        # 4. 选最高权限角色 (admin > agent > user)
        role_priority = {"admin": 3, "agent": 2, "user": 1}
        best_role = max(roles, key=lambda r: role_priority.get(r.name, 0))
        role_name = best_role.name

        # 5. 生成 SSO token(随机 + Redis 存 8 小时)
        sso_token = secrets.token_urlsafe(SSO_TOKEN_BYTES)
        sso_payload = {
            "user_id": user_id,
            "name": name,
            "role": role_name,
            "created_at": datetime.now().isoformat(),
        }
        import json
        try:
            await redis_client.setex(
                f"wecom_sso:token:{sso_token}",
                TOKEN_TTL_SECONDS,
                json.dumps(sso_payload, ensure_ascii=False).encode("utf-8"),
            )
        except Exception as e:
            logger.error(f"SSO callback 存储 token 失败: {e}")
            return RedirectResponse(url=_get_error_redirect_url("redis_error", "服务暂不可用，请稍后重试", next_path), status_code=302)

        # 6. 记录登录日志
        try:
            await record_audit_log(
                db=db,
                employee_id=user_id,
                action="sso_login",
                resource="auth",
                resource_id=user_id,
                details={"name": name, "role": role_name, "login_method": "wecom_sso"},
                result="success",
                request=None,  # callback 请求没有直接可用的 request 对象
            )
            await db.commit()
        except Exception as e:
            logger.warning(f"SSO callback 记录登录日志失败: {e}")  # 不阻塞登录流程

        # 7. 跳转到 next + token
        separator = "&" if "?" in next_path else "?"
        redirect_url = f"{next_path}{separator}sso_token={sso_token}"

        logger.info(f"SSO 成功: user_id={user_id}, role={role_name}, next={next_path}")
        return RedirectResponse(url=redirect_url, status_code=302)

    except Exception as e:
        # 捕获所有未处理的异常，记录详细日志（包含 traceback）并重定向到错误页
        error_details = {
            "error": str(e),
            "error_type": type(e).__name__,
            "code": code[:8] + "..." if code else None,
            "state": state[:8] + "..." if state else None,
            "next": next_path,
        }
        logger.error(
            f"SSO callback 未处理的异常: {error_details}\n"
            f"traceback: {traceback.format_exc()}"
        )
        return RedirectResponse(
            url=_get_error_redirect_url("oauth_failed", "登录过程出现异常，请重试"),
            status_code=302
        )


@router.get("/sso/verify")
async def sso_verify(
    request: Request,
    sso_token: str = Query(..., description="SSO token"),
    redis_client = Depends(get_redis),
    db: AsyncSession = Depends(get_db),
):
    """前端用 SSO token 换用户身份(token 一次性使用,用完删除)。

    后端第二道防线：非企微环境拒绝验证。
    """
    # 后端第二道防线：非企微环境拒绝验证
    _require_wework_ua(request)

    import json
    token_raw = await redis_client.get(f"wecom_sso:token:{sso_token}")
    if not token_raw:
        raise AppException(1005, "SSO token 已过期或无效")

    # 一次性 token(防止泄漏后被滥用)
    await redis_client.delete(f"wecom_sso:token:{sso_token}")

    payload = json.loads(token_raw.decode("utf-8"))
    return {
        "code": 0,
        "data": payload,
    }


@router.post("/refresh")
async def refresh_token(
    token: str = Query(..., description="当前 Bearer token"),
    redis_client = Depends(get_redis),
):
    """刷新 Token TTL。

    前端在 Token 过期前 5 分钟自动调用此接口，实现静默刷新。
    如果 Token 无效或已过期，返回 401 错误。

    Returns:
        刷新成功：{ code: 0, data: { token: "新token", expires_in: 28800 } }
    """
    import json

    # 1. 尝试统一格式 Token
    token_key = f"user:token:{token}"
    token_data_raw = await redis_client.get(token_key)

    if token_data_raw:
        try:
            user_info = json.loads(token_data_raw)
            # 更新最后活跃时间
            user_info["last_active"] = datetime.now().isoformat()

            # 延长 TTL（重新设置 8 小时）
            await redis_client.setex(
                token_key,
                TOKEN_TTL_SECONDS,
                json.dumps(user_info, ensure_ascii=False),
            )

            logger.info(f"Token 刷新成功: employee_id={user_info.get('employee_id')}")
            return {
                "code": 0,
                "data": {
                    "token": token,  # 复用同一个 token，只延长 TTL
                    "expires_in": TOKEN_TTL_SECONDS,
                },
            }
        except json.JSONDecodeError:
            pass

    # 2. 尝试旧格式 Token (employee:token)
    employee_key = f"employee:token:{token}"
    employee_id = await redis_client.get(employee_key)
    if employee_id:
        # 延长 TTL
        await redis_client.expire(employee_key, TOKEN_TTL_SECONDS)
        logger.info(f"Token 刷新成功(employee): employee_id={employee_id}")
        return {
            "code": 0,
            "data": {
                "token": token,
                "expires_in": TOKEN_TTL_SECONDS,
            },
        }

    # 3. 尝试旧格式 Token (agent:token)
    agent_key = f"agent:token:{token}"
    agent_id = await redis_client.get(agent_key)
    if agent_id:
        await redis_client.expire(agent_key, TOKEN_TTL_SECONDS)
        logger.info(f"Token 刷新成功(agent): agent_id={agent_id}")
        return {
            "code": 0,
            "data": {
                "token": token,
                "expires_in": TOKEN_TTL_SECONDS,
            },
        }

    # Token 无效或已过期
    logger.warning(f"Token 刷新失败: token 不存在或已过期")
    raise AppException(401, "Token 已过期，请重新登录")


# --------------------------------------------------------------------------
# 别名路由：支持前端 /api/auth/refresh 调用（与 /api/auth_wecom/refresh 等效）
# --------------------------------------------------------------------------
# 前端 H5/坐席/管理后台调用 /api/auth/refresh，后端响应 /api/auth_wecom/refresh
# 为兼容前端习惯，添加此别名路由
# --------------------------------------------------------------------------

# 创建别名路由器（无 prefix）
alias_router = APIRouter(tags=["认证"])


@alias_router.post("/auth/refresh")
async def refresh_token_alias(
    token: str = Query(..., description="当前 Bearer token"),
    redis_client = Depends(get_redis),
):
    """Token 刷新接口别名。

    前端调用 /api/auth/refresh，后端实际处理逻辑与 /api/auth_wecom/refresh 相同。
    这是为了兼容前端的调用习惯。

    Returns:
        刷新成功：{ code: 0, data: { token: "新token", expires_in: 28800 } }
    """
    import json

    # 1. 尝试统一格式 Token
    token_key = f"user:token:{token}"
    token_data_raw = await redis_client.get(token_key)

    if token_data_raw:
        try:
            user_info = json.loads(token_data_raw)
            user_info["last_active"] = datetime.now().isoformat()
            await redis_client.setex(
                token_key,
                TOKEN_TTL_SECONDS,
                json.dumps(user_info, ensure_ascii=False),
            )
            logger.info(f"Token 刷新成功(alias): employee_id={user_info.get('employee_id')}")
            return {
                "code": 0,
                "data": {
                    "token": token,
                    "expires_in": TOKEN_TTL_SECONDS,
                },
            }
        except json.JSONDecodeError:
            pass

    # 2. 尝试旧格式 Token
    employee_key = f"employee:token:{token}"
    employee_id = await redis_client.get(employee_key)
    if employee_id:
        await redis_client.expire(employee_key, TOKEN_TTL_SECONDS)
        logger.info(f"Token 刷新成功(alias employee): employee_id={employee_id}")
        return {
            "code": 0,
            "data": {
                "token": token,
                "expires_in": TOKEN_TTL_SECONDS,
            },
        }

    # 3. 尝试 agent token
    agent_key = f"agent:token:{token}"
    agent_id = await redis_client.get(agent_key)
    if agent_id:
        await redis_client.expire(agent_key, TOKEN_TTL_SECONDS)
        logger.info(f"Token 刷新成功(alias agent): agent_id={agent_id}")
        return {
            "code": 0,
            "data": {
                "token": token,
                "expires_in": TOKEN_TTL_SECONDS,
            },
        }

    logger.warning(f"Token 刷新失败(alias): token 不存在或已过期")
    raise AppException(401, "Token 已过期，请重新登录")