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 分钟
STATE_KEY_PREFIX = "wechat:state:"  # Redis key 前缀（独立逻辑命名空间）


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]


class RedisStateStore:
    """Redis state 存储（多实例/容器重启安全）。

    复用现有 Redis（celery broker 同实例），key 前缀 wechat:state:，
    TTL 10 分钟，SET NX EX + GETDEL 保证一次性消费。
    Redis 不可用时降级为内存存储，保证登录流程不中断（单节点场景）。
    """

    def __init__(
        self,
        redis_url: str = "",
        ttl_seconds: int = STATE_TTL_SECONDS,
        key_prefix: str = STATE_KEY_PREFIX,
        client=None,
    ):
        self._ttl = ttl_seconds
        self._prefix = key_prefix
        self._fallback = MemoryStateStore(ttl_seconds=ttl_seconds)
        self._redis = None
        if client is not None:
            # 测试/显式注入
            self._redis = client
            return
        try:
            import redis

            self._redis = redis.Redis.from_url(redis_url, decode_responses=True)
            self._redis.ping()
            logger.info(
                "微信 state 存储使用 Redis: %s db=%s",
                self._redis.connection_pool.connection_kwargs.get("host"),
                self._redis.connection_pool.connection_kwargs.get("db"),
            )
        except Exception as e:  # noqa: BLE001 — Redis 不可用降级内存，登录流程不中断
            logger.warning("微信 state Redis 不可用，降级为内存存储: %s", e)
            self._redis = None

    def _key(self, state: str) -> str:
        return f"{self._prefix}{state}"

    def put(self, state: str) -> None:
        if self._redis is None:
            self._fallback.put(state)
            return
        try:
            # SET key 1 NX EX ttl：不存在才写入，自带过期
            self._redis.set(self._key(state), "1", nx=True, ex=self._ttl)
        except Exception as e:  # noqa: BLE001
            logger.warning("微信 state 写入 Redis 失败，降级内存: %s", e)
            self._fallback.put(state)

    # Lua：原子读取并删除（单线程执行），兼容所有 Redis 版本（GETDEL 需 6.2+）
    _CONSUME_LUA = """
local v = redis.call('GET', KEYS[1])
if v then redis.call('DEL', KEYS[1]) end
return v
"""

    def verify_and_consume(self, state: str) -> bool:
        if self._redis is None:
            return self._fallback.verify_and_consume(state)
        try:
            try:
                val = self._redis.eval(self._CONSUME_LUA, 1, self._key(state))
            except Exception:  # noqa: BLE001 — eval 不可用时退化 GET+DELETE
                val = self._redis.get(self._key(state))
                if val is not None:
                    self._redis.delete(self._key(state))
            return val is not None
        except Exception as e:  # noqa: BLE001
            logger.warning("微信 state 校验 Redis 失败，降级内存: %s", e)
            return self._fallback.verify_and_consume(state)


@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)
            # 微信响应头不带 charset，requests 默认按 ISO-8859-1 解码会导致中文乱码
            token_resp.encoding = "utf-8"
            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)
            # 同上：显式 UTF-8 解码，保证中文昵称/unionid 等不乱码
            user_resp.encoding = "utf-8"
            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 _build_default_state_store():
    """默认 state 存储：优先 Redis（多实例/重启安全），不可用由 store 内部降级内存。"""
    redis_url = ""
    try:
        from app.config import get_settings

        redis_url = get_settings().CELERY_BROKER_URL or get_settings().REDIS_URL
    except Exception:  # noqa: BLE001 — API 配置不可用时退回环境变量
        redis_url = os.environ.get("CELERY_BROKER_URL", "") or os.environ.get("REDIS_URL", "")
    if redis_url:
        return RedisStateStore(redis_url)
    return MemoryStateStore()


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