""" 微信 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