feat: 支持多面板

This commit is contained in:
xream 2024-12-09 22:31:23 +08:00
parent df8078b380
commit e572b249d2
No known key found for this signature in database
GPG Key ID: 1D2C5225471789F9
2 changed files with 293 additions and 24 deletions

172
bot.py
View File

@ -30,7 +30,7 @@ TELEGRAM_TOKEN = os.getenv('TELEGRAM_TOKEN')
DATABASE_PATH = 'users.db' 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) SEARCH_SERVER = range(1)
# 初始化数据库 # 初始化数据库
@ -95,6 +95,7 @@ async def help_command(update: Update, context: ContextTypes.DEFAULT_TYPE):
可用命令 可用命令
/bind - 绑定账号 /bind - 绑定账号
/unbind - 解绑账号 /unbind - 解绑账号
/dashboard - 管理面板
/overview - 查看服务器状态总览 /overview - 查看服务器状态总览
/server - 查看单台服务器状态 /server - 查看单台服务器状态
/cron - 执行计划任务 /cron - 执行计划任务
@ -108,11 +109,6 @@ async def bind_start(update: Update, context: ContextTypes.DEFAULT_TYPE):
await update.message.reply_text("请与机器人私聊进行绑定操作,\n避免机密信息泄露。") await update.message.reply_text("请与机器人私聊进行绑定操作,\n避免机密信息泄露。")
return ConversationHandler.END 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("请输入您的用户名:") await update.message.reply_text("请输入您的用户名:")
return BIND_USERNAME return BIND_USERNAME
@ -129,9 +125,16 @@ async def bind_password(update: Update, context: ContextTypes.DEFAULT_TYPE):
async def bind_dashboard(update: Update, context: ContextTypes.DEFAULT_TYPE): async def bind_dashboard(update: Update, context: ContextTypes.DEFAULT_TYPE):
dashboard_url = update.message.text.strip() dashboard_url = update.message.text.strip()
context.user_data['dashboard_url'] = dashboard_url 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 telegram_id = update.effective_user.id
username = context.user_data['username'] username = context.user_data['username']
password = context.user_data['password'] password = context.user_data['password']
dashboard_url = context.user_data['dashboard_url']
# 测试连接 # 测试连接
try: try:
@ -143,17 +146,29 @@ async def bind_dashboard(update: Update, context: ContextTypes.DEFAULT_TYPE):
return ConversationHandler.END 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("绑定成功!您现在可以使用机器人的功能了。") await update.message.reply_text("绑定成功!您现在可以使用机器人的功能了。")
return ConversationHandler.END return ConversationHandler.END
async def unbind(update: Update, context: ContextTypes.DEFAULT_TYPE): async def unbind(update: Update, context: ContextTypes.DEFAULT_TYPE):
user = await db.get_user(update.effective_user.id) dashboards = await db.get_all_dashboards(update.effective_user.id)
if user: if not dashboards:
await db.delete_user(update.effective_user.id) await update.message.reply_text("您尚未绑定任何面板。")
await update.message.reply_text("已解绑。") return
else:
await update.message.reply_text("您尚未绑定账号。") 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): async def overview(update: Update, context: ContextTypes.DEFAULT_TYPE):
user = await db.get_user(update.effective_user.id) 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 query = update.callback_query
data = query.data 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) user = await db.get_user(query.from_user.id)
if not user: if not user:
await query.answer("请先使用 /bind 命令绑定您的账号。", show_alert=True) 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': elif data == 'refresh_availability':
await view_availability(query, context, api) 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): async def view_loop_traffic(query, context, api):
# 获取服务状态 # 获取服务状态
try: try:
@ -625,6 +754,21 @@ async def services_overview(update: Update, context: ContextTypes.DEFAULT_TYPE):
reply_markup = InlineKeyboardMarkup(keyboard) reply_markup = InlineKeyboardMarkup(keyboard)
await update.message.reply_text("请选择要查看的服务信息:", reply_markup=reply_markup) 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(): def main():
application = ApplicationBuilder().token(TELEGRAM_TOKEN).build() application = ApplicationBuilder().token(TELEGRAM_TOKEN).build()
@ -642,6 +786,7 @@ def main():
application.add_handler(CommandHandler('overview', overview)) application.add_handler(CommandHandler('overview', overview))
application.add_handler(CommandHandler('cron', cron_jobs)) application.add_handler(CommandHandler('cron', cron_jobs))
application.add_handler(CommandHandler('services', services_overview)) application.add_handler(CommandHandler('services', services_overview))
application.add_handler(CommandHandler('dashboard', dashboard))
# 绑定命令的会话处理 # 绑定命令的会话处理
bind_handler = ConversationHandler( bind_handler = ConversationHandler(
@ -650,6 +795,7 @@ def main():
BIND_USERNAME: [MessageHandler(filters.TEXT & ~filters.COMMAND, bind_username)], BIND_USERNAME: [MessageHandler(filters.TEXT & ~filters.COMMAND, bind_username)],
BIND_PASSWORD: [MessageHandler(filters.TEXT & ~filters.COMMAND, bind_password)], BIND_PASSWORD: [MessageHandler(filters.TEXT & ~filters.COMMAND, bind_password)],
BIND_DASHBOARD: [MessageHandler(filters.TEXT & ~filters.COMMAND, bind_dashboard)], BIND_DASHBOARD: [MessageHandler(filters.TEXT & ~filters.COMMAND, bind_dashboard)],
BIND_ALIAS: [MessageHandler(filters.TEXT & ~filters.COMMAND, bind_alias)],
}, },
fallbacks=[] fallbacks=[]
) )

View File

@ -8,34 +8,157 @@ class Database:
async def initialize(self): async def initialize(self):
async with aiosqlite.connect(self.db_path) as db: async with aiosqlite.connect(self.db_path) as db:
# 创建用户表
await db.execute(''' await db.execute('''
CREATE TABLE IF NOT EXISTS users ( CREATE TABLE IF NOT EXISTS users (
telegram_id INTEGER PRIMARY KEY, 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, username TEXT NOT NULL,
password 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() 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: async with aiosqlite.connect(self.db_path) as db:
await db.execute(''' await db.execute('''
INSERT OR REPLACE INTO users (telegram_id, username, password, dashboard_url) UPDATE dashboards
VALUES (?, ?, ?, ?) SET alias = ?
''', (telegram_id, username, password, dashboard_url)) WHERE id = ?
''', (alias, dashboard_id))
await db.commit() await db.commit()
async def get_user(self, telegram_id): async def get_user(self, telegram_id):
async with aiosqlite.connect(self.db_path) as db: 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() row = await cursor.fetchone()
if row: if row:
return {'username': row[0], 'password': row[1], 'dashboard_url': row[2]} return {'username': row[0], 'password': row[1], 'dashboard_url': row[2], 'alias': row[3]}
else:
return None 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 def delete_user(self, telegram_id):
async with aiosqlite.connect(self.db_path) as db: 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.execute('DELETE FROM users WHERE telegram_id = ?', (telegram_id,))
await db.commit() await db.commit()