Pine
fa06678e92
- PINEAGENTS_ENV=prod|dev 判定运行环境; /auth/status 回显 env + wxacode_env 以便确认 - PINEAGENTS_WXACODE_ENV=release|trial|develop 决定小程序码指向(默认 prod→release, dev→trial=体验版) - get_wxacode 用该 env_version(支持体验版测试)
226 lines
8.9 KiB
Python
226 lines
8.9 KiB
Python
# -*- 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)
|
||
|