feat: send message lock

This commit is contained in:
xtaodada 2024-11-11 16:57:10 +08:00
parent f1f10b454a
commit 18b57480a4
Signed by: xtaodada
GPG Key ID: 4CBB3F4FA8C85659
11 changed files with 95 additions and 51 deletions

View File

@ -1,4 +1,8 @@
import asyncio
import contextlib import contextlib
from asyncio import sleep
from time import time
from typing import Dict
from cashews import cache from cashews import cache
from pyrogram import filters from pyrogram import filters
@ -7,27 +11,58 @@ from pyrogram.types import ChatMemberUpdated
from pyromod.utils.errors import TimeoutConversationError from pyromod.utils.errors import TimeoutConversationError
from sticker.languages import MSG_PUBLIC, ADMIN_MSG, MSG, VERIFY_TIME 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.single_utils import Client, Message
from sticker import bot, log, LogAction 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): async def start_verify(client: "Client", chat, user):
key = f"wait:{chat.id}:{user.id}" key = f"wait:{chat.id}:{user.id}"
key2 = f"msg:{chat.id}:{user.id}"
await cache.set(key, "True", expire=VERIFY_TIME + 5) await cache.set(key, "True", expire=VERIFY_TIME + 5)
try: client.loop.create_task(send_message(client, chat, user))
msg: "Message" = await client.send_message(
chat.id, MSG % (user.mention, user.mention)
)
except Exception:
return
await log(chat, user, LogAction.REQUEST)
try: try:
msg_: "Message" = await client.listen( 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) await msg_.delay_delete(1)
if not msg_.sticker: if not msg_.sticker:
add_ban_chat_member_job(chat.id, user.id) 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 cache.delete(key)
await log(chat, user, LogAction.ACCEPT) await log(chat, user, LogAction.ACCEPT)
except TimeoutConversationError: 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) add_ban_chat_member_job(chat.id, user.id)
await log(chat, user, LogAction.FAIL_TIMEOUT) await log(chat, user, LogAction.FAIL_TIMEOUT)
await ServiceMessage.try_delete(user.id, chat.id) await ServiceMessage.try_delete(user.id, chat.id)

View File

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

View File

@ -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 PyYAML==6.0.1
coloredlogs>=15.0.1 coloredlogs>=15.0.1
sqlitedict==2.1.0 sqlitedict==2.1.0
apscheduler==3.10.4 apscheduler==3.10.4
cashews==6.2.0 cashews==7.2.0
pytz pytz

View File

@ -7,9 +7,8 @@ from datetime import datetime, timezone
from logging import getLogger, StreamHandler, CRITICAL, INFO, basicConfig, DEBUG from logging import getLogger, StreamHandler, CRITICAL, INFO, basicConfig, DEBUG
from coloredlogs import ColoredFormatter from coloredlogs import ColoredFormatter
import pyromod.listen
from pyrogram import Client
from sticker.bot import bot
from sticker.config import Config from sticker.config import Config
from sticker.scheduler import scheduler from sticker.scheduler import scheduler
@ -44,16 +43,6 @@ with contextlib.suppress(ImportError):
if not scheduler.running: if not scheduler.running:
scheduler.start() 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): class LogAction(str, Enum):

View File

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

15
sticker/bot.py Normal file
View File

@ -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"},
)

View File

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

View File

View File

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

View File

@ -18,6 +18,15 @@ async def delete_message(message: Message) -> bool:
return False 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): async def decline_request(chat_join_request: ChatJoinRequest):
with contextlib.suppress(Exception): with contextlib.suppress(Exception):
await chat_join_request.decline() await chat_join_request.decline()
@ -36,6 +45,19 @@ async def ban_chat_member(chat_id: int, user_id: int):
return False 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): def add_delete_message_job(message: Message, delete_seconds: int = 60):
scheduler.add_job( scheduler.add_job(
delete_message, delete_message,

View File

@ -1,25 +1,8 @@
import contextlib
from os import sep, remove, mkdir
from os.path import exists
from typing import Optional from typing import Optional
from pyrogram import Client from pyrogram import Client
from pyrogram.types import Message 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 class Client(Client): # noqa
async def listen(self, chat_id, filters=None, timeout=None) -> Optional[Message]: async def listen(self, chat_id, filters=None, timeout=None) -> Optional[Message]: