118 lines
4.4 KiB
Python
118 lines
4.4 KiB
Python
import asyncio
|
|
import time
|
|
from typing import Dict, List, Tuple, Optional, Callable
|
|
from loguru import logger
|
|
|
|
|
|
class MessageRevokeManager:
|
|
"""消息撤回管理器，用于管理需要延时撤回的消息"""
|
|
|
|
def __init__(self, default_delay: float = 90.0):
|
|
"""
|
|
初始化消息撤回管理器
|
|
|
|
Args:
|
|
default_delay: 默认撤回延时时间（秒），默认为90秒
|
|
"""
|
|
self.revoke_queue: List[Tuple[float, str, int, int, int]] = [] # (撤回时间, wxid, client_msg_id, create_time, new_msg_id)
|
|
self.default_delay = default_delay
|
|
self.running = False
|
|
self.task = None
|
|
self.logging = logger
|
|
self.revoke_callback = None
|
|
|
|
def add_message_to_revoke(self, wxid: str, client_msg_id: int, create_time: int, new_msg_id: int,
|
|
delay: Optional[float] = None) -> None:
|
|
"""
|
|
添加消息到撤回队列
|
|
|
|
Args:
|
|
wxid: 接收人wxid
|
|
client_msg_id: 消息ID
|
|
create_time: 消息创建时间
|
|
new_msg_id: 新消息ID
|
|
delay: 延迟撤回时间（秒），如果为None则使用默认值
|
|
"""
|
|
if delay is None:
|
|
delay = self.default_delay
|
|
|
|
revoke_time = time.time() + delay
|
|
self.revoke_queue.append((revoke_time, wxid, client_msg_id, create_time, new_msg_id))
|
|
self.revoke_queue.sort(key=lambda x: x[0]) # 按撤回时间排序
|
|
|
|
# 如果撤回任务未运行，则启动
|
|
if not self.running:
|
|
self.start()
|
|
|
|
def start(self) -> None:
|
|
"""启动消息撤回任务"""
|
|
if self.running:
|
|
return
|
|
|
|
self.running = True
|
|
self.task = asyncio.create_task(self._process_revoke_queue())
|
|
self.logging.info("消息自动撤回任务已启动")
|
|
|
|
def stop(self) -> None:
|
|
"""停止消息撤回任务"""
|
|
if not self.running:
|
|
return
|
|
|
|
self.running = False
|
|
if self.task:
|
|
self.task.cancel()
|
|
self.logging.info("消息自动撤回任务已停止")
|
|
|
|
def set_revoke_callback(self, callback: Callable) -> None:
|
|
"""
|
|
设置撤回回调函数
|
|
|
|
Args:
|
|
callback: 撤回消息的回调函数，接受wxid, client_msg_id, create_time, new_msg_id参数
|
|
"""
|
|
self.revoke_callback = callback
|
|
|
|
def set_default_delay(self, delay: float) -> None:
|
|
"""
|
|
设置默认撤回延时
|
|
|
|
Args:
|
|
delay: 默认撤回延时（秒）
|
|
"""
|
|
self.default_delay = delay
|
|
|
|
async def _process_revoke_queue(self) -> None:
|
|
"""处理撤回队列的异步任务"""
|
|
try:
|
|
while self.running:
|
|
now = time.time()
|
|
|
|
# 检查是否有需要撤回的消息
|
|
while self.revoke_queue and self.revoke_queue[0][0] <= now:
|
|
revoke_time, wxid, client_msg_id, create_time, new_msg_id = self.revoke_queue.pop(0)
|
|
try:
|
|
# 通过回调函数调用实际的撤回方法
|
|
if self.revoke_callback and callable(self.revoke_callback):
|
|
await self.revoke_callback(wxid, client_msg_id, create_time, new_msg_id)
|
|
self.logging.info(f"已自动撤回消息: wxid={wxid}, msg_id={client_msg_id}")
|
|
except Exception as e:
|
|
self.logging.error(f"自动撤回消息失败: {e}")
|
|
|
|
# 如果队列为空，等待一段时间
|
|
if not self.revoke_queue:
|
|
self.running = False
|
|
self.logging.info("撤回队列为空，撤回任务已停止")
|
|
break
|
|
|
|
# 计算下一个需要撤回的消息的等待时间
|
|
next_revoke_time = self.revoke_queue[0][0]
|
|
wait_time = max(0.1, next_revoke_time - time.time())
|
|
|
|
# 等待到下一个撤回时间
|
|
await asyncio.sleep(wait_time)
|
|
|
|
except asyncio.CancelledError:
|
|
self.logging.info("消息撤回任务被取消")
|
|
except Exception as e:
|
|
self.logging.error(f"消息撤回任务异常: {e}")
|
|
self.running = False |