test_bot_config.py 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314
  1. from __future__ import annotations
  2. import json
  3. import pytest
  4. from aiohttp import CookieJar
  5. from aiohttp.test_utils import TestClient, TestServer
  6. def test_multiple_bot_profiles_are_atomic_and_redacted(app_modules, tmp_path):
  7. config = app_modules.load("wbb.admin.bot_config")
  8. path = tmp_path / "bots.json"
  9. first_token = "123456:" + "A" * 30
  10. second_token = "654321:" + "B" * 30
  11. config.update_telegram_config(path, {"api_id": 12345, "api_hash": "a" * 32})
  12. first = config.create_bot_profile(
  13. path,
  14. {
  15. "label": "Primary",
  16. "bot_token": first_token,
  17. "sudo_users_id": "100 200 100",
  18. "message_dump_chat": -100123,
  19. "business_assistant": {
  20. "base_url": "https://models.example.com/v1",
  21. "api_key": "secret-assistant-key",
  22. "model": "support-model",
  23. },
  24. },
  25. )
  26. second = config.create_bot_profile(
  27. path,
  28. {"label": "Secondary", "bot_token": second_token, "enabled": False},
  29. )
  30. status = config.telegram_config_status(path)
  31. assert status["api_ready"] is True
  32. assert len(status["bots"]) == 2
  33. assert status["bots"][0]["sudo_users_id"] == ["100", "200"]
  34. assert status["bots"][0]["role_ids"] == ["full_access"]
  35. assert status["bots"][0]["permissions"]
  36. assert "teacher_directory.manage" in status["bots"][0]["permissions"]
  37. assert status["bots"][0]["business_assistant"] == {
  38. "base_url": "https://models.example.com/v1",
  39. "api_key_configured": True,
  40. "model": "support-model",
  41. "timeout_seconds": 30,
  42. "max_output_tokens": 600,
  43. "configured": True,
  44. }
  45. serialized = json.dumps(status)
  46. assert first_token not in serialized
  47. assert second_token not in serialized
  48. assert "a" * 32 not in serialized
  49. assert "secret-assistant-key" not in serialized
  50. assert path.stat().st_mode & 0o777 == 0o600
  51. config.update_bot_profile(path, first["bot_id"], {"label": "Renamed", "bot_token": ""})
  52. stored = config.get_bot_profile_secrets(path, first["bot_id"])
  53. assert stored["bot_token"] == first_token
  54. assert stored["label"] == "Renamed"
  55. config.delete_bot_profile(path, second["bot_id"])
  56. assert len(config.telegram_config_status(path)["bots"]) == 1
  57. def test_custom_roles_are_assignable_and_permissions_are_merged(app_modules, tmp_path):
  58. config = app_modules.load("wbb.admin.bot_config")
  59. path = tmp_path / "bots.json"
  60. config.update_telegram_config(path, {"api_id": 12345, "api_hash": "a" * 32})
  61. role = config.create_bot_role(
  62. path,
  63. {
  64. "name": "积分与抽奖",
  65. "description": "社区运营",
  66. "permissions": ["points.manage", "giveaways.manage"],
  67. },
  68. )
  69. profile = config.create_bot_profile(
  70. path,
  71. {
  72. "label": "Operator",
  73. "bot_token": "123456:" + "A" * 30,
  74. "role_ids": [role["role_id"], "karma_manager"],
  75. },
  76. )
  77. assert profile["role_ids"] == [role["role_id"], "karma_manager"]
  78. assert set(profile["permissions"]) == {
  79. "points.manage",
  80. "giveaways.manage",
  81. "karma.manage",
  82. }
  83. config.update_bot_role(
  84. path,
  85. role["role_id"],
  86. {"permissions": ["automation.manage"]},
  87. )
  88. updated = config.telegram_config_status(path)["bots"][0]
  89. assert set(updated["permissions"]) == {"automation.manage", "karma.manage"}
  90. with pytest.raises(config.BotConfigError) as error:
  91. config.delete_bot_role(path, role["role_id"])
  92. assert error.value.code == "role_in_use"
  93. config.update_bot_profile(path, profile["bot_id"], {"role_ids": []})
  94. config.delete_bot_role(path, role["role_id"])
  95. assert all(
  96. item["role_id"] != role["role_id"]
  97. for item in config.telegram_config_status(path)["roles"]
  98. )
  99. def test_bot_profile_validation(app_modules, tmp_path):
  100. config = app_modules.load("wbb.admin.bot_config")
  101. path = tmp_path / "bots.json"
  102. with pytest.raises(config.BotConfigError) as error:
  103. config.create_bot_profile(path, {"label": "Bad", "bot_token": "not-a-token"})
  104. assert error.value.code == "invalid_bot_token"
  105. with pytest.raises(config.BotConfigError) as error:
  106. config.update_telegram_config(
  107. path,
  108. {"api_id": 12345, "api_hash": "123456:" + "A" * 30},
  109. )
  110. assert error.value.code == "invalid_api_hash"
  111. assert "不能使用机器人令牌" in str(error.value)
  112. def test_role_permissions_gate_telegram_modules(app_modules):
  113. permissions = app_modules.load("wbb.services.bot_permissions")
  114. assert permissions.module_allowed("points", {"points.manage"}) is True
  115. assert permissions.module_allowed("points", {"giveaways.manage"}) is False
  116. assert permissions.module_allowed("admin", {"chat.members"}) is True
  117. assert permissions.module_allowed("admin", {"chat.invites"}) is True
  118. assert permissions.module_allowed("admin", set()) is False
  119. assert permissions.module_allowed("chat_watcher", set()) is True
  120. assert permissions.TELEGRAM_COMMAND_PERMISSIONS["ban"] == "chat.members"
  121. assert permissions.TELEGRAM_COMMAND_PERMISSIONS["pin"] == "chat.announcements"
  122. assert permissions.TELEGRAM_COMMAND_PERMISSIONS["invite"] == "chat.invites"
  123. def test_supervisor_assigns_each_worker_an_isolated_database(app_modules, tmp_path):
  124. supervisor_module = app_modules.load("wbb.admin.supervisor")
  125. supervisor = supervisor_module.BotSupervisor(tmp_path / "bots.json", project_root=tmp_path)
  126. profile = {
  127. "bot_id": "bot.one/secondary",
  128. "bot_token": "123456:" + "A" * 30,
  129. "api_id": 12345,
  130. "api_hash": "a" * 32,
  131. "sudo_users_id": [],
  132. "log_group_id": 0,
  133. "gban_log_group_id": 0,
  134. "message_dump_chat": 0,
  135. "permissions": ["points.manage", "giveaways.manage"],
  136. "business_assistant": {
  137. "base_url": "https://models.example.com/v1",
  138. "api_key": "worker-only-key",
  139. "model": "support-model",
  140. },
  141. }
  142. environment = supervisor._worker_environment(profile, 18088)
  143. assert environment["WBB_BOT_PROFILE_ID"] == "bot.one/secondary"
  144. assert environment["WBB_BOT_DATABASE"] == "wbb_bot_bot_one_secondary"
  145. assert environment["WBB_BOT_PERMISSIONS"] == "points.manage,giveaways.manage"
  146. assert environment["BUSINESS_ASSISTANT_OPENAI_API_KEY"] == "worker-only-key"
  147. assert environment["BUSINESS_ASSISTANT_OPENAI_MODEL"] == "support-model"
  148. async def test_bot_admin_api_crud_and_token_test_are_redacted(app_modules, monkeypatch):
  149. admin_api = app_modules.load("wbb.admin.api")
  150. application = admin_api.build_admin_application()
  151. await application["admin_api"].initialize()
  152. client = TestClient(TestServer(application), cookie_jar=CookieJar(unsafe=True))
  153. await client.start_server()
  154. token = "123456:" + "C" * 30
  155. try:
  156. login = await client.post(
  157. "/api/admin/v1/auth/login",
  158. json={"username": "admin", "password": "qwe0.123456"},
  159. )
  160. login_data = (await login.json())["data"]
  161. changed = await client.put(
  162. "/api/admin/v1/auth/password",
  163. headers={"X-CSRF-Token": login_data["csrf_token"]},
  164. json={
  165. "current_password": "qwe0.123456",
  166. "new_password": "changed-pass-123",
  167. },
  168. )
  169. csrf = (await changed.json())["data"]["csrf_token"]
  170. headers = {"X-CSRF-Token": csrf}
  171. telegram = await client.put(
  172. "/api/admin/v1/settings/telegram",
  173. headers=headers,
  174. json={"api_id": 12345, "api_hash": "d" * 32, "confirm": True},
  175. )
  176. assert telegram.status == 200
  177. assert (await telegram.json())["data"]["api_hash_configured"] is True
  178. created_role = await client.post(
  179. "/api/admin/v1/roles",
  180. headers=headers,
  181. json={
  182. "name": "积分专员",
  183. "description": "只负责积分",
  184. "permissions": ["points.manage"],
  185. "confirm": True,
  186. },
  187. )
  188. assert created_role.status == 201
  189. role = (await created_role.json())["data"]
  190. created = await client.post(
  191. "/api/admin/v1/bots",
  192. headers=headers,
  193. json={
  194. "label": "Test Bot",
  195. "bot_token": token,
  196. "enabled": False,
  197. "role_ids": [role["role_id"]],
  198. "confirm": True,
  199. },
  200. )
  201. assert created.status == 201
  202. created_payload = await created.json()
  203. profile = created_payload["data"]
  204. assert profile["bot_token_configured"] is True
  205. assert profile["permissions"] == ["points.manage"]
  206. assert token not in json.dumps(created_payload)
  207. async def fake_test_bot_token(_token: str):
  208. assert _token == token
  209. return {
  210. "id": "123456",
  211. "username": "test_bot",
  212. "name": "Test Bot",
  213. "can_connect_to_business": True,
  214. }
  215. monkeypatch.setattr(admin_api, "test_bot_token", fake_test_bot_token)
  216. tested = await client.post(
  217. f"/api/admin/v1/bots/{profile['bot_id']}/test",
  218. headers=headers,
  219. json={"confirm": True},
  220. )
  221. assert tested.status == 200
  222. assert (await tested.json())["data"]["username"] == "test_bot"
  223. assert (await tested.json())["data"]["can_connect_to_business"] is True
  224. listed = await client.get("/api/admin/v1/bots")
  225. listed_payload = await listed.json()
  226. assert listed_payload["data"]["items"][0]["identity"]["username"] == "test_bot"
  227. assert token not in json.dumps(listed_payload)
  228. deleted = await client.delete(
  229. f"/api/admin/v1/bots/{profile['bot_id']}",
  230. headers=headers,
  231. json={"confirm": True},
  232. )
  233. assert deleted.status == 200
  234. deleted_role = await client.delete(
  235. f"/api/admin/v1/roles/{role['role_id']}",
  236. headers=headers,
  237. json={"confirm": True},
  238. )
  239. assert deleted_role.status == 200
  240. finally:
  241. await client.close()
  242. async def test_bot_role_permission_is_enforced_by_worker_api(app_modules):
  243. admin_api = app_modules.load("wbb.admin.api")
  244. app_modules.wbb.BOT_PERMISSIONS = set()
  245. application = admin_api.build_admin_application()
  246. await application["admin_api"].initialize()
  247. client = TestClient(TestServer(application), cookie_jar=CookieJar(unsafe=True))
  248. await client.start_server()
  249. try:
  250. login = await client.post(
  251. "/api/admin/v1/auth/login",
  252. json={"username": "admin", "password": "qwe0.123456"},
  253. )
  254. login_data = (await login.json())["data"]
  255. changed = await client.put(
  256. "/api/admin/v1/auth/password",
  257. headers={"X-CSRF-Token": login_data["csrf_token"]},
  258. json={
  259. "current_password": "qwe0.123456",
  260. "new_password": "changed-pass-123",
  261. },
  262. )
  263. assert changed.status == 200
  264. denied = await client.get(
  265. "/api/admin/v1/points/accounts?chat_id=-100"
  266. )
  267. assert denied.status == 403
  268. payload = await denied.json()
  269. assert payload["error"]["code"] == "bot_role_permission_denied"
  270. assert payload["error"]["details"]["required_permission"] == "points.manage"
  271. app_modules.wbb.BOT_PERMISSIONS = {"points.manage"}
  272. allowed = await client.get(
  273. "/api/admin/v1/points/accounts?chat_id=-100"
  274. )
  275. assert allowed.status == 200
  276. finally:
  277. await client.close()