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 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.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


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)
    return f"https://open.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/
    """
    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
    callback_url = _get_oauth_callback_url(request)
    oauth_url = _build_oauth_url(state, callback_url)

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


@router.get("/sso/callback")
async def sso_callback(
    code: str = Query(..., description="企微 OAuth2 授权 code"),
    state: str = Query(..., description="防 CSRF state"),
    redis_client = Depends(get_redis),
    db: AsyncSession = Depends(get_db),
):
    """企微 OAuth 回调: 用 code 换 userid → 查 role → 生成 token → 跳 next。"""
    # 1. 校验 state(防 CSRF)
    state_key = f"wecom_sso:state:{state}"
    state_raw = await redis_client.get(state_key)
    if not state_raw:
        raise AppException(1002, "SSO state 已过期或无效, 请重新进入")

    # 删除 state(一次性)
    await redis_client.delete(state_key)

    import ast
    state_data = ast.literal_eval(state_raw.decode("utf-8"))
    next_path = state_data.get("next", "/itdesk/")

    # 2. 用 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:
            raise AppException(1003, "企微 OAuth 返回 userid 为空")

        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}")
        raise AppException(1004, f"企微身份识别失败: {str(e)}")
    finally:
        try:
            await wecom.close()
        except Exception:
            pass

    # 3. 查 role (user/agent/admin)
    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()

    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
    await redis_client.setex(
        f"wecom_sso:token:{sso_token}",
        8 * 3600,  # 8 小时
        json.dumps(sso_payload, ensure_ascii=False).encode("utf-8"),
    )

    # 6. 跳转到 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)


@router.get("/sso/verify")
async def sso_verify(
    sso_token: str = Query(..., description="SSO token"),
    redis_client = Depends(get_redis),
    db: AsyncSession = Depends(get_db),
):
    """前端用 SSO token 换用户身份(token 一次性使用,用完删除)。"""
    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,
    }