"""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)