From e572b249d2428a03520f42320bc6f45f7e66cec9 Mon Sep 17 00:00:00 2001 From: xream Date: Mon, 9 Dec 2024 22:31:23 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E6=94=AF=E6=8C=81=E5=A4=9A=E9=9D=A2?= =?UTF-8?q?=E6=9D=BF?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- bot.py | 176 +++++++++++++++++++++++++++++++++++++++++++++++----- database.py | 141 ++++++++++++++++++++++++++++++++++++++--- 2 files changed, 293 insertions(+), 24 deletions(-) diff --git a/bot.py b/bot.py index 5473a4b..a922628 100755 --- a/bot.py +++ b/bot.py @@ -30,7 +30,7 @@ TELEGRAM_TOKEN = os.getenv('TELEGRAM_TOKEN') DATABASE_PATH = 'users.db' # 定义阶段 -BIND_USERNAME, BIND_PASSWORD, BIND_DASHBOARD = range(3) +BIND_USERNAME, BIND_PASSWORD, BIND_DASHBOARD, BIND_ALIAS = range(4) SEARCH_SERVER = range(1) # 初始化数据库 @@ -95,6 +95,7 @@ async def help_command(update: Update, context: ContextTypes.DEFAULT_TYPE): 可用命令： /bind - 绑定账号 /unbind - 解绑账号 +/dashboard - 管理面板 /overview - 查看服务器状态总览 /server - 查看单台服务器状态 /cron - 执行计划任务 @@ -108,13 +109,8 @@ async def bind_start(update: Update, context: ContextTypes.DEFAULT_TYPE): await update.message.reply_text("请与机器人私聊进行绑定操作，\n避免机密信息泄露。") return ConversationHandler.END - user = await db.get_user(update.effective_user.id) - if user: - await update.message.reply_text("您已绑定账号，如需重新绑定，请先使用 /unbind 命令解绑。") - return ConversationHandler.END - else: - await update.message.reply_text("请输入您的用户名：") - return BIND_USERNAME + await update.message.reply_text("请输入您的用户名：") + return BIND_USERNAME async def bind_username(update: Update, context: ContextTypes.DEFAULT_TYPE): context.user_data['username'] = update.message.text.strip() @@ -129,9 +125,16 @@ async def bind_password(update: Update, context: ContextTypes.DEFAULT_TYPE): async def bind_dashboard(update: Update, context: ContextTypes.DEFAULT_TYPE): dashboard_url = update.message.text.strip() context.user_data['dashboard_url'] = dashboard_url + await update.message.reply_text("请为这个面板设置一个别名（如：主面板、备用等）：") + return BIND_ALIAS + +async def bind_alias(update: Update, context: ContextTypes.DEFAULT_TYPE): + alias = update.message.text.strip() + context.user_data['alias'] = alias telegram_id = update.effective_user.id username = context.user_data['username'] password = context.user_data['password'] + dashboard_url = context.user_data['dashboard_url'] # 测试连接 try: @@ -143,17 +146,29 @@ async def bind_dashboard(update: Update, context: ContextTypes.DEFAULT_TYPE): return ConversationHandler.END # 保存到数据库 - await db.add_user(telegram_id, username, password, dashboard_url) + await db.add_user(telegram_id, username, password, dashboard_url, alias) await update.message.reply_text("绑定成功！您现在可以使用机器人的功能了。") return ConversationHandler.END async def unbind(update: Update, context: ContextTypes.DEFAULT_TYPE): - user = await db.get_user(update.effective_user.id) - if user: - await db.delete_user(update.effective_user.id) - await update.message.reply_text("已解绑。") - else: - await update.message.reply_text("您尚未绑定账号。") + dashboards = await db.get_all_dashboards(update.effective_user.id) + if not dashboards: + await update.message.reply_text("您尚未绑定任何面板。") + return + + keyboard = [] + # 添加每个 dashboard 的解绑选项 + for dashboard in dashboards: + default_mark = "（默认）" if dashboard['is_default'] else "" + button_text = f"解绑 {dashboard['alias']}{default_mark}" + keyboard.append([InlineKeyboardButton(button_text, callback_data=f"unbind_{dashboard['id']}")]) + + # 添加解绑所有的选项 + if len(dashboards) > 1: + keyboard.append([InlineKeyboardButton("解绑所有面板", callback_data="unbind_all")]) + + reply_markup = InlineKeyboardMarkup(keyboard) + await update.message.reply_text("请选择要解绑的面板：", reply_markup=reply_markup) async def overview(update: Update, context: ContextTypes.DEFAULT_TYPE): user = await db.get_user(update.effective_user.id) @@ -246,6 +261,73 @@ async def button_handler(update: Update, context: ContextTypes.DEFAULT_TYPE): query = update.callback_query data = query.data + if data.startswith('unbind_'): + if data == 'unbind_all': + await db.delete_user(query.from_user.id) + await query.edit_message_text("已解绑所有面板，您可以使用 /bind 重新绑定。") + else: + dashboard_id = int(data.split('_')[-1]) + # 获取当前面板信息，用于判断是否是默认面板 + dashboards = await db.get_all_dashboards(query.from_user.id) + current_dashboard = next((d for d in dashboards if d['id'] == dashboard_id), None) + was_default = current_dashboard and current_dashboard['is_default'] + + has_remaining = await db.delete_dashboard(query.from_user.id, dashboard_id) + + if not has_remaining: + await query.edit_message_text("已解绑最后一个面板，您可以使用 /bind 重新绑定。") + else: + # 新面板列表 + dashboards = await db.get_all_dashboards(query.from_user.id) + keyboard = [] + + # 如果解绑的是默认面板，显示新的默认面板提示 + if was_default: + new_default = next((d for d in dashboards if d['is_default']), None) + message = f"已解绑面板，新的默认面板已设置为：{new_default['alias']}\n\n请选择要解绑的面板：" + else: + message = "请选择要解绑的面板：" + + for dashboard in dashboards: + default_mark = "（默认）" if dashboard['is_default'] else "" + button_text = f"解绑 {dashboard['alias']}{default_mark}" + keyboard.append([InlineKeyboardButton(button_text, callback_data=f"unbind_{dashboard['id']}")]) + + if len(dashboards) > 1: + keyboard.append([InlineKeyboardButton("解绑所有面板", callback_data="unbind_all")]) + + reply_markup = InlineKeyboardMarkup(keyboard) + await query.edit_message_text(message, reply_markup=reply_markup) + return + + elif data.startswith('set_default_'): + dashboard_id = int(data.split('_')[-1]) + dashboards = await db.get_all_dashboards(query.from_user.id) + selected_dashboard = next((d for d in dashboards if d['id'] == dashboard_id), None) + + if not selected_dashboard: + await query.answer("未找到该面板", show_alert=True) + return + + if selected_dashboard['is_default']: + await query.answer("这已经是默认面板了", show_alert=True) + return + + # 直接切换默认面板 + await db.set_default_dashboard(query.from_user.id, dashboard_id) + + # 更新面板列表 + dashboards = await db.get_all_dashboards(query.from_user.id) + keyboard = [] + for dashboard in dashboards: + default_mark = "（当前默认）" if dashboard['is_default'] else "" + button_text = f"{dashboard['alias']}{default_mark}" + keyboard.append([InlineKeyboardButton(button_text, callback_data=f"set_default_{dashboard['id']}")]) + + reply_markup = InlineKeyboardMarkup(keyboard) + await query.edit_message_text("您的面板列表：", reply_markup=reply_markup) + return + user = await db.get_user(query.from_user.id) if not user: await query.answer("请先使用 /bind 命令绑定您的账号。", show_alert=True) @@ -491,6 +573,53 @@ async def button_handler(update: Update, context: ContextTypes.DEFAULT_TYPE): elif data == 'refresh_availability': await view_availability(query, context, api) + elif data.startswith('set_default_'): + dashboard_id = int(data.split('_')[-1]) + await db.set_default_dashboard(query.from_user.id, dashboard_id) + await query.edit_message_text("已更新默认面板。") + return + + elif data.startswith('dashboard_'): + dashboard_id = int(data.split('_')[-1]) + dashboards = await db.get_all_dashboards(query.from_user.id) + selected_dashboard = next((d for d in dashboards if d['id'] == dashboard_id), None) + + if not selected_dashboard: + await query.answer("未找到该面板", show_alert=True) + return + + if selected_dashboard['is_default']: + await query.answer("这已经是默认面板了", show_alert=True) + return + + # 直接切换默认面板 + await db.set_default_dashboard(query.from_user.id, dashboard_id) + + # 更新面板列表 + dashboards = await db.get_all_dashboards(query.from_user.id) + keyboard = [] + for dashboard in dashboards: + default_mark = "（当前默认）" if dashboard['is_default'] else "" + button_text = f"{dashboard['alias']}{default_mark}" + keyboard.append([InlineKeyboardButton(button_text, callback_data=f"set_default_{dashboard['id']}")]) + + reply_markup = InlineKeyboardMarkup(keyboard) + await query.edit_message_text("您的面板列表：", reply_markup=reply_markup) + return + + elif data == "dashboard_back": + # 返回面板列表 + dashboards = await db.get_all_dashboards(query.from_user.id) + keyboard = [] + for dashboard in dashboards: + default_mark = "（当前默认）" if dashboard['is_default'] else "" + button_text = f"{dashboard['alias']}{default_mark}" + keyboard.append([InlineKeyboardButton(button_text, callback_data=f"set_default_{dashboard['id']}")]) + + reply_markup = InlineKeyboardMarkup(keyboard) + await query.edit_message_text("您的面板列表：", reply_markup=reply_markup) + return + async def view_loop_traffic(query, context, api): # 获取服务状态 try: @@ -625,6 +754,21 @@ async def services_overview(update: Update, context: ContextTypes.DEFAULT_TYPE): reply_markup = InlineKeyboardMarkup(keyboard) await update.message.reply_text("请选择要查看的服务信息：", reply_markup=reply_markup) +async def dashboard(update: Update, context: ContextTypes.DEFAULT_TYPE): + dashboards = await db.get_all_dashboards(update.effective_user.id) + if not dashboards: + await update.message.reply_text("您还没有绑定任何面板。") + return + + keyboard = [] + for dashboard in dashboards: + default_mark = "（当前默认）" if dashboard['is_default'] else "" + button_text = f"{dashboard['alias']}{default_mark}" + keyboard.append([InlineKeyboardButton(button_text, callback_data=f"set_default_{dashboard['id']}")]) + + reply_markup = InlineKeyboardMarkup(keyboard) + await update.message.reply_text("您的面板列表：", reply_markup=reply_markup) + def main(): application = ApplicationBuilder().token(TELEGRAM_TOKEN).build() @@ -642,6 +786,7 @@ def main(): application.add_handler(CommandHandler('overview', overview)) application.add_handler(CommandHandler('cron', cron_jobs)) application.add_handler(CommandHandler('services', services_overview)) + application.add_handler(CommandHandler('dashboard', dashboard)) # 绑定命令的会话处理 bind_handler = ConversationHandler( @@ -650,6 +795,7 @@ def main(): BIND_USERNAME: [MessageHandler(filters.TEXT & ~filters.COMMAND, bind_username)], BIND_PASSWORD: [MessageHandler(filters.TEXT & ~filters.COMMAND, bind_password)], BIND_DASHBOARD: [MessageHandler(filters.TEXT & ~filters.COMMAND, bind_dashboard)], + BIND_ALIAS: [MessageHandler(filters.TEXT & ~filters.COMMAND, bind_alias)], }, fallbacks=[] ) diff --git a/database.py b/database.py index 89a0f41..ab56964 100755 --- a/database.py +++ b/database.py @@ -8,34 +8,157 @@ class Database: async def initialize(self): async with aiosqlite.connect(self.db_path) as db: + # 创建用户表 await db.execute(''' CREATE TABLE IF NOT EXISTS users ( telegram_id INTEGER PRIMARY KEY, + default_dashboard_id INTEGER, + FOREIGN KEY (default_dashboard_id) REFERENCES dashboards (id) + ) + ''') + # 创建 dashboard 表 + await db.execute(''' + CREATE TABLE IF NOT EXISTS dashboards ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + telegram_id INTEGER, username TEXT NOT NULL, password TEXT NOT NULL, - dashboard_url TEXT NOT NULL + dashboard_url TEXT NOT NULL, + alias TEXT, + FOREIGN KEY (telegram_id) REFERENCES users (telegram_id) ) ''') await db.commit() - async def add_user(self, telegram_id, username, password, dashboard_url): + async def add_user(self, telegram_id, username, password, dashboard_url, alias=None): + async with aiosqlite.connect(self.db_path) as db: + # 首先确保用户存在 + await db.execute('INSERT OR IGNORE INTO users (telegram_id) VALUES (?)', (telegram_id,)) + + # 如果没有提供别名，使用 URL 的第一部分作为默认别名 + if not alias: + try: + alias = dashboard_url.split('://')[1].split('.')[0].upper() + except: + alias = "NEZHA" + + # 添加 dashboard + cursor = await db.execute(''' + INSERT INTO dashboards (telegram_id, username, password, dashboard_url, alias) + VALUES (?, ?, ?, ?, ?) + RETURNING id + ''', (telegram_id, username, password, dashboard_url, alias)) + dashboard_id = (await cursor.fetchone())[0] + + # 如果用户还没有默认 dashboard，设置这个为默认 + await db.execute(''' + UPDATE users + SET default_dashboard_id = COALESCE(default_dashboard_id, ?) + WHERE telegram_id = ? + ''', (dashboard_id, telegram_id)) + + await db.commit() + return dashboard_id + + async def update_alias(self, dashboard_id, alias): async with aiosqlite.connect(self.db_path) as db: await db.execute(''' - INSERT OR REPLACE INTO users (telegram_id, username, password, dashboard_url) - VALUES (?, ?, ?, ?) - ''', (telegram_id, username, password, dashboard_url)) + UPDATE dashboards + SET alias = ? + WHERE id = ? + ''', (alias, dashboard_id)) await db.commit() async def get_user(self, telegram_id): async with aiosqlite.connect(self.db_path) as db: - async with db.execute('SELECT username, password, dashboard_url FROM users WHERE telegram_id = ?', (telegram_id,)) as cursor: + # 获取用户的默认 dashboard + async with db.execute(''' + SELECT d.username, d.password, d.dashboard_url, d.alias + FROM users u + JOIN dashboards d ON d.id = u.default_dashboard_id + WHERE u.telegram_id = ? + ''', (telegram_id,)) as cursor: row = await cursor.fetchone() if row: - return {'username': row[0], 'password': row[1], 'dashboard_url': row[2]} - else: - return None + return {'username': row[0], 'password': row[1], 'dashboard_url': row[2], 'alias': row[3]} + return None + + async def get_all_dashboards(self, telegram_id): + async with aiosqlite.connect(self.db_path) as db: + async with db.execute(''' + SELECT d.id, d.username, d.password, d.dashboard_url, d.alias, + CASE WHEN u.default_dashboard_id = d.id THEN 1 ELSE 0 END as is_default + FROM dashboards d + LEFT JOIN users u ON u.telegram_id = d.telegram_id + WHERE d.telegram_id = ? + ''', (telegram_id,)) as cursor: + rows = await cursor.fetchall() + return [ + { + 'id': row[0], + 'username': row[1], + 'password': row[2], + 'dashboard_url': row[3], + 'alias': row[4], + 'is_default': bool(row[5]) + } + for row in rows + ] + + async def set_default_dashboard(self, telegram_id, dashboard_id): + async with aiosqlite.connect(self.db_path) as db: + await db.execute(''' + UPDATE users + SET default_dashboard_id = ? + WHERE telegram_id = ? + ''', (dashboard_id, telegram_id)) + await db.commit() + + async def delete_dashboard(self, telegram_id, dashboard_id): + async with aiosqlite.connect(self.db_path) as db: + # 检查是否是默认面板 + async with db.execute(''' + SELECT default_dashboard_id + FROM users + WHERE telegram_id = ? + ''', (telegram_id,)) as cursor: + row = await cursor.fetchone() + is_default = row and row[0] == dashboard_id + + # 删除 dashboard + await db.execute(''' + DELETE FROM dashboards + WHERE id = ? AND telegram_id = ? + ''', (dashboard_id, telegram_id)) + + # 检查是否还有其他面板 + async with db.execute(''' + SELECT id + FROM dashboards + WHERE telegram_id = ? + ORDER BY id ASC + ''', (telegram_id,)) as cursor: + remaining_dashboards = await cursor.fetchall() + + if not remaining_dashboards: + # 如果没有面板了，删除用户 + await db.execute('DELETE FROM users WHERE telegram_id = ?', (telegram_id,)) + elif is_default and remaining_dashboards: + # 如果删除的是默认面板且还有其他面板，设置第一个面板为默认 + new_default_id = remaining_dashboards[0][0] + await db.execute(''' + UPDATE users + SET default_dashboard_id = ? + WHERE telegram_id = ? + ''', (new_default_id, telegram_id)) + + await db.commit() + return bool(remaining_dashboards) # 返回是否还有其他面板 async def delete_user(self, telegram_id): async with aiosqlite.connect(self.db_path) as db: + # 删除用户的所有 dashboard + await db.execute('DELETE FROM dashboards WHERE telegram_id = ?', (telegram_id,)) + # 删除用户 await db.execute('DELETE FROM users WHERE telegram_id = ?', (telegram_id,)) await db.commit()