iShotaBot/modules/banme.py

import contextlib
import re
import random

from datetime import datetime, timedelta
from pyrogram import Client, filters
from pyrogram.enums import ChatMemberStatus
from pyrogram.types import Message, ChatPermissions, CallbackQuery, InlineKeyboardMarkup, InlineKeyboardButton
from init import bot
from scheduler import reply_message


def gen_cancel_button(uid: int):
    return InlineKeyboardMarkup(
        [
            [
                InlineKeyboardButton(text="别口球我！", callback_data=f"banme_cancel_{uid}")
            ]
        ]
    )


@bot.on_message(
    filters.incoming
    & filters.group
    & filters.command(["banme", f"banme@{bot.me.username}"])
)
async def ban_me_command(client: Client, message: Message):
    args = str(message.text).strip()
    # 检查是否有倍数参数
    if multiple_text := re.search(r"^(\d+)倍$", args):
        multiple = int(multiple_text.groups()[0])
    else:
        multiple = 1
    if multiple > 5 or multiple < 1:
        multiple = 1

    # 检查bot和用户身份
    if (
            await client.get_chat_member(message.chat.id, "self")
    ).status != ChatMemberStatus.ADMINISTRATOR:
        await message.reply("Bot非群管理员, 无法执行禁言操作QAQ")
        return
    if not message.from_user:
        # 频道
        await reply_message(message, "你是个频道, 别来凑热闹OvO")
        return

    member = (
        await client.get_chat_member(message.chat.id, message.from_user.id)
    ).status
    if member in [ChatMemberStatus.ADMINISTRATOR, ChatMemberStatus.OWNER]:
        await reply_message(message, "你也是个管理, 别来凑热闹OvO")
        return

    # 随机禁言时间
    random_time = 2 * int(random.gauss(128 * multiple, 640 * multiple // 10))
    act_time = 60 if random_time < 60 else min(random_time, 2591940)
    msg = f"既然你那么想被口球的话, 那我就成全你吧!\n送你一份{act_time // 60}分{act_time % 60}秒禁言套餐哦, 谢谢惠顾~"

    await client.restrict_chat_member(
        message.chat.id,
        message.from_user.id,
        ChatPermissions(),
        datetime.now() + timedelta(seconds=act_time),
    )
    await reply_message(message, msg, reply_markup=gen_cancel_button(message.from_user.id))


@bot.on_callback_query(
    filters.regex(r"^banme_cancel_(\d+)$")
)
async def ban_me_cancel(client: Client, callback_query: CallbackQuery):
    if not callback_query.from_user:
        return
    uid = int(callback_query.data.split("_")[-1])
    if callback_query.from_user.id != uid:
        await callback_query.answer("这不是属于你的按钮 ~")
        return
    await callback_query.answer("已撤销，请等待 60s 后发言")
    with contextlib.suppress(Exception):
        await callback_query.message.delete()
    await client.restrict_chat_member(
        callback_query.message.chat.id,
        callback_query.from_user.id,
        ChatPermissions(),
        datetime.now() + timedelta(seconds=60),
    )