73 lines
2.5 KiB
Python
73 lines
2.5 KiB
Python
import json
|
||
import asyncio
|
||
from channels.generic.websocket import AsyncWebsocketConsumer
|
||
from django.core.cache import cache
|
||
|
||
|
||
class QRStatusConsumer(AsyncWebsocketConsumer):
|
||
"""
|
||
扫码登录 WebSocket Consumer.
|
||
前端连接 ws://host/ws/qr-status/{token}/，
|
||
后端在扫码/确认/取消/过期时实时推送状态。
|
||
"""
|
||
|
||
async def connect(self):
|
||
self.qr_token = self.scope['url_route']['kwargs']['token']
|
||
self.group_name = f"qr_{self.qr_token}"
|
||
|
||
# 加入该 token 对应的 group
|
||
await self.channel_layer.group_add(self.group_name, self.channel_name)
|
||
await self.accept()
|
||
|
||
# 检查 token 是否已过期/不存在
|
||
data = cache.get(f"qr_token:{self.qr_token}")
|
||
if not data:
|
||
await self.send(text_data=json.dumps({"status": "expired"}))
|
||
await self.close()
|
||
return
|
||
|
||
# 启动过期倒计时任务
|
||
self._expire_task = asyncio.create_task(self._expire_watcher(data.get("expires_at")))
|
||
|
||
async def disconnect(self, close_code):
|
||
if hasattr(self, 'group_name'):
|
||
await self.channel_layer.group_discard(self.group_name, self.channel_name)
|
||
if hasattr(self, '_expire_task'):
|
||
self._expire_task.cancel()
|
||
|
||
async def receive(self, text_data):
|
||
"""客户端不需要发送消息，忽略即可"""
|
||
pass
|
||
|
||
# ── group 消息处理器：由 views 通过 group_send 触发 ──
|
||
|
||
async def qr_status_update(self, event):
|
||
"""接收扫码/确认/取消状态推送"""
|
||
message = {"status": event["status"]}
|
||
if event.get("username"):
|
||
message["username"] = event["username"]
|
||
if event.get("auth"):
|
||
message["auth"] = event["auth"]
|
||
await self.send(text_data=json.dumps(message))
|
||
|
||
# confirmed 或 cancelled 后关闭连接
|
||
if event["status"] in ("confirmed", "cancelled", "expired"):
|
||
await self.close()
|
||
|
||
# ── 内部方法 ──
|
||
|
||
async def _expire_watcher(self, expires_at):
|
||
"""Token 过期时主动推送 expired 并关闭"""
|
||
if not expires_at:
|
||
return
|
||
import time
|
||
wait = max(expires_at - time.time(), 0) + 1 # 多等 1 秒确保 Redis 已清除
|
||
try:
|
||
await asyncio.sleep(wait)
|
||
except asyncio.CancelledError:
|
||
return
|
||
|
||
# token 已过期
|
||
await self.send(text_data=json.dumps({"status": "expired"}))
|
||
await self.close()
|