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


# ---------------------------------------------------------------------------
# 小程序码（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  # 会话有效期（秒）
    TTL_SECONDS = _TTL  # 对外下发（前端二维码倒计时）

    def __init__(self) -> None:
        self._data: dict[str, dict] = {}
        self._lock = threading.Lock()

    def start(self, scene: str, bind_user_id: str | None = None) -> None:
        with self._lock:
            self._data[scene] = {
                "status": "pending", "token": "", "profile": None,
                "bind_user_id": bind_user_id,  # 绑定态：要把小程序身份绑到哪个平台账号
                "message": "",
                "expires_at": time.monotonic() + self._TTL,
            }

    def complete(self, scene: str, token: str = "", profile: dict | None = None, message: str = "") -> None:
        with self._lock:
            if scene in self._data:
                self._data[scene].update(
                    {"status": "done", "token": token, "profile": profile, "message": message},
                )

    def bind_user(self, scene: str) -> str | None:
        """返回绑定态会话要绑定的平台账号 id（非绑定态返回 None）。"""
        with self._lock:
            rec = self._data.get(scene)
            return (rec or {}).get("bind_user_id") if isinstance(rec, dict) else None

    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"], "message": rec["message"]}
            return {"status": rec["status"], "token": "", "profile": None, "message": ""}


def new_scene() -> str:
    return secrets.token_urlsafe(24)