From 56a3311c4582ffee63a7b50d84ffa81a9ee87ccb Mon Sep 17 00:00:00 2001
From: wuzilong <704541015@qq.com>
Date: Fri, 18 Sep 2026 19:55:00 +0800
Subject: [PATCH] feat(auth): add wx-qr session state machine for web login.

Co-authored-by: Cursor <cursoragent@cursor.com>
---
 backend/app/services/wx_qr_session.py | 172 ++++++++++++++++++++++++++
 backend/tests/test_wx_qr_session.py   |  53 ++++++++
 2 files changed, 225 insertions(+)
 create mode 100644 backend/app/services/wx_qr_session.py
 create mode 100644 backend/tests/test_wx_qr_session.py

diff --git a/backend/app/services/wx_qr_session.py b/backend/app/services/wx_qr_session.py
new file mode 100644
index 000000000..b245cce00
--- /dev/null
+++ b/backend/app/services/wx_qr_session.py
@@ -0,0 +1,172 @@
+"""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)
diff --git a/backend/tests/test_wx_qr_session.py b/backend/tests/test_wx_qr_session.py
new file mode 100644
index 000000000..52c55f53d
--- /dev/null
+++ b/backend/tests/test_wx_qr_session.py
@@ -0,0 +1,53 @@
+"""wx_qr_session 状态机与一次性消费。"""
+import time
+
+from app.services import wx_qr_session as s
+
+
+def setup_function():
+    s._memory_clear_for_tests()
+
+
+def test_create_pending_and_get():
+    sid, short = s.create_session(ttl_sec=60)
+    assert sid and short and len(short) <= 32
+    data = s.get_session(sid)
+    assert data["status"] == "pending"
+    assert s.resolve_session_id(short) == sid
+
+
+def test_mark_scanned():
+    sid, _ = s.create_session(ttl_sec=60)
+    assert s.mark_scanned(sid) is True
+    assert s.get_session(sid)["status"] == "scanned"
+
+
+def test_mark_success_consume_once():
+    sid, _ = s.create_session(ttl_sec=60)
+    payload = {"token": "t", "refreshToken": "r", "expiresIn": 1, "userInfo": {}}
+    assert s.mark_success(sid, payload) is True
+    first = s.consume_success(sid)
+    assert first["token"] == "t"
+    assert s.consume_success(sid) is None
+    assert s.get_session(sid)["status"] == "consumed"
+
+
+def test_mark_failed():
+    sid, _ = s.create_session(ttl_sec=60)
+    assert s.mark_failed(sid, "未找到员工") is True
+    assert s.get_session(sid)["status"] == "failed"
+    assert s.get_session(sid)["message"] == "未找到员工"
+
+
+def test_expire_old_session():
+    sid, _ = s.create_session(ttl_sec=1)
+    time.sleep(1.1)
+    data = s.get_session(sid)
+    assert data is None or data["status"] == "expired"
+
+
+def test_invalidate():
+    sid, _ = s.create_session(ttl_sec=60)
+    s.invalidate(sid)
+    data = s.get_session(sid)
+    assert data is None or data["status"] == "expired"