bot_config.py 18 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512
  1. from __future__ import annotations
  2. import json
  3. import re
  4. import tempfile
  5. import threading
  6. from datetime import UTC, datetime
  7. from pathlib import Path
  8. from typing import Any
  9. from uuid import uuid4
  10. from aiohttp import ClientError, ClientSession, ClientTimeout
  11. from wbb.services.bot_permissions import (
  12. ALL_BOT_PERMISSIONS,
  13. builtin_roles,
  14. normalize_permissions,
  15. permission_catalog,
  16. )
  17. BOT_TOKEN_PATTERN = re.compile(r"^\d{5,}:[A-Za-z0-9_-]{20,}$")
  18. API_HASH_PATTERN = re.compile(r"^[A-Fa-f0-9]{32}$")
  19. PROFILE_FIELDS = {
  20. "sudo_users_id",
  21. "log_group_id",
  22. "gban_log_group_id",
  23. "message_dump_chat",
  24. }
  25. _STORE_LOCK = threading.RLock()
  26. class BotConfigError(ValueError):
  27. def __init__(self, code: str, message: str):
  28. super().__init__(message)
  29. self.code = code
  30. def _empty_document() -> dict[str, Any]:
  31. return {
  32. "version": 2,
  33. "telegram": {"api_id": 0, "api_hash": ""},
  34. "roles": [],
  35. "bots": [],
  36. }
  37. def _read_document(path: str | Path) -> dict[str, Any]:
  38. config_path = Path(path)
  39. if not config_path.exists():
  40. return _empty_document()
  41. try:
  42. data = json.loads(config_path.read_text(encoding="utf-8"))
  43. except (OSError, json.JSONDecodeError) as exc:
  44. raise BotConfigError(
  45. "bot_config_corrupt",
  46. "机器人配置文件无法读取。",
  47. ) from exc
  48. if not isinstance(data, dict) or not isinstance(data.get("bots", []), list):
  49. raise BotConfigError("bot_config_corrupt", "机器人配置文件结构无效。")
  50. telegram = data.get("telegram")
  51. if not isinstance(telegram, dict):
  52. telegram = {}
  53. return {
  54. "version": 2,
  55. "telegram": {
  56. "api_id": int(telegram.get("api_id") or 0),
  57. "api_hash": str(telegram.get("api_hash") or ""),
  58. },
  59. "roles": [item for item in data.get("roles", []) if isinstance(item, dict)],
  60. "bots": [item for item in data.get("bots", []) if isinstance(item, dict)],
  61. }
  62. def _write_document(path: str | Path, data: dict[str, Any]) -> None:
  63. config_path = Path(path).resolve()
  64. config_path.parent.mkdir(parents=True, exist_ok=True)
  65. with tempfile.NamedTemporaryFile(
  66. mode="w",
  67. encoding="utf-8",
  68. dir=config_path.parent,
  69. delete=False,
  70. ) as handle:
  71. temporary_path = Path(handle.name)
  72. json.dump(data, handle, ensure_ascii=True, indent=2)
  73. handle.write("\n")
  74. try:
  75. temporary_path.chmod(0o600)
  76. temporary_path.replace(config_path)
  77. config_path.chmod(0o600)
  78. finally:
  79. temporary_path.unlink(missing_ok=True)
  80. def _integer(value: Any, name: str, *, positive: bool = False) -> int:
  81. try:
  82. parsed = int(value)
  83. except (TypeError, ValueError) as exc:
  84. raise BotConfigError("invalid_bot_config", f"{name} 必须是整数。") from exc
  85. if positive and parsed <= 0:
  86. raise BotConfigError("invalid_bot_config", f"{name} 必须大于 0。")
  87. return parsed
  88. def _normalize_sudoers(value: Any) -> list[int]:
  89. items = value if isinstance(value, list) else str(value or "").replace(",", " ").split()
  90. normalized = [_integer(item, "SUDOERS 用户 ID", positive=True) for item in items]
  91. return list(dict.fromkeys(normalized))
  92. def _find_profile(data: dict[str, Any], bot_id: str) -> dict[str, Any]:
  93. profile = next(
  94. (item for item in data["bots"] if str(item.get("bot_id")) == str(bot_id)),
  95. None,
  96. )
  97. if profile is None:
  98. raise BotConfigError("bot_not_found", "未找到该机器人配置。")
  99. return profile
  100. def _all_roles(data: dict[str, Any]) -> list[dict[str, Any]]:
  101. roles = builtin_roles()
  102. builtin_ids = {str(item["role_id"]) for item in roles}
  103. roles.extend(
  104. {
  105. "role_id": str(item.get("role_id") or ""),
  106. "name": str(item.get("name") or ""),
  107. "description": str(item.get("description") or ""),
  108. "permissions": [
  109. permission
  110. for permission in item.get("permissions", [])
  111. if permission in ALL_BOT_PERMISSIONS
  112. ],
  113. "builtin": False,
  114. "created_at": item.get("created_at"),
  115. "updated_at": item.get("updated_at"),
  116. }
  117. for item in data.get("roles", [])
  118. if str(item.get("role_id") or "") not in builtin_ids
  119. )
  120. return roles
  121. def _find_role(data: dict[str, Any], role_id: str) -> dict[str, Any]:
  122. role = next(
  123. (item for item in _all_roles(data) if item["role_id"] == str(role_id)),
  124. None,
  125. )
  126. if role is None:
  127. raise BotConfigError("role_not_found", "未找到该机器人角色。")
  128. return role
  129. def _normalize_role_ids(
  130. value: Any,
  131. data: dict[str, Any],
  132. *,
  133. legacy_default: bool = False,
  134. ) -> list[str]:
  135. if value is None and legacy_default:
  136. return ["full_access"]
  137. items = value if isinstance(value, list) else str(value or "").replace(",", " ").split()
  138. normalized = list(dict.fromkeys(str(item).strip() for item in items if str(item).strip()))
  139. known = {str(item["role_id"]) for item in _all_roles(data)}
  140. unknown = sorted(set(normalized) - known)
  141. if unknown:
  142. raise BotConfigError(
  143. "invalid_bot_roles",
  144. f"包含不存在的机器人角色:{', '.join(unknown)}。",
  145. )
  146. return normalized
  147. def _profile_role_ids(profile: dict[str, Any], data: dict[str, Any]) -> list[str]:
  148. raw = profile.get("role_ids")
  149. if raw is None:
  150. return ["full_access"]
  151. values = raw if isinstance(raw, list) else str(raw or "").replace(",", " ").split()
  152. known = {str(item["role_id"]) for item in _all_roles(data)}
  153. return [
  154. item
  155. for item in dict.fromkeys(str(value) for value in values)
  156. if item in known
  157. ]
  158. def _effective_permissions(
  159. profile: dict[str, Any],
  160. data: dict[str, Any],
  161. ) -> list[str]:
  162. roles = {str(item["role_id"]): item for item in _all_roles(data)}
  163. granted: set[str] = set()
  164. for role_id in _profile_role_ids(profile, data):
  165. granted.update(roles[role_id].get("permissions", []))
  166. return [
  167. item["key"]
  168. for item in permission_catalog()
  169. if item["key"] in granted
  170. ]
  171. def _public_profile(
  172. profile: dict[str, Any],
  173. *,
  174. data: dict[str, Any],
  175. api_ready: bool,
  176. runtime: dict[str, Any] | None = None,
  177. ) -> dict[str, Any]:
  178. token_configured = bool(BOT_TOKEN_PATTERN.fullmatch(str(profile.get("bot_token") or "")))
  179. identity = profile.get("identity") if isinstance(profile.get("identity"), dict) else None
  180. return {
  181. "bot_id": str(profile.get("bot_id") or ""),
  182. "label": str(profile.get("label") or ""),
  183. "enabled": bool(profile.get("enabled", True)),
  184. "bot_token_configured": token_configured,
  185. "sudo_users_id": [str(item) for item in profile.get("sudo_users_id", [])],
  186. "log_group_id": str(profile.get("log_group_id") or 0),
  187. "gban_log_group_id": str(profile.get("gban_log_group_id") or 0),
  188. "message_dump_chat": str(profile.get("message_dump_chat") or 0),
  189. "role_ids": _profile_role_ids(profile, data),
  190. "permissions": _effective_permissions(profile, data),
  191. "identity": identity,
  192. "ready_to_connect": api_ready and token_configured,
  193. "created_at": profile.get("created_at"),
  194. "updated_at": profile.get("updated_at"),
  195. "runtime": runtime or {"state": "stopped"},
  196. }
  197. def telegram_config_status(
  198. path: str | Path,
  199. *,
  200. runtimes: dict[str, dict[str, Any]] | None = None,
  201. ) -> dict[str, Any]:
  202. with _STORE_LOCK:
  203. data = _read_document(path)
  204. telegram = data["telegram"]
  205. api_id = int(telegram.get("api_id") or 0)
  206. api_hash_configured = bool(API_HASH_PATTERN.fullmatch(str(telegram.get("api_hash") or "")))
  207. api_ready = api_id > 0 and api_hash_configured
  208. runtime_map = runtimes or {}
  209. return {
  210. "api_id": api_id or None,
  211. "api_hash_configured": api_hash_configured,
  212. "api_ready": api_ready,
  213. "roles": _all_roles(data),
  214. "permission_catalog": permission_catalog(),
  215. "bots": [
  216. _public_profile(
  217. profile,
  218. data=data,
  219. api_ready=api_ready,
  220. runtime=runtime_map.get(str(profile.get("bot_id"))),
  221. )
  222. for profile in data["bots"]
  223. ],
  224. }
  225. def update_telegram_config(path: str | Path, body: dict[str, Any]) -> dict[str, Any]:
  226. with _STORE_LOCK:
  227. data = _read_document(path)
  228. telegram = data["telegram"]
  229. changed = False
  230. if "api_id" in body:
  231. telegram["api_id"] = _integer(
  232. body.get("api_id"), "Telegram App API ID", positive=True
  233. )
  234. changed = True
  235. api_hash = str(body.get("api_hash") or "").strip()
  236. if api_hash:
  237. if not API_HASH_PATTERN.fullmatch(api_hash):
  238. raise BotConfigError(
  239. "invalid_api_hash",
  240. "Telegram 应用密钥必须是 my.telegram.org/apps 提供的 "
  241. "32 位十六进制字符串,不能使用机器人令牌。",
  242. )
  243. telegram["api_hash"] = api_hash
  244. changed = True
  245. if not changed:
  246. raise BotConfigError(
  247. "bot_config_unchanged",
  248. "没有可保存的 Telegram 应用接口配置。",
  249. )
  250. _write_document(path, data)
  251. return telegram_config_status(path)
  252. def create_bot_profile(path: str | Path, body: dict[str, Any]) -> dict[str, Any]:
  253. label = str(body.get("label") or "").strip()
  254. token = str(body.get("bot_token") or "").strip()
  255. if not label or len(label) > 60:
  256. raise BotConfigError(
  257. "invalid_bot_label",
  258. "机器人名称长度需要在 1 到 60 个字符之间。",
  259. )
  260. if not BOT_TOKEN_PATTERN.fullmatch(token):
  261. raise BotConfigError("invalid_bot_token", "机器人令牌格式无效。")
  262. with _STORE_LOCK:
  263. data = _read_document(path)
  264. now = datetime.now(UTC).isoformat().replace("+00:00", "Z")
  265. profile = {
  266. "bot_id": uuid4().hex,
  267. "label": label,
  268. "bot_token": token,
  269. "enabled": bool(body.get("enabled", True)),
  270. "sudo_users_id": _normalize_sudoers(body.get("sudo_users_id", [])),
  271. "log_group_id": _integer(body.get("log_group_id", 0), "日志群 ID"),
  272. "gban_log_group_id": _integer(
  273. body.get("gban_log_group_id", 0), "全局封禁日志群 ID"
  274. ),
  275. "message_dump_chat": _integer(
  276. body.get("message_dump_chat", 0), "媒体中转群 ID"
  277. ),
  278. "role_ids": _normalize_role_ids(
  279. body.get("role_ids"),
  280. data,
  281. legacy_default="role_ids" not in body,
  282. ),
  283. "identity": None,
  284. "created_at": now,
  285. "updated_at": now,
  286. }
  287. data["bots"].append(profile)
  288. _write_document(path, data)
  289. api_ready = bool(data["telegram"]["api_id"] and data["telegram"]["api_hash"])
  290. return _public_profile(profile, data=data, api_ready=api_ready)
  291. def update_bot_profile(path: str | Path, bot_id: str, body: dict[str, Any]) -> dict[str, Any]:
  292. with _STORE_LOCK:
  293. data = _read_document(path)
  294. profile = _find_profile(data, bot_id)
  295. if "label" in body:
  296. label = str(body.get("label") or "").strip()
  297. if not label or len(label) > 60:
  298. raise BotConfigError(
  299. "invalid_bot_label",
  300. "机器人名称长度需要在 1 到 60 个字符之间。",
  301. )
  302. profile["label"] = label
  303. token = str(body.get("bot_token") or "").strip()
  304. if token:
  305. if not BOT_TOKEN_PATTERN.fullmatch(token):
  306. raise BotConfigError("invalid_bot_token", "机器人令牌格式无效。")
  307. profile["bot_token"] = token
  308. profile["identity"] = None
  309. if "enabled" in body:
  310. profile["enabled"] = bool(body.get("enabled"))
  311. if "sudo_users_id" in body:
  312. profile["sudo_users_id"] = _normalize_sudoers(body.get("sudo_users_id"))
  313. if "role_ids" in body:
  314. profile["role_ids"] = _normalize_role_ids(body.get("role_ids"), data)
  315. for field in PROFILE_FIELDS - {"sudo_users_id"}:
  316. if field in body:
  317. profile[field] = _integer(body.get(field), field)
  318. profile["updated_at"] = datetime.now(UTC).isoformat().replace("+00:00", "Z")
  319. _write_document(path, data)
  320. api_ready = bool(data["telegram"]["api_id"] and data["telegram"]["api_hash"])
  321. return _public_profile(profile, data=data, api_ready=api_ready)
  322. def delete_bot_profile(path: str | Path, bot_id: str) -> None:
  323. with _STORE_LOCK:
  324. data = _read_document(path)
  325. _find_profile(data, bot_id)
  326. data["bots"] = [
  327. item for item in data["bots"] if str(item.get("bot_id")) != str(bot_id)
  328. ]
  329. _write_document(path, data)
  330. def get_bot_profile_secrets(path: str | Path, bot_id: str) -> dict[str, Any]:
  331. with _STORE_LOCK:
  332. data = _read_document(path)
  333. profile = dict(_find_profile(data, bot_id))
  334. profile["role_ids"] = _profile_role_ids(profile, data)
  335. profile["permissions"] = _effective_permissions(profile, data)
  336. telegram = dict(data["telegram"])
  337. return {**profile, **telegram}
  338. def create_bot_role(path: str | Path, body: dict[str, Any]) -> dict[str, Any]:
  339. name = str(body.get("name") or "").strip()
  340. description = str(body.get("description") or "").strip()
  341. if not name or len(name) > 60:
  342. raise BotConfigError(
  343. "invalid_role_name",
  344. "角色名称长度需要在 1 到 60 个字符之间。",
  345. )
  346. if len(description) > 200:
  347. raise BotConfigError("invalid_role_description", "角色说明不能超过 200 个字符。")
  348. try:
  349. permissions = normalize_permissions(body.get("permissions", []))
  350. except ValueError as exc:
  351. raise BotConfigError("invalid_role_permissions", str(exc)) from exc
  352. now = datetime.now(UTC).isoformat().replace("+00:00", "Z")
  353. role = {
  354. "role_id": uuid4().hex,
  355. "name": name,
  356. "description": description,
  357. "permissions": permissions,
  358. "created_at": now,
  359. "updated_at": now,
  360. }
  361. with _STORE_LOCK:
  362. data = _read_document(path)
  363. data["roles"].append(role)
  364. _write_document(path, data)
  365. return {**role, "builtin": False}
  366. def update_bot_role(
  367. path: str | Path,
  368. role_id: str,
  369. body: dict[str, Any],
  370. ) -> dict[str, Any]:
  371. with _STORE_LOCK:
  372. data = _read_document(path)
  373. current = _find_role(data, role_id)
  374. if current.get("builtin"):
  375. raise BotConfigError("builtin_role_immutable", "内置角色不能修改。")
  376. role = next(
  377. item
  378. for item in data["roles"]
  379. if str(item.get("role_id")) == str(role_id)
  380. )
  381. if "name" in body:
  382. name = str(body.get("name") or "").strip()
  383. if not name or len(name) > 60:
  384. raise BotConfigError(
  385. "invalid_role_name",
  386. "角色名称长度需要在 1 到 60 个字符之间。",
  387. )
  388. role["name"] = name
  389. if "description" in body:
  390. description = str(body.get("description") or "").strip()
  391. if len(description) > 200:
  392. raise BotConfigError(
  393. "invalid_role_description",
  394. "角色说明不能超过 200 个字符。",
  395. )
  396. role["description"] = description
  397. if "permissions" in body:
  398. try:
  399. role["permissions"] = normalize_permissions(body.get("permissions"))
  400. except ValueError as exc:
  401. raise BotConfigError("invalid_role_permissions", str(exc)) from exc
  402. role["updated_at"] = datetime.now(UTC).isoformat().replace("+00:00", "Z")
  403. _write_document(path, data)
  404. return {**role, "builtin": False}
  405. def delete_bot_role(path: str | Path, role_id: str) -> None:
  406. with _STORE_LOCK:
  407. data = _read_document(path)
  408. role = _find_role(data, role_id)
  409. if role.get("builtin"):
  410. raise BotConfigError("builtin_role_immutable", "内置角色不能删除。")
  411. assigned = [
  412. str(profile.get("label") or profile.get("bot_id"))
  413. for profile in data["bots"]
  414. if str(role_id) in _profile_role_ids(profile, data)
  415. ]
  416. if assigned:
  417. raise BotConfigError(
  418. "role_in_use",
  419. f"角色仍分配给以下机器人:{', '.join(assigned)}。",
  420. )
  421. data["roles"] = [
  422. item
  423. for item in data["roles"]
  424. if str(item.get("role_id")) != str(role_id)
  425. ]
  426. _write_document(path, data)
  427. def store_bot_identity(path: str | Path, bot_id: str, identity: dict[str, Any]) -> None:
  428. with _STORE_LOCK:
  429. data = _read_document(path)
  430. profile = _find_profile(data, bot_id)
  431. profile["identity"] = {
  432. "id": str(identity.get("id") or ""),
  433. "username": str(identity.get("username") or ""),
  434. "name": str(identity.get("name") or ""),
  435. }
  436. profile["updated_at"] = datetime.now(UTC).isoformat().replace("+00:00", "Z")
  437. _write_document(path, data)
  438. async def test_bot_token(token: str) -> dict[str, Any]:
  439. if not BOT_TOKEN_PATTERN.fullmatch(token):
  440. raise BotConfigError("invalid_bot_token", "机器人令牌格式无效。")
  441. timeout = ClientTimeout(total=12)
  442. try:
  443. async with ClientSession(timeout=timeout) as session:
  444. async with session.post(f"https://api.telegram.org/bot{token}/getMe") as response:
  445. payload = await response.json(content_type=None)
  446. except (ClientError, TimeoutError, ValueError) as exc:
  447. raise BotConfigError(
  448. "telegram_unreachable",
  449. "无法连接 Telegram 机器人接口。",
  450. ) from exc
  451. if response.status != 200 or not payload.get("ok"):
  452. raise BotConfigError("telegram_auth_failed", "机器人令牌验证失败。")
  453. result = payload.get("result") or {}
  454. return {
  455. "id": str(result.get("id") or ""),
  456. "username": result.get("username") or "",
  457. "name": " ".join(
  458. part for part in (result.get("first_name"), result.get("last_name")) if part
  459. ),
  460. }