server-core/app/services/wechat.py

# -*- 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


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)