feat: send message lock

plugins/chat_member_update.py

@@ -1,4 +1,8 @@
+import asyncio
 import contextlib
+from asyncio import sleep
+from time import time
+from typing import Dict
 
 from cashews import cache
 from pyrogram import filters
@@ -7,27 +11,58 @@ from pyrogram.types import ChatMemberUpdated
 
 from pyromod.utils.errors import TimeoutConversationError
 from sticker.languages import MSG_PUBLIC, ADMIN_MSG, MSG, VERIFY_TIME
-from sticker.scheduler import add_ban_chat_member_job
+from sticker.scheduler import add_ban_chat_member_job, add_delete_message_id_job
-from sticker.service_message import ServiceMessage
+from sticker.functions.service_message import ServiceMessage
 from sticker.single_utils import Client, Message
 from sticker import bot, log, LogAction
 
 
-async def start_verify(client: "Client", chat, user):
+lock_map_lock = asyncio.Lock()
-    key = f"wait:{chat.id}:{user.id}"
+lock_map: Dict[int, asyncio.Lock] = {}
-    await cache.set(key, "True", expire=VERIFY_TIME + 5)
+
+
+async def get_lock(chat_id: int):
+    async with lock_map_lock:
+        lock = lock_map.get(chat_id)
+        if not lock:
+            lock = asyncio.Lock()
+            lock_map[chat_id] = lock
+    return lock
+
+
+async def send_message(client: "Client", chat, user):
+    n_time = time()
+    lock = await get_lock(chat.id)
+    async with lock:
+        if time() - n_time > 30:
+            # 认为此任务已过期
+            return
         try:
+            key = f"msg:{chat.id}:{user.id}"
             msg: "Message" = await client.send_message(
                 chat.id, MSG % (user.mention, user.mention)
             )
+            await msg.delay_delete(VERIFY_TIME + 5)
+            await cache.set(key, msg.id, expire=VERIFY_TIME + 5)
         except Exception:
             return
         await log(chat, user, LogAction.REQUEST)
+
+
+async def start_verify(client: "Client", chat, user):
+    key = f"wait:{chat.id}:{user.id}"
+    key2 = f"msg:{chat.id}:{user.id}"
+    await cache.set(key, "True", expire=VERIFY_TIME + 5)
+    client.loop.create_task(send_message(client, chat, user))
     try:
         msg_: "Message" = await client.listen(
-            chat.id, filters=filters.user(user.id), timeout=VERIFY_TIME
+            chat.id,
+            filters=filters.user(user.id) & ~filters.service,
+            timeout=VERIFY_TIME,
         )
-        await msg.delay_delete(1)
+        msg = await cache.get(key2)
+        if msg:
+            add_delete_message_id_job(chat.id, msg)
         await msg_.delay_delete(1)
         if not msg_.sticker:
             add_ban_chat_member_job(chat.id, user.id)
@@ -37,7 +72,9 @@ async def start_verify(client: "Client", chat, user):
             await cache.delete(key)
             await log(chat, user, LogAction.ACCEPT)
     except TimeoutConversationError:
-        await msg.delay_delete(1)
+        msg = await cache.get(key2)
+        if msg:
+            add_delete_message_id_job(chat.id, msg)
         add_ban_chat_member_job(chat.id, user.id)
         await log(chat, user, LogAction.FAIL_TIMEOUT)
         await ServiceMessage.try_delete(user.id, chat.id)

plugins/service.py

@@ -2,7 +2,7 @@ from pyrogram import filters
 from pyrogram.enums import MessageServiceType
 
 from sticker import bot
-from sticker.service_message import ServiceMessage
+from sticker.functions.service_message import ServiceMessage
 from sticker.single_utils import Client, Message
 
 

requirements.txt

@@ -1,8 +1,8 @@
-pyrogram==2.0.106
+git+https://github.com/TeamPGM/pyrogram
-TgCrypto>=1.2.3
+PyroTgCrypto==1.2.6a0
 PyYAML==6.0.1
 coloredlogs>=15.0.1
 sqlitedict==2.1.0
 apscheduler==3.10.4
-cashews==6.2.0
+cashews==7.2.0
 pytz

sticker/__init__.py

@@ -7,9 +7,8 @@ from datetime import datetime, timezone
 from logging import getLogger, StreamHandler, CRITICAL, INFO, basicConfig, DEBUG
 
 from coloredlogs import ColoredFormatter
-import pyromod.listen
-from pyrogram import Client
 
+from sticker.bot import bot
 from sticker.config import Config
 from sticker.scheduler import scheduler
 
@@ -44,16 +43,6 @@ with contextlib.suppress(ImportError):
 
 if not scheduler.running:
     scheduler.start()
-bot = Client(
-    "sticker",
-    bot_token=Config.BOT_TOKEN,
-    session_string=Config.STRING_SESSION,
-    api_id=Config.API_ID,
-    api_hash=Config.API_HASH,
-    ipv6=Config.IPV6,
-    proxy=Config.PROXY,
-    plugins={"root": "plugins"},
-)
 
 
 class LogAction(str, Enum):

sticker/__main__.py

@@ -5,7 +5,8 @@ from sticker import bot, logs
 
 async def main():
     await bot.start()
-    logs.info("bot started.")
+    me = await bot.get_me()
+    logs.info(f"bot @{me.username} started.")
     await idle()
     await bot.stop()
 

sticker/bot.py

@@ -0,0 +1,15 @@
+import pyromod.listen
+from pyrogram import Client
+
+from sticker.config import Config
+
+bot = Client(
+    "sticker",
+    bot_token=Config.BOT_TOKEN,
+    session_string=Config.STRING_SESSION,
+    api_id=Config.API_ID,
+    api_hash=Config.API_HASH,
+    ipv6=Config.IPV6,
+    proxy=Config.PROXY,
+    plugins={"root": "plugins"},
+)

sticker/config.py

@@ -1,7 +1,6 @@
 import os
-from json import load as load_json
 import sys
-from yaml import load, FullLoader, safe_load
+from yaml import load, FullLoader
 from shutil import copyfile
 
 

sticker/functions/service_message.py

@@ -16,9 +16,7 @@ class ServiceMessage:
     @staticmethod
     async def get_cache(uid: int, cid: int) -> List[int]:
         data = await cache.get(f"service_message:{uid}:{cid}")
-        if data:
+        return data or []
-            return data
-        return []
 
     @staticmethod
     async def try_delete(uid: int, cid: int):

sticker/scheduler.py

@@ -18,6 +18,15 @@ async def delete_message(message: Message) -> bool:
     return False
 
 
+async def delete_message_id(chat_id: int, message_id: int) -> bool:
+    with contextlib.suppress(Exception):
+        from sticker.bot import bot
+
+        await bot.delete_messages(chat_id, message_id)
+        return True
+    return False
+
+
 async def decline_request(chat_join_request: ChatJoinRequest):
     with contextlib.suppress(Exception):
         await chat_join_request.decline()
@@ -36,6 +45,19 @@ async def ban_chat_member(chat_id: int, user_id: int):
     return False
 
 
+def add_delete_message_id_job(chat_id: int, message_id: int, delete_seconds: int = 60):
+    scheduler.add_job(
+        delete_message_id,
+        "date",
+        id=f"{chat_id}|{message_id}|delete_message",
+        name=f"{chat_id}|{message_id}|delete_message",
+        args=[chat_id, message_id],
+        run_date=datetime.datetime.now(pytz.timezone("Asia/Shanghai"))
+        + datetime.timedelta(seconds=delete_seconds),
+        replace_existing=True,
+    )
+
+
 def add_delete_message_job(message: Message, delete_seconds: int = 60):
     scheduler.add_job(
         delete_message,

sticker/single_utils.py

@@ -1,25 +1,8 @@
-import contextlib
-from os import sep, remove, mkdir
-from os.path import exists
 from typing import Optional
 
 from pyrogram import Client
 from pyrogram.types import Message
 
-from pyromod.utils.errors import TimeoutConversationError, ListenerCanceled
-
-from sqlitedict import SqliteDict
-
-# init folders
-if not exists("data"):
-    mkdir("data")
-sqlite = SqliteDict(f"data{sep}data.sqlite", autocommit=True)
-
-
-def safe_remove(name: str) -> None:
-    with contextlib.suppress(FileNotFoundError):
-        remove(name)
-
 
 class Client(Client):  # noqa
     async def listen(self, chat_id, filters=None, timeout=None) -> Optional[Message]: