Files
shirobana-bot/bot-utility/main.py
T
RRRosmontis 34089d0042 Initial import of Shirobana bot backend
Include utility modules, shared channels, prompts, read-only skills, wardrobe memory and regression tests. Externalize legacy credentials and exclude private runtime state.
2026-11-11 00:25:30 +08:00

362 lines
16 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
# main.py — TG 主入口
# 启动内容:
# 1. 加载 settings.json(TG token / 代理 / 各通道管理员)
# 2. 共享后台服务: live_monitor Flask(:1986) + 开播通知消费者、
# mail 服务(:9088 绑定确认 + 轮询线程 + 通知消费者)、
# deepseek 余额监控、komari 服务器监控、ginkgo 白果云通知轮询、
# 统一身份库 identity(bridges/accounts)
# 3. NapCat QQ 桥 (bridges/napcat.py, OneBot 11 反向 WS :9089)
# 4. Telegram bot(aiogram 长轮询,经本地 xray 代理出站)
import asyncio
import json
import os
import queue
import sys
import tempfile
import threading
import time
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
# 本目录 (bot/utility) 自动在 sys.path; 父目录 bot/ 提供共享平台层 channel/tg_channel
# chat/ 是独立的 chat bot(bot/chat/main.py 入口),本入口不引用
sys.path.insert(0, os.path.dirname(BASE_DIR))
from channel import (Channel, register_channel, get_channel, send_message,
push_to_user, load_settings, admin_list)
from tg_channel import TGChannel, TG_USER_OFFSET, TG_GROUP_OFFSET
from dispatch import handle_message
import tg_watchdog
import live_monitor
import mail
import ginkgostu
import identity
import for_admin_komari_agent
from deepseek import deepseek_balance
from bridges.napcat import start_napcat_bridge
import help_ui
import notifications
import notification_render
import scheduler
import notebook
from agent.modes import ModeStore
from aiogram import Bot, Dispatcher, types
from aiogram.client.session.aiohttp import AiohttpSession
# ---------------- 开播通知消费者(移植自旧 main.py,通道化) ----------------
async def check_live_notifications():
while True:
try:
while True:
notif = live_monitor.notification_queue.get_nowait()
room_id = notif["room_id"]
site = notif["site"]
chat_type = notif["chat_type"]
uid = notif.get("uid")
label = live_monitor.SITE_LABEL.get(site, "主站")
url = f"https://live.rrrosmontis.icu/players/player.html?site={site}&room={room_id}"
msg = f"注意到 {label} 的 {room_id} 房间似乎正在直播!\n{url}"
chat_id = notif.get("chat_id")
event = notifications.Event(
kind="live",
dedupe_key=f"live:{site}:{room_id}:{chat_type}:{uid or chat_id}:{int(notif.get('ts') or notif.get('time') or time.time())}",
legacy_text=msg,
uid=uid if chat_type == "private" else None,
channel=None if chat_type == "private" else notif.get("channel", "qq"),
msg_type=chat_type,
target_id=None if chat_type == "private" else chat_id,
payload={"room_id": room_id, "site": site, "label": label, "url": url},
)
notifications.center.publish(event)
print(f"[live_monitor] 已入队 {label} 房间 {room_id} 开播通知")
except queue.Empty:
pass
await asyncio.sleep(2)
# ---------------- Telegram 消息处理 ----------------
def _unlink_temp_paths(paths):
for path in paths:
try:
os.unlink(path)
except OSError:
pass
def make_tg_handler(tg_channel: TGChannel):
async def on_message(message: types.Message, bot: Bot):
if not message.from_user or message.from_user.is_bot:
return
downloaded_paths = []
try:
raw_uid = message.from_user.id
ns_uid = TGChannel.ns_user(raw_uid)
image_infos = []
# 群聊 mention 剥离在本地变量上进行(aiogram 3 的 Message 是 frozen model,
# 不能就地改写 message.text)
msg_text = message.text
if message.chat.type == "private":
msg_type = "private"
target_id = ns_uid
is_mention = False
else:
msg_type = "group"
target_id = TGChannel.ns_group(message.chat.id)
is_mention = False
# 回复 bot 的消息
rt = message.reply_to_message
if rt and rt.from_user and rt.from_user.id == bot.id:
is_mention = True
# @bot mention(实体偏移按 UTF-16 计,用 extract_from 取文本,
# 不能用 str 下标切片——非 BMP 字符会错位)
if msg_text:
bot_uname = (tg_channel.tg_username or "").lower()
for ent in (message.entities or []):
if ent.type != "mention":
continue
frag = ent.extract_from(message.text)
if frag and frag.lstrip("@").lower() == bot_uname:
is_mention = True
msg_text = msg_text.replace(frag, "", 1)
# mention 剥离后重新取文本
text = (msg_text or message.caption or "").strip()
# 图片(取最大张)
if message.photo:
ph = message.photo[-1]
suffix = ".jpg"
fd, path = tempfile.mkstemp(suffix=suffix)
os.close(fd)
try:
await bot.download(ph, destination=path)
downloaded_paths.append(path)
image_infos.append({
"url": path, # TG 端已落盘,主模型直接读取
"file": path,
"file_size": str(ph.file_size or os.path.getsize(path)),
})
except Exception as e:
print(f"[tg] 图片下载失败: {e}")
try:
os.unlink(path)
except OSError:
pass
# 用户贴纸:保存原始文件;动态/视频贴纸用缩略图供主模型学习。
sticker = getattr(message, "sticker", None)
if sticker:
is_animated = bool(getattr(sticker, "is_animated", False))
is_video = bool(getattr(sticker, "is_video", False))
suffix = ".tgs" if is_animated else ".webm" if is_video else ".webp"
paths = []
try:
fd, original_path = tempfile.mkstemp(suffix=suffix)
os.close(fd)
paths.append(original_path)
await bot.download(sticker, destination=original_path)
downloaded_paths.append(original_path)
preview_path = original_path
if is_animated or is_video:
thumbnail = getattr(sticker, "thumbnail", None)
if not thumbnail:
raise ValueError("动态贴纸缺少可用缩略图")
fd, preview_path = tempfile.mkstemp(suffix=".jpg")
os.close(fd)
paths.append(preview_path)
await bot.download(thumbnail, destination=preview_path)
downloaded_paths.append(preview_path)
sticker_count = sum(1 for info in image_infos if info.get("kind") == "sticker")
image_infos.append({
"kind": "sticker",
"source_id": f"s{sticker_count + 1}",
"url": preview_path,
"file": original_path,
"original_path": original_path,
"file_size": str(os.path.getsize(preview_path)),
"file_unique_id": getattr(sticker, "file_unique_id", ""),
"emoji": getattr(sticker, "emoji", "") or "",
"set_name": getattr(sticker, "set_name", "") or "",
"source_uid": ns_uid,
"format": "animated" if is_animated else "video" if is_video else "static",
})
except Exception as e:
print(f"[tg] 贴纸下载失败: {e}")
for path in paths:
try:
os.unlink(path)
except OSError:
pass
if not text and not image_infos:
return
# TG 根 Help 在 command/chat 两种模式下都保留原生按钮;
# 仅正文渲染方式不同,模块/QQ 仍走统一 dispatch。
if text.lower() in (".help", "/help") or text.lower().startswith("/help@"):
uid = identity.resolve("tg", ns_uid) or identity.ensure_tg(ns_uid, tg_channel.tg_username)
group_id = target_id if msg_type == "group" else None
include_admin = tg_channel.is_admin(uid)
mode = ModeStore().get_mode("tg", uid, group_id)
if mode == "chat":
content = await help_ui.chat_help(uid, None, include_admin)
else:
content = help_ui.command_index(include_admin)
await bot.send_message(
message.chat.id,
content,
reply_markup=_help_keyboard(include_admin),
)
return
print(f"[tg] {msg_type}/{target_id} uid={raw_uid} mention={is_mention} text={text[:60]!r} imgs={len(image_infos)}")
await handle_message(tg_channel, msg_type, ns_uid, target_id,
text, image_infos, is_mention)
except Exception as e:
import traceback
print(f"[tg] 消息处理异常: {e}")
traceback.print_exc()
finally:
if downloaded_paths:
asyncio.get_running_loop().call_later(900, _unlink_temp_paths, tuple(downloaded_paths))
return on_message
# ---------------- Telegram Help buttons ----------------
def _help_keyboard(include_admin: bool):
from aiogram.types import InlineKeyboardButton, InlineKeyboardMarkup
rows, row = [], []
for index, item in enumerate(help_ui.SERVICES, 1):
if item.admin and not include_admin:
continue
row.append(InlineKeyboardButton(text=item.title, callback_data=f"help:{index}"))
if len(row) == 3:
rows.append(row)
row = []
if row:
rows.append(row)
return InlineKeyboardMarkup(inline_keyboard=rows)
def make_tg_help_callback(tg_channel: TGChannel):
async def on_help(callback: types.CallbackQuery, bot: Bot):
if not callback.from_user or not callback.message:
return
item = help_ui.service((callback.data or "").split(":", 1)[-1])
if item is None:
await callback.answer("这个帮助入口已经失效")
return
ns_uid = TGChannel.ns_user(callback.from_user.id)
uid = identity.resolve("tg", ns_uid) or identity.ensure_tg(ns_uid, tg_channel.tg_username)
include_admin = tg_channel.is_admin(uid)
if item.admin and not include_admin:
await callback.answer("该功能仅管理员可用", show_alert=True)
return
await callback.answer()
chat = callback.message.chat
group_id = TGChannel.ns_group(chat.id) if chat.type != "private" else None
mode = ModeStore().get_mode("tg", uid, group_id)
if mode == "chat":
content = await help_ui.chat_help(uid, item.key, include_admin)
else:
content = help_ui.command_help(item.key, include_admin)
await bot.send_message(
chat.id,
content,
reply_markup=_help_keyboard(include_admin),
)
return on_help
# ---------------- 主入口 ----------------
async def main():
settings = load_settings()
tg_token = settings.get("tg_token", "")
tg_proxy = settings.get("tg_proxy", "")
# 0) 统一身份库(账号映射目录由入口注入,平台层不感知 utility 目录结构)
identity.configure(os.path.join(BASE_DIR, "bridges", "accounts"),
settings.get("admin", {}).get("tg"))
from chat import engine
engine.init_chat()
notifications.center.configure(notification_render.render_event)
notifications.center.register_acknowledger(
"mail", lambda event: mail._advance_last_seen(
int(event.uid), event.payload["addr"], int(event.payload["mid"])))
# 1) 共享后台服务
threading.Thread(target=live_monitor.run_flask, daemon=True).start()
print("[live_monitor] Flask 服务器已启动在端口 1986")
mail.start_services()
print("[mail] 邮件服务已启动(确认 HTTP :9088 + 轮询线程)")
# create_task 必须持有强引用:事件循环只持弱引用,无引用的任务可能被 GC 中途回收
bg_tasks = [
asyncio.create_task(check_live_notifications()),
asyncio.create_task(mail.notification_loop()),
asyncio.create_task(deepseek_balance.monitor_loop()),
asyncio.create_task(for_admin_komari_agent.monitor_loop()),
asyncio.create_task(ginkgostu.monitor_loop()),
asyncio.create_task(notifications.center.worker()),
asyncio.create_task(scheduler.scheduler.worker()),
asyncio.create_task(notebook.notebook.cleanup_worker()),
]
# 2) QQ 桥(NapCat 反向 WS :9089)
bg_tasks.append(asyncio.create_task(start_napcat_bridge(settings.get("admin", {}).get("qq"))))
# 3) Telegram bot
if not tg_token:
print("[tg] settings.json 缺少 tg_token,Telegram 通道不启动")
await asyncio.Future()
return
# aiogram 3.31: 代理走 AiohttpSession(需 aiohttp-socks)
bot = Bot(token=tg_token, session=AiohttpSession(proxy=tg_proxy or None))
tg_channel = TGChannel(bot, admin_ids=settings.get("admin", {}).get("tg"))
for _attempt in range(30):
try:
tg_channel.tg_username = (await bot.get_me()).username or ""
break
except Exception as e:
print(f"[tg] get_me 失败(第{_attempt+1}次): {e},5秒后重试")
await asyncio.sleep(5)
if _attempt == 29:
raise
register_channel(tg_channel)
print(f"[tg] Telegram 通道已启动: @{tg_channel.tg_username}"
+ (f" (proxy={tg_proxy})" if tg_proxy else ""))
# 注册 TG 斜杠命令菜单(.help 仍可用;管理员指令不放入菜单,直接输入即可)
try:
from aiogram.types import BotCommand
await bot.set_my_commands([
BotCommand(command="help", description="查看服务与帮助"),
BotCommand(command="my", description="统一身份与个人设置"),
BotCommand(command="setting", description="切换 command/chat 模式"),
])
print("[tg] 斜杠命令菜单已注册")
except Exception as e:
print(f"[tg] set_my_commands 失败(不影响运行): {e}")
dp = Dispatcher()
dp.message.register(make_tg_handler(tg_channel))
dp.callback_query.register(make_tg_help_callback(tg_channel))
for _attempt in range(30):
try:
await bot.delete_webhook(drop_pending_updates=False)
break
except Exception as e:
print(f"[tg] delete_webhook 失败(第{_attempt+1}次): {e},5秒后重试")
await asyncio.sleep(5)
if _attempt == 29:
raise
print("[tg] 开始长轮询...")
bg_tasks.append(asyncio.create_task(tg_watchdog.tg_watchdog(bot, dp, name="utility")))
await dp.start_polling(bot, close_bot=True)
if __name__ == "__main__":
asyncio.run(main())