TelegramBotAdmin/wbb/__main__.py

"""
MIT License

Copyright (c) 2024 TheHamkerCat

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
"""
import asyncio
import importlib
import os
import re
import time
from contextlib import closing, suppress

import psutil
from pyrogram import StopPropagation, filters, idle
from pyrogram.enums import ChatType, ParseMode
from pyrogram.types import BotCommand, InlineKeyboardButton, InlineKeyboardMarkup
from uvloop import install

from wbb import (
    BOT_NAME,
    BOT_PERMISSIONS,
    BOT_USERNAME,
    LOG_GROUP_ID,
    USERBOT_CONNECTED,
    USERBOT_NAME,
    aiohttpsession,
    app,
    app2,
    bot_start_time,
    log,
)
from wbb.core.keyboard import ikb
from wbb.modules import ALL_MODULES
from wbb.services.bot_permissions import (
    GROUP_MANAGEMENT_PERMISSIONS,
    PRIVATE_MANAGEMENT_PERMISSIONS,
    TELEGRAM_COMMAND_PERMISSIONS,
    has_any_permission,
    has_permission,
)
from wbb.utils import paginate_modules
from wbb.utils.constants import MARKDOWN
from wbb.utils.dbfunctions import clean_restart_stage, get_rules
from wbb.utils.functions import extract_text_and_keyb

loop = asyncio.get_event_loop()

HELPABLE = {}
HELP_MODULE_KEYS = {
    "admin",
    "admin_misc",
    "admin_panel",
    "antiservice",
    "blacklist",
    "chatbot",
    "feds",
    "filters",
    "flood",
    "giveaway",
    "greetings",
    "karma",
    "locks",
    "notes",
    "points",
    "rules",
}

BOT_COMMANDS = [
    BotCommand("start", "打开机器人首页"),
    BotCommand("help", "查看中文功能帮助"),
    BotCommand("manage", "打开私聊群管理菜单"),
    BotCommand("cancel", "取消当前管理流程"),
    BotCommand("points", "查看积分余额"),
    BotCommand("checkin", "领取每日签到积分"),
    BotCommand("points_rank", "查看群积分排行榜"),
    BotCommand("points_history", "查看最近积分流水"),
    BotCommand("giveaway", "创建群抽奖"),
    BotCommand("gjoin", "报名参加抽奖"),
    BotCommand("glist", "查看进行中的抽奖"),
    BotCommand("gparticipants", "查看抽奖参与者"),
    BotCommand("gend", "立即开奖"),
    BotCommand("gcancel", "取消抽奖并退款"),
    BotCommand("greroll", "重新抽取中奖者"),
    BotCommand("rules", "查看本群群规"),
    BotCommand("directory", "打开师生与附近菜单"),
    BotCommand("nearby", "查找附近老师"),
    BotCommand("teacher_rank", "查看老师榜单"),
    BotCommand("teacher_apply", "申请成为老师"),
    BotCommand("teacher_online", "切换老师为上线"),
    BotCommand("teacher_offline", "切换老师为下线"),
    BotCommand("teacher_list", "老师自助上榜"),
    BotCommand("teacher_unlist", "老师自助下榜"),
]
BOT_COMMAND_PERMISSIONS = {
    "manage": "__group_any__",
    "cancel": "__any__",
    "points": "points.manage",
    "checkin": "points.manage",
    "points_rank": "points.manage",
    "points_history": "points.manage",
    "giveaway": "giveaways.manage",
    "gjoin": "giveaways.manage",
    "glist": "giveaways.manage",
    "gparticipants": "giveaways.manage",
    "gend": "giveaways.manage",
    "gcancel": "giveaways.manage",
    "greroll": "giveaways.manage",
    "rules": "chat.rules",
    "directory": "teacher_directory.manage",
    "nearby": "teacher_directory.manage",
    "teacher_rank": "teacher_directory.manage",
    "teacher_apply": "teacher_directory.manage",
    "teacher_online": "teacher_directory.manage",
    "teacher_offline": "teacher_directory.manage",
    "teacher_list": "teacher_directory.manage",
    "teacher_unlist": "teacher_directory.manage",
}


def _visible_bot_commands() -> list[BotCommand]:
    visible = []
    for command in BOT_COMMANDS:
        required = BOT_COMMAND_PERMISSIONS.get(command.command)
        if required == "__group_any__" and not has_any_permission(
            GROUP_MANAGEMENT_PERMISSIONS, BOT_PERMISSIONS
        ):
            continue
        if required == "__any__" and not has_any_permission(
            PRIVATE_MANAGEMENT_PERMISSIONS, BOT_PERMISSIONS
        ):
            continue
        if required and required not in {"__any__", "__group_any__"} and not has_permission(
            required, BOT_PERMISSIONS
        ):
            continue
        visible.append(command)
    return visible


@app.on_message(
    filters.command(list(TELEGRAM_COMMAND_PERMISSIONS)),
    group=-1000,
)
async def bot_role_command_guard(_, message):
    command = (
        str(message.command[0]).lower().split("@", 1)[0]
        if message.command
        else ""
    )
    required = TELEGRAM_COMMAND_PERMISSIONS.get(command)
    if required and not has_permission(required, BOT_PERMISSIONS):
        raise StopPropagation


@app.on_message(
    filters.command(["admins", "admin"], prefixes="@"),
    group=-1000,
)
async def bot_role_at_command_guard(_, __):
    if not has_permission("chat.members", BOT_PERMISSIONS):
        raise StopPropagation


@app.on_callback_query(filters.regex(r"^unwarn_"), group=-1000)
async def bot_role_callback_guard(_, __):
    if not has_permission("chat.members", BOT_PERMISSIONS):
        raise StopPropagation


async def bot_sys_stats() -> str:
    uptime = max(0, int(time.time() - bot_start_time))
    days, remainder = divmod(uptime, 86400)
    hours, remainder = divmod(remainder, 3600)
    minutes, seconds = divmod(remainder, 60)
    duration = f"{days} 天 {hours} 小时 {minutes} 分 {seconds} 秒"
    process = psutil.Process(os.getpid())
    return (
        f"机器人：{BOT_NAME}\n"
        f"运行时长：{duration}\n"
        f"进程内存：{round(process.memory_info().rss / 1024 ** 2)} MB\n"
        f"处理器：{psutil.cpu_percent()}%\n"
        f"系统内存：{psutil.virtual_memory().percent}%\n"
        f"磁盘：{psutil.disk_usage('/').percent}%"
    )


async def start_bot():
    global HELPABLE

    for module in ALL_MODULES:
        imported_module = importlib.import_module("wbb.modules." + module)
        if (
            module in HELP_MODULE_KEYS
            and hasattr(imported_module, "__MODULE__")
            and imported_module.__MODULE__
        ):
            if (
                hasattr(imported_module, "__HELP__")
                and imported_module.__HELP__
            ):
                HELPABLE[module] = imported_module
    bot_modules = ""
    j = 1
    for i in ALL_MODULES:
        if j == 4:
            bot_modules += "|{:<15}|\n".format(i)
            j = 0
        else:
            bot_modules += "|{:<15}".format(i)
        j += 1
    print("+===============================================================+")
    print("|                              WBB                              |")
    print("+===============+===============+===============+===============+")
    print(bot_modules)
    print("+===============+===============+===============+===============+")
    log.info(f"机器人已启动：{BOT_NAME}")
    if USERBOT_CONNECTED:
        log.info(f"用户客户端已启动：{USERBOT_NAME}")

    try:
        await app.set_bot_commands(_visible_bot_commands())
        log.info("已注册中文 Telegram 命令菜单")
    except Exception as exc:
        log.error(f"注册 Telegram 命令菜单失败：{exc}")

    restart_data = await clean_restart_stage()

    try:
        log.info("正在发送上线状态")
        if restart_data:
            await app.edit_message_text(
                restart_data["chat_id"],
                restart_data["message_id"],
                "**机器人已成功重启。**",
            )

        else:
            await app.send_message(LOG_GROUP_ID, "机器人已启动。")
    except Exception:
        pass

    await idle()

    await aiohttpsession.close()
    log.info("正在停止客户端")
    await app.stop()
    if USERBOT_CONNECTED:
        await app2.stop()
    log.info("正在取消异步任务")
    for task in asyncio.all_tasks():
        task.cancel()
    log.info("机器人已停止")


home_actions = [
    InlineKeyboardButton(text="功能帮助", callback_data="bot_commands"),
]
if has_any_permission(GROUP_MANAGEMENT_PERMISSIONS, BOT_PERMISSIONS):
    home_actions.append(
        InlineKeyboardButton(text="群组管理", callback_data="manage_hint")
    )
if has_permission("teacher_directory.manage", BOT_PERMISSIONS):
    home_actions.append(
        InlineKeyboardButton(
            text="师生与附近",
            url=f"https://t.me/{BOT_USERNAME}?start=directory",
        )
    )

home_keyboard_pm = InlineKeyboardMarkup(
    [
        home_actions,
        [
            InlineKeyboardButton(
                text="运行状态",
                callback_data="stats_callback",
            )
        ],
        [
            InlineKeyboardButton(
                text="添加机器人到群组",
                url=f"http://t.me/{BOT_USERNAME}?startgroup=new",
            )
        ],
    ]
)

assigned_features = [
    label
    for permission, label in (
        ("chat.profile", "群资料"),
        ("chat.members", "成员管理"),
        ("chat.permissions", "群权限"),
        ("chat.announcements", "公告"),
        ("chat.invites", "邀请链接"),
        ("chat.rules", "群规"),
        ("chat.notes", "群笔记"),
        ("automation.manage", "自动化与风控"),
        ("points.manage", "积分"),
        ("giveaways.manage", "抽奖"),
        ("karma.manage", "声望"),
        ("teacher_directory.manage", "师生与附近"),
    )
    if has_permission(permission, BOT_PERMISSIONS)
]
home_text_pm = (
    f"你好，我是 {BOT_NAME}。\n\n"
    + (
        f"我当前负责：{'、'.join(assigned_features)}。"
        if assigned_features
        else "我当前尚未分配管理职责。"
    )
    + (
        "群管理员可发送 /manage 打开私聊管理菜单。"
        if has_any_permission(GROUP_MANAGEMENT_PERMISSIONS, BOT_PERMISSIONS)
        else ""
    )
)

keyboard = InlineKeyboardMarkup(
    [
        [
            InlineKeyboardButton(
                text="查看帮助",
                url=f"t.me/{BOT_USERNAME}?start=help",
            ),
            InlineKeyboardButton(
                text="运行状态",
                callback_data="stats_callback",
            ),
        ],
    ]
)


FED_MARKUP = InlineKeyboardMarkup(
    [
        [
            InlineKeyboardButton(
                "联邦创建者命令", callback_data="fed_owner"
            ),
            InlineKeyboardButton(
                "联邦管理员命令", callback_data="fed_admin"
            ),
        ],
        [
            InlineKeyboardButton("普通成员命令", callback_data="fed_user"),
        ],
        [
            InlineKeyboardButton("返回", callback_data="help_back"),
        ],
    ]
)


@app.on_message(filters.command("start"))
async def start(_, message):
    if message.chat.type != ChatType.PRIVATE:
        return await message.reply(
            "请私聊我查看详细功能。", reply_markup=keyboard
        )
    if len(message.text.split()) > 1:
        user = await app.get_users(message.from_user.id)
        name = (message.text.split(None, 1)[1]).lower()
        match = re.match(r"rules_(.*)", name)
        if match:
            chat_id = match.group(1)
            user_id = message.from_user.id
            chat = await app.get_chat(int(chat_id))
            text = f"**`{chat.title}` 的群规：\n\n**"
            rules = await get_rules(int(chat_id))
            if rules:
                text = text + rules
                if "{chat}" in text:
                    text = text.replace("{chat}", chat.title)
                if "{name}" in text:
                    text = text.replace("{name}", user.mention)
                keyb = None
                if re.findall(r"\[.+\,.+\]", text):
                    text, keyb = extract_text_and_keyb(ikb, text)
                await app.send_message(user_id, text=text, reply_markup=keyb)
            else:
                await app.send_message(
                    user_id,
                    "该群管理员尚未设置群规。",
                )
            return
        if name == "directory" and has_permission(
            "teacher_directory.manage", BOT_PERMISSIONS
        ):
            from wbb.modules.teacher_directory import send_directory_menu

            return await send_directory_menu(message)
        if name == "mkdwn_help":
            await message.reply(
                MARKDOWN,
                parse_mode=ParseMode.HTML,
                disable_web_page_preview=True,
            )
        elif "_" in name:
            module = name.split("_", 1)[1]
            if module not in HELPABLE:
                return await message.reply("未找到对应的中文帮助内容。")
            text = (
                f"以下是 **{HELPABLE[module].__MODULE__}** 的使用说明：\n"
                + HELPABLE[module].__HELP__
            )
            if module == "feds":
                return await message.reply(
                    text=text,
                    reply_markup=FED_MARKUP,
                    disable_web_page_preview=True,
                )
            await message.reply(
                text,
                reply_markup=InlineKeyboardMarkup(
                    [[InlineKeyboardButton("返回", callback_data="help_back")]]
                ),
                disable_web_page_preview=True,
            )
        elif name == "help":
            text, keyb = await help_parser(message.from_user.first_name)
            await message.reply(
                text,
                reply_markup=keyb,
            )
    else:
        await message.reply(
            home_text_pm,
            reply_markup=home_keyboard_pm,
        )
        if has_permission("teacher_directory.manage", BOT_PERMISSIONS):
            from wbb.modules.teacher_directory import send_directory_menu

            await send_directory_menu(message)
    return


@app.on_message(filters.command("help"))
async def help_command(_, message):
    if message.chat.type != ChatType.PRIVATE:
        if len(message.command) >= 2:
            name = (message.text.split(None, 1)[1]).replace(" ", "_").lower()
            if str(name) in HELPABLE:
                key = InlineKeyboardMarkup(
                    [
                        [
                            InlineKeyboardButton(
                                text="查看帮助",
                                url=f"t.me/{BOT_USERNAME}?start=help_{name}",
                            )
                        ],
                    ]
                )
                await message.reply(
                    f"点击下方按钮查看“{HELPABLE[name].__MODULE__}”的帮助。",
                    reply_markup=key,
                )
            else:
                await message.reply(
                    "请私聊我查看详细功能。", reply_markup=keyboard
                )
        else:
            await message.reply(
                "请私聊我查看详细功能。", reply_markup=keyboard
            )
    else:
        if len(message.command) >= 2:
            name = (message.text.split(None, 1)[1]).replace(" ", "_").lower()
            if str(name) in HELPABLE:
                text = (
                    f"以下是 **{HELPABLE[name].__MODULE__}** 的使用说明：\n"
                    + HELPABLE[name].__HELP__
                )
                await message.reply(text, disable_web_page_preview=True)
            else:
                text, help_keyboard = await help_parser(
                    message.from_user.first_name
                )
                await message.reply(
                    text,
                    reply_markup=help_keyboard,
                    disable_web_page_preview=True,
                )
        else:
            text, help_keyboard = await help_parser(
                message.from_user.first_name
            )
            await message.reply(
                text, reply_markup=help_keyboard, disable_web_page_preview=True
            )
    return


async def help_parser(name, keyboard=None):
    if not keyboard:
        keyboard = InlineKeyboardMarkup(paginate_modules(0, HELPABLE, "help"))
    return (
        """你好，{first_name}，我是 {bot_name}。
我可以管理群组、自动回复、积分和抽奖。
请选择下方功能分类查看中文命令说明。
""".format(
            first_name=name,
            bot_name=BOT_NAME,
        ),
        keyboard,
    )


@app.on_callback_query(filters.regex("bot_commands"))
async def commands_callbacc(_, CallbackQuery):
    text, keyboard = await help_parser(CallbackQuery.from_user.mention)
    await app.send_message(
        CallbackQuery.message.chat.id,
        text=text,
        reply_markup=keyboard,
    )

    await CallbackQuery.message.delete()


@app.on_callback_query(filters.regex("^manage_hint$"))
async def manage_hint_callback(_, query):
    if not has_any_permission(PRIVATE_MANAGEMENT_PERMISSIONS, BOT_PERMISSIONS):
        return await query.answer("当前机器人未分配私聊管理职责。", show_alert=True)
    await query.answer("请发送 /manage 打开群组管理菜单。", show_alert=True)


@app.on_callback_query(filters.regex("stats_callback"))
async def stats_callbacc(_, CallbackQuery):
    text = await bot_sys_stats()
    await app.answer_callback_query(CallbackQuery.id, text, show_alert=True)


@app.on_callback_query(filters.regex(r"help_(.*?)"))
async def help_button(client, query):
    home_match = re.match(r"help_home\((.+?)\)", query.data)
    mod_match = re.match(r"help_module\((.+?)\)", query.data)
    prev_match = re.match(r"help_prev\((.+?)\)", query.data)
    next_match = re.match(r"help_next\((.+?)\)", query.data)
    back_match = re.match(r"help_back", query.data)
    create_match = re.match(r"help_create", query.data)
    top_text = f"""
你好，{query.from_user.first_name}，我是 {BOT_NAME}。
我可以管理群组、自动回复、积分和抽奖。
请选择下方功能分类查看中文命令说明。

常用命令：
 - /start：打开机器人首页
 - /help：查看功能帮助
 - /manage：打开私聊群管理菜单
 """
    if mod_match:
        module = (mod_match.group(1)).replace(" ", "_")
        text = (
            "{} **{}**:\n".format(
                "以下是使用说明", HELPABLE[module].__MODULE__
            )
            + HELPABLE[module].__HELP__
        )
        if module == "feds":
            return await query.message.edit(
                text=text,
                reply_markup=FED_MARKUP,
                disable_web_page_preview=True,
            )
        await query.message.edit(
            text=text,
            reply_markup=InlineKeyboardMarkup(
                [[InlineKeyboardButton("返回", callback_data="help_back")]]
            ),
            disable_web_page_preview=True,
        )
    elif home_match:
        await app.send_message(
            query.from_user.id,
            text=home_text_pm,
            reply_markup=home_keyboard_pm,
        )
        await query.message.delete()
    elif prev_match:
        curr_page = int(prev_match.group(1))
        await query.message.edit(
            text=top_text,
            reply_markup=InlineKeyboardMarkup(
                paginate_modules(curr_page - 1, HELPABLE, "help")
            ),
            disable_web_page_preview=True,
        )

    elif next_match:
        next_page = int(next_match.group(1))
        await query.message.edit(
            text=top_text,
            reply_markup=InlineKeyboardMarkup(
                paginate_modules(next_page + 1, HELPABLE, "help")
            ),
            disable_web_page_preview=True,
        )

    elif back_match:
        await query.message.edit(
            text=top_text,
            reply_markup=InlineKeyboardMarkup(
                paginate_modules(0, HELPABLE, "help")
            ),
            disable_web_page_preview=True,
        )

    elif create_match:
        text, keyboard = await help_parser(query)
        await query.message.edit(
            text=text,
            reply_markup=keyboard,
            disable_web_page_preview=True,
        )

    return await client.answer_callback_query(query.id)


if __name__ == "__main__":
    install()
    with closing(loop):
        with suppress(asyncio.exceptions.CancelledError):
            loop.run_until_complete(start_bot())
        loop.run_until_complete(asyncio.sleep(3.0))  # task cancel wait