From 18b57480a4420e4bc2292762475c5765e946d0b7 Mon Sep 17 00:00:00 2001
From: xtaodada <xtao@xtaolink.cn>
Date: Mon, 11 Nov 2024 16:57:10 +0800
Subject: [PATCH] feat: send message lock

---
 plugins/chat_member_update.py              | 61 +++++++++++++++++-----
 plugins/service.py                         |  2 +-
 requirements.txt                           |  6 +--
 sticker/__init__.py                        | 13 +----
 sticker/__main__.py                        |  3 +-
 sticker/bot.py                             | 15 ++++++
 sticker/config.py                          |  3 +-
 sticker/functions/__init__.py              |  0
 sticker/{ => functions}/service_message.py |  4 +-
 sticker/scheduler.py                       | 22 ++++++++
 sticker/single_utils.py                    | 17 ------
 11 files changed, 95 insertions(+), 51 deletions(-)
 create mode 100644 sticker/bot.py
 create mode 100644 sticker/functions/__init__.py
 rename sticker/{ => functions}/service_message.py (92%)

diff --git a/plugins/chat_member_update.py b/plugins/chat_member_update.py
index 87a4464..224678e 100644
--- a/plugins/chat_member_update.py
+++ b/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.service_message import ServiceMessage
+from sticker.scheduler import add_ban_chat_member_job, add_delete_message_id_job
+from sticker.functions.service_message import ServiceMessage
 from sticker.single_utils import Client, Message
 from sticker import bot, log, LogAction
 
 
+lock_map_lock = asyncio.Lock()
+lock_map: Dict[int, asyncio.Lock] = {}
+
+
+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)
-    try:
-        msg: "Message" = await client.send_message(
-            chat.id, MSG % (user.mention, user.mention)
-        )
-    except Exception:
-        return
-    await log(chat, user, LogAction.REQUEST)
+    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)
diff --git a/plugins/service.py b/plugins/service.py
index f544840..27e497a 100644
--- a/plugins/service.py
+++ b/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
 
 
diff --git a/requirements.txt b/requirements.txt
index 6f1228d..392cf0d 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,8 +1,8 @@
-pyrogram==2.0.106
-TgCrypto>=1.2.3
+git+https://github.com/TeamPGM/pyrogram
+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
diff --git a/sticker/__init__.py b/sticker/__init__.py
index 87d528a..51f9022 100644
--- a/sticker/__init__.py
+++ b/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):
diff --git a/sticker/__main__.py b/sticker/__main__.py
index ba81e35..2143ffd 100644
--- a/sticker/__main__.py
+++ b/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()
 
diff --git a/sticker/bot.py b/sticker/bot.py
new file mode 100644
index 0000000..fbced67
--- /dev/null
+++ b/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"},
+)
diff --git a/sticker/config.py b/sticker/config.py
index 5a8de00..630c4db 100644
--- a/sticker/config.py
+++ b/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
 
 
diff --git a/sticker/functions/__init__.py b/sticker/functions/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/sticker/service_message.py b/sticker/functions/service_message.py
similarity index 92%
rename from sticker/service_message.py
rename to sticker/functions/service_message.py
index a2eb422..f16284f 100644
--- a/sticker/service_message.py
+++ b/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
-        return []
+        return data or []
 
     @staticmethod
     async def try_delete(uid: int, cid: int):
diff --git a/sticker/scheduler.py b/sticker/scheduler.py
index 9b259d7..d7f34be 100644
--- a/sticker/scheduler.py
+++ b/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,
diff --git a/sticker/single_utils.py b/sticker/single_utils.py
index 4580799..cc3576f 100644
--- a/sticker/single_utils.py
+++ b/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]: