# -*- coding: utf-8 -*- """微信登录服务（小程序 code2session + 开放平台网站应用扫码登录）。 - ``code2session``：小程序 ``wx.login`` 的 code 换 openid。 - ``qr_connect_url``：生成开放平台「网站应用扫码登录」授权链接（前端渲染成二维码）。 - ``oauth_code2openid``：扫码授权回跳的 code 换 openid（OAuth2 access_token 端点）。 未配置对应凭据（小程序 WECHAT_APPID/SECRET、扫码 WECHAT_OPEN_APPID/SECRET）时走 「测试直通」：把 code 当 openid 返回，便于无真实凭据联调（生产必须配置真实凭据）。 """ from __future__ import annotations import secrets import threading import time from urllib.parse import urlencode import httpx from .. import config _WX_SESSION_URL = "https://api.weixin.qq.com/sns/jscode2session" _WX_OAUTH_TOKEN_URL = "https://api.weixin.qq.com/sns/oauth2/access_token" _WX_QR_BASE_URL = "https://open.weixin.qq.com/connect/qrconnect" class WechatError(Exception): """微信登录异常（code 无效 / 网络错误等）。""" async def code2session(code: str) -> str: """小程序登录 code 换 openid。未配置 appid 时直通：``openid = code``。""" code = code.strip() if not config.WECHAT_APPID or not config.WECHAT_SECRET: return code try: async with httpx.AsyncClient(timeout=10) as client: resp = await client.get( _WX_SESSION_URL, params={ "appid": config.WECHAT_APPID, "secret": config.WECHAT_SECRET, "js_code": code, "grant_type": "authorization_code", }, ) data = resp.json() except Exception as exc: # noqa: BLE001 raise WechatError(f"wechat code2session request failed: {exc}") from exc errcode = data.get("errcode", 0) if errcode: raise WechatError( f"wechat code2session error {errcode}: {data.get('errmsg', '')}" ) openid = data.get("openid") if not openid: raise WechatError("wechat code2session returned no openid") return openid # --------------------------------------------------------------------------- # 开放平台「网站应用」扫码登录 # --------------------------------------------------------------------------- def qr_connect_url(state: str, redirect_uri: str) -> str: """生成扫码登录授权链接（前端据其渲染二维码）。 扫码后在微信内完成授权 → 微信重定向到 ``redirect_uri?code=&state=``， 服务端 ``/auth/wx-callback`` 收到后换 openid 并存到对应 state 会话。 """ params = { "appid": config.WECHAT_OPEN_APPID, "redirect_uri": redirect_uri, "response_type": "code", "scope": "snsapi_login", "state": state, } return f"{_WX_QR_BASE_URL}?{urlencode(params)}#wechat_redirect" async def oauth_code2openid(code: str) -> str: """扫码授权回跳的 code 换 openid。未配置开放平台凭据时直通：``openid = code + '_qr'``。""" code = code.strip() if not config.WECHAT_OPEN_APPID or not config.WECHAT_OPEN_SECRET: return f"{code}_qr" try: async with httpx.AsyncClient(timeout=10) as client: resp = await client.get( _WX_OAUTH_TOKEN_URL, params={ "appid": config.WECHAT_OPEN_APPID, "secret": config.WECHAT_OPEN_SECRET, "code": code, "grant_type": "authorization_code", }, ) data = resp.json() except Exception as exc: # noqa: BLE001 raise WechatError(f"wechat oauth token request failed: {exc}") from exc errcode = data.get("errcode", 0) or data.get("errCode", 0) if errcode: raise WechatError(f"wechat oauth error {errcode}: {data.get('errmsg', '')}") openid = data.get("openid") if not openid: raise WechatError("wechat oauth returned no openid") return openid # --------------------------------------------------------------------------- # 小程序码（wxacode）—— 微信「扫一扫」自动打开小程序到指定页面 + 携带 scene # --------------------------------------------------------------------------- _WX_ACODE_URL = "https://api.weixin.qq.com/wxa/getwxacodeunlimit?access_token={token}" _WX_TOKEN_URL = "https://api.weixin.qq.com/cgi-bin/token?grant_type=client_credential&appid={appid}&secret={secret}" # 缓存 access_token（模块级，线程安全；微信 token 有效期约 7200s） _token_cache: dict = {"token": "", "expires_at": 0.0} _token_lock = threading.Lock() async def _wx_access_token() -> str: """获取小程序 access_token（缓存）。未配置 WECHAT_APPID/SECRET 时抛 WechatError。""" if not config.WECHAT_APPID or not config.WECHAT_SECRET: raise WechatError("WECHAT_APPID/WECHAT_SECRET 未配置，无法生成小程序码") now = time.monotonic() with _token_lock: if _token_cache["token"] and now < _token_cache["expires_at"]: return _token_cache["token"] try: async with httpx.AsyncClient(timeout=10) as client: resp = await client.get( _WX_TOKEN_URL.format(appid=config.WECHAT_APPID, secret=config.WECHAT_SECRET), ) data = resp.json() except Exception as exc: # noqa: BLE001 raise WechatError(f"wechat access_token request failed: {exc}") from exc token = data.get("access_token", "") if not token: raise WechatError(f"wechat access_token error: {data.get('errmsg', data)}") with _token_lock: _token_cache["token"] = token _token_cache["expires_at"] = now + int(data.get("expires_in", 7200)) - 300 return token async def get_wxacode(scene: str, page: str = "pages/scan-login/index", env_version: str | None = None) -> bytes: """生成小程序码（wxacode）图片，扫码自动打开小程序指定页面并携带 scene。 env_version：release(正式版) | trial(体验版) | develop(开发版)，默认取 config.WECHAT_WXACODE_ENV （随 APP_ENV：prod→release，dev→trial，可显式 PINEAGENTS_WXACODE_ENV 覆盖）。 scene 限 32 字符（本服务用短 token）。返回 PNG 字节。 """ token = await _wx_access_token() env = env_version or config.WECHAT_WXACODE_ENV try: async with httpx.AsyncClient(timeout=15) as client: resp = await client.post( _WX_ACODE_URL.format(token=token), json={"scene": scene, "page": page, "check_path": False, "env_version": env}, ) except Exception as exc: # noqa: BLE001 raise WechatError(f"wechat getwxacodeunlimit failed: {exc}") from exc if resp.status_code != 200 or resp.headers.get("content-type", "").startswith("application/json"): # 出错时返回 JSON 错误体 try: err = resp.json() raise WechatError(f"wechat getwxacodeunlimit error: {err.get('errcode')} {err.get('errmsg','')}") except WechatError: raise except Exception: raise WechatError("wechat getwxacodeunlimit returned non-image") return resp.content class WxQrSessionStore: """扫码登录会话：scene → {status, token, profile, expires_at}（内存 + TTL）。 - ``start``(scene)：pending。 - ``complete``(scene, token, profile)：done（回调存登录结果）。 - ``poll``(scene)：done 时返回并消费（一次性），否则 pending/expired。 进程重启即失效——生产可在 Redis 持久化（见架构 §8.4）。 """ _TTL = 120 # 会话有效期（秒） def __init__(self) -> None: self._data: dict[str, dict] = {} self._lock = threading.Lock() def start(self, scene: str) -> None: with self._lock: self._data[scene] = { "status": "pending", "token": "", "profile": None, "expires_at": time.monotonic() + self._TTL, } def complete(self, scene: str, token: str, profile: dict) -> None: with self._lock: if scene in self._data: self._data[scene].update( {"status": "done", "token": token, "profile": profile}, ) def poll(self, scene: str) -> dict: with self._lock: rec = self._data.get(scene) if not rec: return {"status": "expired"} if time.monotonic() > rec["expires_at"]: self._data.pop(scene, None) return {"status": "expired"} if rec["status"] == "done": # 消费一次（防重复取），但保留短暂窗口 rec["status"] = "consumed" return {"status": "done", "token": rec["token"], "profile": rec["profile"]} return {"status": rec["status"], "token": "", "profile": None} def new_scene() -> str: return secrets.token_urlsafe(24)