PagerMaid_Plugins_Pyro/restore_pin/main.py

from asyncio import sleep
from typing import Dict, List

from pyrogram import filters
from pyrogram.errors import ChatAdminRequired, UserAdminInvalid, FloodWait
from pyrogram.raw.functions.channels import GetAdminLog, GetMessages
from pyrogram.raw.types import ChannelAdminLogEventsFilter, ChannelAdminLogEventActionUpdatePinned, InputMessageID
from pyrogram.raw.types.channels import AdminLogResults
from pyrogram.raw.types.messages import Messages

from pagermaid.enums import Message
from pagermaid.listener import listener
from pagermaid.services import bot


async def get_admin_log(cid: int) -> AdminLogResults:
    d: AdminLogResults = await bot.invoke(
        GetAdminLog(
            channel=await bot.resolve_peer(cid),
            q="",
            max_id=0,
            min_id=0,
            limit=100,
            events_filter=ChannelAdminLogEventsFilter(pinned=True),
        )
    )
    return d


def get_num_map(events: AdminLogResults) -> Dict[int, List[int]]:
    num_map: Dict[int, List[int]] = {}
    for event in events.events:
        if isinstance(event.action, ChannelAdminLogEventActionUpdatePinned):
            if event.action.message.pinned:
                continue
            num = num_map.get(event.user_id, [])
            num.append(event.action.message.id)
            num_map[event.user_id] = num
    return num_map


async def try_ask_admin(message: Message, num_map: Dict[int, List[int]]) -> int:
    nums = sorted(num_map.keys(), key=lambda x: len(num_map[x]), reverse=True)
    text = "请发送执行误取消置顶操作的管理员 id：\n\n"
    for idx in nums:
        text += f"`{idx}` - 取消 {len(num_map[idx])} 条\n"
    await message.edit(text)
    try:
        async with bot.conversation(
                message.chat.id, filters=filters.user(message.from_user.id)
        ) as conv:
            await sleep(.1)
            res: Message = await conv.get_response()
            await res.safe_delete()
            uid = int(res.text)
            if uid not in num_map:
                raise ValueError
    except ValueError:
        await message.edit("错误：管理员 id 不正确")
        return 0
    return uid


async def pin_one(message: Message, mid: int):
    try:
        await bot.pin_chat_message(message.chat.id, mid, disable_notification=True)
    except ChatAdminRequired:
        return
    except UserAdminInvalid:
        return
    except FloodWait as e:
        await message.edit(f"触发限制，睡眠 {e.value} 秒")
        await sleep(e.value)
        await pin_one(message, mid)


async def get_unpin_messages(cid: int, ids: List[int]) -> List[int]:
    ids = [InputMessageID(id=i) for i in ids]
    r: Messages = await bot.invoke(GetMessages(channel=await bot.resolve_peer(cid), id=ids))
    new_ids = []
    for i in r.messages:
        if not i.pinned:
            new_ids.append(i.id)
    return new_ids


async def try_restore_pin(message: Message, ids: List[int]):
    new_ids = await get_unpin_messages(message.chat.id, ids)
    if not new_ids:
        return await message.edit("没有需要恢复的置顶")
    error = ""
    for idx, i in enumerate(new_ids):
        if (idx + 1) % 5 == 0:
            await message.edit(f"正在恢复第 {idx + 1} 条置顶...")
        try:
            await pin_one(message, i)
        except Exception as e:
            error += f"恢复第 {idx + 1} 条置顶失败 ：{e}\n"
        await sleep(3)
    if error:
        await message.edit(error)
    else:
        await message.edit("已恢复所有消息的置顶")


@listener(
    command="restore_pin",
    description="恢复管理员误取消的置顶",
    groups_only=True,
    need_admin=True,
)
async def restore_pin(message: Message):
    if not message.from_user:
        return
    message = await message.edit("正在获取管理员日志...")
    try:
        events = await get_admin_log(message.chat.id)
    except Exception as e:
        return await message.edit(f"请求管理员日志失败：{e}")
    if not events.events:
        return await message.edit("管理员日志为空，无法恢复")
    num_map = get_num_map(events)
    if not num_map:
        return await message.edit("管理员日志为空，无法恢复")
    admin_id = await try_ask_admin(message, num_map)
    if admin_id not in num_map:
        return
    await message.edit("尝试恢复置顶...")
    await try_restore_pin(message, num_map[admin_id])