293 lines
10 KiB
Python
293 lines
10 KiB
Python
"""LinuxAwesomeBot 入口：启动 OneBot 反向 WS 服务端 + 命令处理 + 高负载告警。"""
|
||
import asyncio
|
||
import logging
|
||
import re
|
||
import signal
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
import monitor
|
||
import templates
|
||
from alerter import Alerter, Reporter
|
||
from onebot import OneBotServer
|
||
from updater import Updater
|
||
|
||
logging.basicConfig(
|
||
level=logging.INFO,
|
||
format="%(asctime)s [%(name)s] %(levelname)s: %(message)s",
|
||
datefmt="%H:%M:%S",
|
||
)
|
||
log = logging.getLogger("bot")
|
||
|
||
|
||
# ---------- 配置 ----------
|
||
|
||
def load_config() -> dict:
|
||
import tomllib
|
||
import shutil
|
||
try:
|
||
with open("config.toml", "rb") as f:
|
||
return tomllib.load(f)
|
||
except FileNotFoundError:
|
||
shutil.copy("config.example.toml", "config.toml")
|
||
log.warning("未找到 config.toml，已从模板生成，请修改后再运行（先退出）")
|
||
sys.exit(1)
|
||
|
||
|
||
# ---------- 消息格式化（模板见 templates/） ----------
|
||
|
||
def _fmt_uptime(seconds: float) -> str:
|
||
d, r = divmod(int(seconds), 86400)
|
||
h, r = divmod(r, 3600)
|
||
m = r // 60
|
||
return f"{d}天{h}小时{m}分" if d else f"{h}小时{m}分"
|
||
|
||
|
||
def _hottest(temps: dict):
|
||
flat = [t for entries in temps.values() for t in entries]
|
||
return max(flat, key=lambda t: t["current"]) if flat else None
|
||
|
||
|
||
def fmt_status(s: dict) -> str:
|
||
c, m = s["cpu"], s["memory"]
|
||
hot = _hottest(s["temps"])
|
||
net = s["net"]
|
||
d = s["disk"][0] if s["disk"] else None
|
||
return templates.render(
|
||
"status",
|
||
cpu_total=f"{sum(c['per_core']):.1f}%",
|
||
cpu_max=f"{c['cores'] * 100}%",
|
||
cpu_avg=f"{c['percent']}",
|
||
cores=c["cores"],
|
||
load=f"{c['load'][0]:.2f}/{c['load'][1]:.2f}/{c['load'][2]:.2f}",
|
||
memory=f"{m['used_gb']}/{m['total_gb']}GB ({m['percent']}%) Swap {m['swap_percent']}%",
|
||
temp=f"{hot['label']} {hot['current']}C" if hot else "无数据",
|
||
net=(f"上行 {net['up_mbps']}Mbps / 下行 {net['down_mbps']}Mbps"
|
||
if net["up_mbps"] is not None else "无数据"),
|
||
disk_mount=d["mount"] if d else "-",
|
||
disk_used=f"{d['used_gb']}/{d['total_gb']}GB" if d else "-",
|
||
disk_percent=f"{d['percent']}%" if d else "-",
|
||
uptime=_fmt_uptime(s["uptime_s"]),
|
||
)
|
||
|
||
|
||
def fmt_top(s: dict) -> str:
|
||
if not s["top"]:
|
||
return templates.render("top", top_n=10, top_list="无数据")
|
||
lines = [
|
||
f"{i}. {p['name']} (PID {p['pid']}) CPU {p['cpu']}% 内存 {p['mem_mb']}MB"
|
||
for i, p in enumerate(s["top"], 1)
|
||
]
|
||
return templates.render("top", top_n=len(s["top"]), top_list="\n".join(lines))
|
||
|
||
|
||
def fmt_temps(s: dict) -> str:
|
||
if not s["temps"]:
|
||
return templates.render("temps", temp_list="未获取到 (需要 lm-sensors，虚拟机通常不支持)")
|
||
lines = [
|
||
f" {name}/{t['label']}: {t['current']}C"
|
||
for name, entries in s["temps"].items() for t in entries
|
||
]
|
||
return templates.render("temps", temp_list="\n".join(lines))
|
||
|
||
|
||
def fmt_net(net: dict) -> str:
|
||
return templates.render(
|
||
"net", up=f"{net['up_mbps']}Mbps", down=f"{net['down_mbps']}Mbps")
|
||
|
||
|
||
# ---------- 命令处理 ----------
|
||
|
||
# 命令别名 -> 动作；只有此表内的词会触发，其他消息一律忽略
|
||
COMMANDS = {
|
||
"status": "status", "状态": "status",
|
||
"top": "top", "进程": "top",
|
||
"temp": "temp", "温度": "temp",
|
||
"net": "net", "网络": "net",
|
||
"help": "help", "帮助": "help",
|
||
}
|
||
|
||
|
||
class Bot:
|
||
def __init__(self, cfg: dict):
|
||
self.cfg = cfg
|
||
self.ob = OneBotServer(
|
||
cfg["server"]["host"], cfg["server"]["port"],
|
||
cfg["server"].get("access_token", ""),
|
||
)
|
||
self._self_id: int | None = None # 机器人自己的 QQ 号，群聊 @ 校验用
|
||
|
||
async def on_event(self, ev: dict):
|
||
if ev.get("post_type") != "message":
|
||
return
|
||
msg_type = ev.get("message_type")
|
||
if msg_type == "group":
|
||
# 群聊必须 @ 机器人本人才响应
|
||
if not await self._at_me(ev.get("message")):
|
||
return
|
||
text = self._text_after_at(ev.get("message"))
|
||
elif msg_type == "private":
|
||
text = self._extract_text(ev.get("message"))
|
||
if not text: # 纯图片/表情等无文本消息不响应
|
||
return
|
||
else:
|
||
return
|
||
|
||
reply = await self.handle_command(text)
|
||
if reply is None:
|
||
return
|
||
if msg_type == "group":
|
||
await self.ob.send_group_msg(ev["group_id"], reply)
|
||
else:
|
||
await self.ob.send_private_msg(ev["user_id"], reply)
|
||
|
||
@staticmethod
|
||
def _extract_text(message) -> str:
|
||
if isinstance(message, str):
|
||
return message.strip()
|
||
if isinstance(message, list): # CQ 段数组
|
||
return "".join(
|
||
seg.get("data", {}).get("text", "")
|
||
for seg in message if isinstance(seg, dict) and seg.get("type") == "text"
|
||
).strip()
|
||
return ""
|
||
|
||
async def _get_self_id(self) -> int | None:
|
||
"""获取并缓存机器人自己的 QQ 号。"""
|
||
if self._self_id is None:
|
||
try:
|
||
info = await self.ob.call_api("get_login_info")
|
||
self._self_id = int(info.get("user_id", 0)) or None
|
||
except Exception:
|
||
log.exception("获取机器人登录信息失败")
|
||
return self._self_id
|
||
|
||
async def _at_me(self, message) -> bool:
|
||
"""校验消息中 @ 的是机器人本人（@其他人/全体不算）。"""
|
||
if self._self_id is None:
|
||
await self._get_self_id()
|
||
ats: list[str] = []
|
||
if isinstance(message, list):
|
||
ats = [str(seg.get("data", {}).get("qq", ""))
|
||
for seg in message if isinstance(seg, dict) and seg.get("type") == "at"]
|
||
else:
|
||
ats = re.findall(r"\[CQ:at,qq=(\d+)", message or "")
|
||
return str(self._self_id or "") in ats
|
||
|
||
def _text_after_at(self, message) -> str:
|
||
"""取 @ 之后剩余的文本。"""
|
||
if isinstance(message, list):
|
||
after = False
|
||
parts = []
|
||
for seg in message:
|
||
if not isinstance(seg, dict):
|
||
continue
|
||
if seg.get("type") == "at":
|
||
after = True
|
||
continue
|
||
if after and seg.get("type") == "text":
|
||
parts.append(seg.get("data", {}).get("text", ""))
|
||
return "".join(parts).strip()
|
||
if isinstance(message, str) and "[CQ:at" in message:
|
||
return message.split("]", 1)[-1].strip()
|
||
return self._extract_text(message)
|
||
|
||
async def handle_command(self, text: str) -> str | None:
|
||
cmd = text.lstrip("/!").strip().lower()
|
||
cmd = cmd.split()[0] if cmd else ""
|
||
action = COMMANDS.get(cmd)
|
||
if action is None: # 非命令消息一律忽略
|
||
return None
|
||
try:
|
||
if action == "status":
|
||
return fmt_status(monitor.get_status())
|
||
if action == "top":
|
||
return fmt_top({"top": monitor.get_top_processes(10)})
|
||
if action == "temp":
|
||
return fmt_temps({"temps": monitor.get_temps()})
|
||
if action == "net":
|
||
return fmt_net(monitor.get_net(interval=1.0))
|
||
if action == "help":
|
||
return templates.render("help")
|
||
return None
|
||
except Exception:
|
||
log.exception("处理命令失败: %s", text)
|
||
return "[错误] 命令执行出错，请查看服务端日志。"
|
||
|
||
async def run(self):
|
||
await self.ob.start(self.on_event)
|
||
alerter = Alerter(self.ob, self.cfg.get("alert", {}))
|
||
alerter.start()
|
||
reporter = Reporter(self.ob, self.cfg.get("report", {}))
|
||
reporter.start()
|
||
updater = Updater(self.ob, self.cfg.get("update", {}))
|
||
if self.cfg.get("update", {}).get("enabled", False):
|
||
updater.start()
|
||
watcher = asyncio.get_running_loop().create_task(self._watch_config(alerter))
|
||
log.info("机器人就绪，等待 OneBot 实现端连接...")
|
||
try:
|
||
await self._wait_for_stop()
|
||
finally:
|
||
log.info("正在关闭...")
|
||
watcher.cancel()
|
||
await updater.stop()
|
||
await reporter.stop()
|
||
await alerter.stop()
|
||
await self.ob.stop()
|
||
log.info("已退出")
|
||
|
||
async def _wait_for_stop(self):
|
||
"""阻塞直到收到 Ctrl-C / SIGTERM。"""
|
||
stop = asyncio.Event()
|
||
loop = asyncio.get_running_loop()
|
||
|
||
def request_stop():
|
||
loop.call_soon_threadsafe(stop.set)
|
||
|
||
sigs = [signal.SIGINT, signal.SIGTERM]
|
||
if sys.platform == "win32" and hasattr(signal, "SIGBREAK"):
|
||
sigs.append(signal.SIGBREAK) # CTRL_BREAK / 任务管理器停止服务
|
||
for sig in sigs:
|
||
try:
|
||
loop.add_signal_handler(sig, request_stop)
|
||
except NotImplementedError: # Windows 回退
|
||
signal.signal(sig, lambda s, f: request_stop())
|
||
try:
|
||
await stop.wait()
|
||
except KeyboardInterrupt: # Windows 上 Ctrl-C 可能直接中断 await
|
||
pass
|
||
|
||
async def _watch_config(self, alerter: Alerter, interval: float = 5.0):
|
||
"""监测 config.toml 修改并热重载（端口/token 等需重启的项不生效）。"""
|
||
path = Path("config.toml")
|
||
last = path.stat().st_mtime
|
||
while True:
|
||
await asyncio.sleep(interval)
|
||
try:
|
||
mtime = path.stat().st_mtime
|
||
if mtime == last:
|
||
continue
|
||
last = mtime
|
||
cfg = load_config()
|
||
alerter.apply_config(cfg.get("alert", {}))
|
||
log.info("配置已热重载 (端口/access_token 变更需重启生效)")
|
||
except FileNotFoundError:
|
||
continue
|
||
except Exception:
|
||
log.exception("配置热重载失败")
|
||
|
||
|
||
def main():
|
||
if sys.platform == "win32":
|
||
asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
|
||
cfg = load_config()
|
||
try:
|
||
asyncio.run(Bot(cfg).run())
|
||
except KeyboardInterrupt:
|
||
print("已退出")
|
||
sys.exit(0)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|