xiaoxia-saas/packages/application/auth/wechat_oauth_service.py

"""
微信 OAuth 服务
- 生成授权链接（网页扫码登录）
- 处理回调，用 code 换 access_token + 用户信息
"""

from __future__ import annotations

import logging
import os
import time
import urllib.parse
from dataclasses import dataclass
from threading import Lock
from typing import Optional
from uuid import uuid4

import requests

logger = logging.getLogger(__name__)

STATE_TTL_SECONDS = 600  # state 有效期 10 分钟


class MemoryStateStore:
    """内存 state 存储（简单实现，单节点可用）

    多实例部署时建议替换为 Redis 实现。
    """

    def __init__(self, ttl_seconds: int = STATE_TTL_SECONDS):
        self._ttl = ttl_seconds
        self._states: dict[str, float] = {}  # state -> expire_at
        self._lock = Lock()

    def put(self, state: str) -> None:
        with self._lock:
            self._clean_expired()
            self._states[state] = time.time() + self._ttl

    def verify_and_consume(self, state: str) -> bool:
        with self._lock:
            self._clean_expired()
            if state in self._states:
                del self._states[state]
                return True
            return False

    def _clean_expired(self) -> None:
        now = time.time()
        expired = [s for s, exp in self._states.items() if exp < now]
        for s in expired:
            del self._states[s]


@dataclass
class WechatUserInfo:
    """微信用户信息"""

    openid: str
    unionid: str = ""
    nickname: str = ""
    avatar_url: str = ""


class WechatOAuthService:
    """微信开放平台 OAuth 服务（网页扫码登录）"""

    def __init__(
        self,
        app_id: str | None = None,
        app_secret: str | None = None,
        redirect_uri: str | None = None,
        state_store=None,
    ):
        self.app_id = app_id or os.environ.get("WECHAT_OPEN_APP_ID", "")
        self.app_secret = app_secret or os.environ.get("WECHAT_OPEN_APP_SECRET", "")
        self.redirect_uri = redirect_uri or os.environ.get("WECHAT_OPEN_REDIRECT_URI", "")
        # state 存储（CSRF 防护），默认内存实现
        self._state_store = state_store or MemoryStateStore()

    def is_configured(self) -> bool:
        """检查微信配置是否完整"""
        return bool(self.app_id and self.app_secret and self.redirect_uri)

    def generate_auth_url(self, scope: str = "snsapi_login") -> tuple[str, str]:
        """
        生成微信授权链接

        Returns:
            (授权URL, state)
        """
        state = uuid4().hex
        # 保存 state 用于回调校验（防 CSRF）
        self._state_store.put(state)

        if not self.is_configured():
            # 未配置时返回 mock URL，方便前端联调
            mock_params = urllib.parse.urlencode(
                {
                    "app_id": "mock",
                    "redirect_uri": self.redirect_uri,
                    "scope": scope,
                    "state": state,
                }
            )
            return f"/mock/wechat/auth?{mock_params}", state

        params = {
            "appid": self.app_id,
            "redirect_uri": self.redirect_uri,
            "response_type": "code",
            "scope": scope,
            "state": state,
        }
        url = "https://open.weixin.qq.com/connect/qrconnect?" + urllib.parse.urlencode(params) + "#wechat_redirect"
        return url, state

    def handle_callback(self, code: str, state: str) -> tuple[Optional[WechatUserInfo], Optional[str]]:
        """
        处理微信回调

        Args:
            code: 微信授权码
            state: 防 CSRF 状态

        Returns:
            (微信用户信息, 错误信息)
        """
        if not code:
            return None, "缺少授权码"

        # 校验 state（防 CSRF）—— 一次性使用
        if not state or not self._state_store.verify_and_consume(state):
            logger.warning("微信回调 state 校验失败: state=%s", state)
            return None, "无效的 state 参数，请求可能已过期或被篡改"

        if not self.is_configured():
            # 开发模式：返回 mock 用户信息
            logger.info("微信未配置，使用 mock 用户信息")
            return (
                WechatUserInfo(
                    openid=f"mock_{code[:20]}",
                    unionid=f"mock_union_{code[:16]}",
                    nickname="微信测试用户",
                    avatar_url="",
                ),
                None,
            )

        try:
            # 1. 用 code 换 access_token
            token_url = "https://api.weixin.qq.com/sns/oauth2/access_token"
            token_params = {
                "appid": self.app_id,
                "secret": self.app_secret,
                "code": code,
                "grant_type": "authorization_code",
            }
            token_resp = requests.get(token_url, params=token_params, timeout=10)
            token_data = token_resp.json()

            if "errcode" in token_data and token_data["errcode"] != 0:
                logger.error("微信获取 access_token 失败: %s", token_data)
                return None, f"微信授权失败: {token_data.get('errmsg', '未知错误')}"

            access_token = token_data["access_token"]
            openid = token_data["openid"]
            unionid = token_data.get("unionid", "")

            # 2. 获取用户信息
            user_url = "https://api.weixin.qq.com/sns/userinfo"
            user_params = {
                "access_token": access_token,
                "openid": openid,
                "lang": "zh_CN",
            }
            user_resp = requests.get(user_url, params=user_params, timeout=10)
            user_data = user_resp.json()

            if "errcode" in user_data and user_data["errcode"] != 0:
                logger.error("微信获取用户信息失败: %s", user_data)
                return None, f"获取用户信息失败: {user_data.get('errmsg', '未知错误')}"

            return (
                WechatUserInfo(
                    openid=openid,
                    unionid=unionid,
                    nickname=user_data.get("nickname", ""),
                    avatar_url=user_data.get("headimgurl", ""),
                ),
                None,
            )

        except requests.RequestException as e:
            logger.error("微信 OAuth 请求异常: %s", e, exc_info=True)
            return None, "微信服务暂不可用，请稍后再试"
        except Exception as e:
            logger.error("微信回调处理异常: %s", e, exc_info=True)
            return None, "微信登录处理失败"


# 模块级单例：state 存储必须跨请求共享，否则 /wechat/url 生成的 state
# 与 /wechat/callback 校验时不在同一个 MemoryStateStore，回调必然 400。
# 多实例部署时应替换为 Redis state store（单容器多 worker 也需如此）。
_oauth_service_singleton: WechatOAuthService | None = None


def get_wechat_oauth_service() -> WechatOAuthService:
    """获取微信 OAuth 服务单例（state store 跨请求共享）"""
    global _oauth_service_singleton
    if _oauth_service_singleton is None:
        _oauth_service_singleton = WechatOAuthService()
    return _oauth_service_singleton