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("find_technician", "浏览并咨询技师"),
    BotCommand("technician_profile", "管理技师服务资料"),
    BotCommand("review_teacher", "主动评价技师"),
    BotCommand("technician_top", "查看审核评价排行"),
]
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",
    "find_technician": "teacher_directory.manage",
    "my_services": "teacher_directory.manage",
    "technician_profile": "teacher_directory.manage",
    "review_teacher": "teacher_directory.manage",
    "technician_top": "teacher_directory.manage",
    "review_qr": "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

    loaded_modules = {}
    for module in ALL_MODULES:
        imported_module = importlib.import_module("wbb.modules." + module)
        loaded_modules[module] = imported_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 += f"|{i:<15}|\n"
            j = 0
        else:
            bot_modules += f"|{i:<15}"
        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}")

    business_assistant_module = loaded_modules.get("business_assistant")
    if business_assistant_module and hasattr(
        business_assistant_module, "start_business_assistant_runtime"
    ):
        try:
            await business_assistant_module.start_business_assistant_runtime()
        except Exception as exc:
            log.error(f"启动 Telegram Business 智能接待失败：{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()

    if business_assistant_module and hasattr(
        business_assistant_module, "stop_business_assistant_runtime"
    ):
        with suppress(Exception):
            await business_assistant_module.stop_business_assistant_runtime()

    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=services",
        )
    )

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", "技师咨询与评价"),
        ("business_assistant.manage", "Telegram Business 智能接待"),
    )
    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)
        start_parameter = message.text.split(None, 1)[1]
        if start_parameter.startswith("review_") and has_permission(
            "teacher_directory.manage", BOT_PERMISSIONS
        ):
            from wbb.modules.technician_services import open_review_qr

            return await open_review_qr(
                message,
                start_parameter.removeprefix("review_"),
            )
        name = start_parameter.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
        directory_match = re.fullmatch(r"directory_(-?\d+)", name)
        if directory_match and has_permission(
            "teacher_directory.manage", BOT_PERMISSIONS
        ):
            from wbb.modules.teacher_directory import send_directory_menu

            return await send_directory_menu(
                message,
                source_chat_id=int(directory_match.group(1)),
            )
        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 == "services" and has_permission(
            "teacher_directory.manage", BOT_PERMISSIONS
        ):
            from wbb.modules.technician_services import send_service_home

            return await send_service_home(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.technician_services import send_service_home

            await send_service_home(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 (
        f"""你好，{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