shanxia2.0/backend/app/services/wx_qr_session.py

"""web-admin 扫码登录会话：短码映射 + 状态机；优先 Redis，否则进程内存。"""
from __future__ import annotations

import json
import secrets
import time
import uuid
from typing import Any, Dict, Optional

from app.core.config import settings

_PREFIX = "wxqr:sess:"
_SHORT_PREFIX = "wxqr:short:"
_memory: Dict[str, Dict[str, Any]] = {}
_short_map: Dict[str, str] = {}


def _memory_clear_for_tests() -> None:
    _memory.clear()
    _short_map.clear()


def _ttl() -> int:
    return int(getattr(settings, "WX_QR_SESSION_TTL_SEC", 180) or 180)


def _redis():
    """复用 auth 侧 Redis 客户端模式；无 REDIS_HOST 返回 None。"""
    host = (settings.REDIS_HOST or "").strip()
    if not host:
        return None
    try:
        import redis

        return redis.Redis(
            host=host,
            port=int(settings.REDIS_PORT or 6379),
            db=int(settings.REDIS_DB or 0),
            password=settings.REDIS_PASSWORD or None,
            decode_responses=True,
            socket_timeout=2,
        )
    except Exception:
        return None


def create_session(*, ttl_sec: Optional[int] = None) -> tuple:
    ttl = int(ttl_sec if ttl_sec is not None else _ttl())
    session_id = str(uuid.uuid4())
    short = secrets.token_urlsafe(16)[:32]
    now = int(time.time())
    data = {
        "sessionId": session_id,
        "shortCode": short,
        "status": "pending",
        "message": "",
        "createdAt": now,
        "expiresAt": now + ttl,
        "tokens": None,
    }
    r = _redis()
    if r:
        pipe = r.pipeline()
        pipe.setex(_PREFIX + session_id, ttl, json.dumps(data, ensure_ascii=False))
        pipe.setex(_SHORT_PREFIX + short, ttl, session_id)
        pipe.execute()
    else:
        _memory[session_id] = data
        _short_map[short] = session_id
    return session_id, short


def resolve_session_id(short_or_id: str) -> Optional[str]:
    key = (short_or_id or "").strip()
    if not key:
        return None
    r = _redis()
    if r:
        sid = r.get(_SHORT_PREFIX + key)
        if sid:
            return sid
        if r.exists(_PREFIX + key):
            return key
        return None
    if key in _short_map:
        return _short_map[key]
    if key in _memory:
        return key
    return None


def get_session(session_id: str) -> Optional[Dict[str, Any]]:
    sid = (session_id or "").strip()
    if not sid:
        return None
    r = _redis()
    if r:
        raw = r.get(_PREFIX + sid)
        if not raw:
            return None
        data = json.loads(raw)
    else:
        data = _memory.get(sid)
        if not data:
            return None
        data = dict(data)
    if int(time.time()) >= int(data.get("expiresAt") or 0):
        data["status"] = "expired"
        return data
    return data


def _save(data: Dict[str, Any]) -> None:
    sid = data["sessionId"]
    ttl = max(1, int(data["expiresAt"]) - int(time.time()))
    r = _redis()
    if r:
        r.setex(_PREFIX + sid, ttl, json.dumps(data, ensure_ascii=False))
    else:
        _memory[sid] = data


def mark_scanned(session_id: str) -> bool:
    data = get_session(session_id)
    if not data or data["status"] in ("expired", "success", "consumed", "failed"):
        return False
    if data["status"] == "pending":
        data["status"] = "scanned"
        _save(data)
    return True


def mark_success(session_id: str, tokens: Dict[str, Any]) -> bool:
    data = get_session(session_id)
    if not data or data["status"] in ("expired", "consumed", "failed"):
        return False
    data["status"] = "success"
    data["tokens"] = tokens
    data["message"] = ""
    _save(data)
    return True


def mark_failed(session_id: str, message: str) -> bool:
    data = get_session(session_id)
    if not data or data["status"] in ("expired", "consumed", "success"):
        return False
    data["status"] = "failed"
    data["message"] = message or "登录失败"
    data["tokens"] = None
    _save(data)
    return True


def consume_success(session_id: str) -> Optional[Dict[str, Any]]:
    data = get_session(session_id)
    if not data or data["status"] != "success":
        return None
    tokens = data.get("tokens")
    data["status"] = "consumed"
    data["tokens"] = None
    _save(data)
    return tokens


def invalidate(session_id: str) -> None:
    data = get_session(session_id)
    if not data:
        return
    data["status"] = "expired"
    data["tokens"] = None
    _save(data)