directory.py 15 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475
  1. from __future__ import annotations
  2. import asyncio
  3. from contextlib import suppress
  4. from typing import Any
  5. from aiohttp import ClientError, ClientSession, ClientTimeout
  6. from pyrogram.enums import ChatMemberStatus
  7. import wbb
  8. from wbb import BOT_PROFILE_ID, SUDOERS, app, log
  9. from wbb.utils.dbdirectory import (
  10. DirectoryDataError,
  11. aware_utc,
  12. clear_directory_location,
  13. decide_teacher_application,
  14. expire_teacher_presence,
  15. get_directory_location,
  16. list_directory_teachers,
  17. list_membership_candidates,
  18. save_directory_location,
  19. set_geocode_consent,
  20. set_teacher_state,
  21. submit_teacher_application,
  22. teacher_is_online,
  23. upsert_directory_identity,
  24. upsert_directory_membership,
  25. )
  26. ACTIVE_MEMBER_STATUSES = {
  27. ChatMemberStatus.OWNER,
  28. ChatMemberStatus.ADMINISTRATOR,
  29. ChatMemberStatus.MEMBER,
  30. ChatMemberStatus.RESTRICTED,
  31. }
  32. ADMIN_MEMBER_STATUSES = {
  33. ChatMemberStatus.OWNER,
  34. ChatMemberStatus.ADMINISTRATOR,
  35. }
  36. _presence_task: asyncio.Task[None] | None = None
  37. class DirectoryServiceError(RuntimeError):
  38. def __init__(self, code: str, message: str, *, status: int = 400):
  39. super().__init__(message)
  40. self.code = code
  41. self.status = status
  42. def status_value(value: Any) -> str:
  43. return str(getattr(value, "value", value or "")).lower()
  44. def is_sudoer(user_id: int) -> bool:
  45. try:
  46. return int(user_id) in SUDOERS
  47. except (TypeError, ValueError):
  48. return False
  49. def user_identity(user: Any) -> dict[str, Any]:
  50. return {
  51. "user_id": int(user.id),
  52. "username": getattr(user, "username", None),
  53. "first_name": getattr(user, "first_name", None),
  54. "last_name": getattr(user, "last_name", None),
  55. "display_name": " ".join(
  56. value.strip()
  57. for value in (
  58. getattr(user, "first_name", None),
  59. getattr(user, "last_name", None),
  60. )
  61. if value and value.strip()
  62. )
  63. or (
  64. f"@{user.username}"
  65. if getattr(user, "username", None)
  66. else f"用户 {user.id}"
  67. ),
  68. }
  69. async def observe_directory_user(user: Any, *, source: str = "telegram") -> dict[str, Any]:
  70. identity = user_identity(user)
  71. return await upsert_directory_identity(
  72. user_id=identity["user_id"],
  73. username=identity["username"],
  74. first_name=identity["first_name"],
  75. last_name=identity["last_name"],
  76. source=source,
  77. )
  78. async def observe_group_member(
  79. *,
  80. chat_id: int,
  81. chat_title: str,
  82. user: Any,
  83. status: str = "member",
  84. active: bool = True,
  85. verified: bool = False,
  86. ) -> dict[str, Any]:
  87. identity = user_identity(user)
  88. await observe_directory_user(user)
  89. return await upsert_directory_membership(
  90. bot_id=str(BOT_PROFILE_ID),
  91. chat_id=int(chat_id),
  92. user_id=identity["user_id"],
  93. status=status,
  94. active=active,
  95. chat_title=chat_title,
  96. username=identity["username"],
  97. display_name_value=identity["display_name"],
  98. verified=verified,
  99. )
  100. async def verify_local_membership(
  101. *,
  102. user_id: int,
  103. chat_ids: list[int],
  104. require_admin: bool = False,
  105. ) -> dict[str, Any]:
  106. accepted = ADMIN_MEMBER_STATUSES if require_admin else ACTIVE_MEMBER_STATUSES
  107. for chat_id in dict.fromkeys(int(value) for value in chat_ids):
  108. try:
  109. member = await app.get_chat_member(int(chat_id), int(user_id))
  110. except Exception:
  111. continue
  112. status = getattr(member, "status", None)
  113. active = status in ACTIVE_MEMBER_STATUSES
  114. user = getattr(member, "user", None)
  115. if user is not None:
  116. with suppress(Exception):
  117. await observe_group_member(
  118. chat_id=int(chat_id),
  119. chat_title="",
  120. user=user,
  121. status=status_value(status),
  122. active=active,
  123. verified=True,
  124. )
  125. if status in accepted:
  126. return {
  127. "allowed": True,
  128. "bot_id": str(BOT_PROFILE_ID),
  129. "chat_id": int(chat_id),
  130. "status": status_value(status),
  131. }
  132. return {"allowed": False}
  133. async def verify_platform_membership(
  134. *,
  135. user_id: int,
  136. require_admin: bool = False,
  137. ) -> dict[str, Any]:
  138. if require_admin and is_sudoer(int(user_id)):
  139. return {
  140. "allowed": True,
  141. "bot_id": str(BOT_PROFILE_ID),
  142. "chat_id": None,
  143. "status": "sudoer",
  144. }
  145. candidates = await list_membership_candidates(int(user_id))
  146. if not candidates:
  147. return {"allowed": False, "reason": "membership_evidence_required"}
  148. internal_url = str(
  149. getattr(wbb, "SUPERVISOR_INTERNAL_URL", "") or ""
  150. ).rstrip("/")
  151. internal_token = str(getattr(wbb, "INTERNAL_TOKEN", "") or "")
  152. if internal_url and internal_token:
  153. try:
  154. async with ClientSession(timeout=ClientTimeout(total=20)) as session:
  155. async with session.post(
  156. f"{internal_url}/api/internal/v1/directory/verify-platform",
  157. headers={"Authorization": f"Bearer {internal_token}"},
  158. json={
  159. "user_id": str(user_id),
  160. "require_admin": bool(require_admin),
  161. },
  162. ) as response:
  163. if response.status != 200:
  164. return {"allowed": False, "reason": "verification_unavailable"}
  165. payload = await response.json()
  166. return payload.get("data") or {"allowed": False}
  167. except (ClientError, TimeoutError, ValueError) as exc:
  168. log.error(f"跨 Bot 成员身份校验失败:{exc}")
  169. return {"allowed": False, "reason": "verification_unavailable"}
  170. local_chat_ids = [
  171. int(item["chat_id"])
  172. for item in candidates
  173. if str(item.get("bot_id")) == str(BOT_PROFILE_ID)
  174. ]
  175. if not local_chat_ids:
  176. return {"allowed": False, "reason": "verification_unavailable"}
  177. return await verify_local_membership(
  178. user_id=int(user_id),
  179. chat_ids=local_chat_ids,
  180. require_admin=require_admin,
  181. )
  182. async def require_platform_membership(
  183. user_id: int,
  184. *,
  185. require_admin: bool = False,
  186. ) -> dict[str, Any]:
  187. result = await verify_platform_membership(
  188. user_id=int(user_id),
  189. require_admin=require_admin,
  190. )
  191. if result.get("allowed"):
  192. return result
  193. reason = result.get("reason")
  194. if reason == "membership_evidence_required":
  195. raise DirectoryServiceError(
  196. "membership_verification_required",
  197. "请先在任一受管群发送 /directory 完成群成员身份验证。",
  198. status=403,
  199. )
  200. if reason == "verification_unavailable":
  201. raise DirectoryServiceError(
  202. "membership_verification_unavailable",
  203. "暂时无法通过对应机器人核验群身份,请稍后重试。",
  204. status=503,
  205. )
  206. raise DirectoryServiceError(
  207. "managed_group_member_required",
  208. "该功能仅对受管群的当前成员开放。",
  209. status=403,
  210. )
  211. async def update_user_location(
  212. *,
  213. user: Any,
  214. longitude: float,
  215. latitude: float,
  216. source: str = "telegram_private",
  217. ) -> tuple[dict[str, Any], bool | None]:
  218. await require_platform_membership(int(user.id))
  219. identity = user_identity(user)
  220. await observe_directory_user(user)
  221. previous = await get_directory_location(int(user.id))
  222. consent = previous.get("geocode_consent") if previous else None
  223. location = await save_directory_location(
  224. user_id=int(user.id),
  225. longitude=longitude,
  226. latitude=latitude,
  227. source=source,
  228. actor_id=int(user.id),
  229. actor_name=identity["display_name"],
  230. geocode_consent=consent,
  231. bot_id=str(BOT_PROFILE_ID),
  232. )
  233. return location, consent
  234. async def update_user_geocode_consent(
  235. *,
  236. user: Any,
  237. consent: bool,
  238. ) -> dict[str, Any]:
  239. identity = user_identity(user)
  240. return await set_geocode_consent(
  241. user_id=int(user.id),
  242. consent=consent,
  243. source="telegram_private",
  244. actor_id=int(user.id),
  245. actor_name=identity["display_name"],
  246. )
  247. async def clear_user_location(*, user: Any) -> bool:
  248. identity = user_identity(user)
  249. return await clear_directory_location(
  250. user_id=int(user.id),
  251. actor_id=int(user.id),
  252. actor_name=identity["display_name"],
  253. source="telegram_private",
  254. reason="用户主动清除位置",
  255. )
  256. async def apply_as_teacher(*, user: Any) -> dict[str, Any]:
  257. await require_platform_membership(int(user.id))
  258. await observe_directory_user(user)
  259. try:
  260. return await submit_teacher_application(
  261. user_id=int(user.id),
  262. source="telegram_private",
  263. bot_id=str(BOT_PROFILE_ID),
  264. )
  265. except DirectoryDataError as exc:
  266. raise DirectoryServiceError(exc.code, str(exc), status=409) from exc
  267. async def change_own_teacher_state(
  268. *,
  269. user: Any,
  270. action: str,
  271. source: str = "telegram_private",
  272. idempotency_key: str | None = None,
  273. chat_id: int | None = None,
  274. ) -> tuple[dict[str, Any], bool]:
  275. await require_platform_membership(int(user.id))
  276. identity = user_identity(user)
  277. await observe_directory_user(user)
  278. try:
  279. return await set_teacher_state(
  280. user_id=int(user.id),
  281. action=action,
  282. actor_id=int(user.id),
  283. actor_name=identity["display_name"],
  284. source=source,
  285. reason="老师自助操作",
  286. idempotency_key=idempotency_key,
  287. chat_id=chat_id,
  288. )
  289. except DirectoryDataError as exc:
  290. raise DirectoryServiceError(exc.code, str(exc), status=409) from exc
  291. async def admin_decide_teacher(
  292. *,
  293. actor: Any,
  294. user_id: int,
  295. action: str,
  296. reason: str = "",
  297. ) -> dict[str, Any]:
  298. authorization = await require_platform_membership(
  299. int(actor.id),
  300. require_admin=True,
  301. )
  302. identity = user_identity(actor)
  303. try:
  304. return await decide_teacher_application(
  305. user_id=int(user_id),
  306. action=action,
  307. actor_id=int(actor.id),
  308. actor_name=identity["display_name"],
  309. source="telegram_private",
  310. reason=reason,
  311. authorization_chat_id=authorization.get("chat_id"),
  312. )
  313. except DirectoryDataError as exc:
  314. raise DirectoryServiceError(exc.code, str(exc), status=409) from exc
  315. async def admin_change_teacher_state(
  316. *,
  317. actor: Any,
  318. user_id: int,
  319. action: str,
  320. reason: str,
  321. ) -> dict[str, Any]:
  322. await require_platform_membership(int(actor.id), require_admin=True)
  323. if not reason.strip():
  324. raise DirectoryServiceError("reason_required", "必须填写操作原因。")
  325. identity = user_identity(actor)
  326. try:
  327. profile, _ = await set_teacher_state(
  328. user_id=int(user_id),
  329. action=action,
  330. actor_id=int(actor.id),
  331. actor_name=identity["display_name"],
  332. source="telegram_private",
  333. reason=reason,
  334. )
  335. return profile
  336. except DirectoryDataError as exc:
  337. raise DirectoryServiceError(exc.code, str(exc), status=409) from exc
  338. async def directory_for_user(
  339. *,
  340. user: Any,
  341. radius_km: int | None,
  342. page: int = 1,
  343. page_size: int = 10,
  344. ) -> tuple[list[dict[str, Any]], int]:
  345. await require_platform_membership(int(user.id))
  346. await observe_directory_user(user)
  347. location = await get_directory_location(int(user.id))
  348. if not location:
  349. raise DirectoryServiceError(
  350. "location_required",
  351. "请先点击“更新位置”发送定位,再查看老师。",
  352. status=409,
  353. )
  354. try:
  355. return await list_directory_teachers(
  356. longitude=float(location["longitude"]),
  357. latitude=float(location["latitude"]),
  358. max_distance_meters=(
  359. float(radius_km) * 1000 if radius_km is not None else None
  360. ),
  361. page=max(1, int(page)),
  362. page_size=max(1, min(int(page_size), 20)),
  363. )
  364. except DirectoryDataError as exc:
  365. raise DirectoryServiceError(exc.code, str(exc)) from exc
  366. def public_teacher(profile: dict[str, Any]) -> dict[str, Any]:
  367. region = profile.get("region") or {}
  368. return {
  369. "user_id": int(profile["user_id"]),
  370. "username": profile.get("username"),
  371. "display_name": profile.get("display_name") or f"用户 {profile['user_id']}",
  372. "online": teacher_is_online(profile),
  373. "online_until": aware_utc(profile.get("online_until")),
  374. "region": str(region.get("label") or "区域暂不可用"),
  375. "distance_meters": float(profile.get("distance_meters") or 0),
  376. "location_updated_at": aware_utc(profile.get("location_updated_at")),
  377. }
  378. def profile_summary(profile: dict[str, Any], location: dict[str, Any] | None) -> dict[str, Any]:
  379. return {
  380. "user_id": int(profile["user_id"]),
  381. "username": profile.get("username"),
  382. "display_name": profile.get("display_name") or f"用户 {profile['user_id']}",
  383. "application_status": profile.get("application_status", "none"),
  384. "listed": bool(profile.get("listed")),
  385. "online": teacher_is_online(profile),
  386. "online_until": aware_utc(profile.get("online_until")),
  387. "has_location": bool(location),
  388. "region": (
  389. str((location.get("region") or {}).get("label") or "区域暂不可用")
  390. if location
  391. else ""
  392. ),
  393. "geocode_consent": location.get("geocode_consent") if location else None,
  394. "location_updated_at": aware_utc(location.get("updated_at")) if location else None,
  395. }
  396. def format_distance(distance_meters: float) -> str:
  397. if distance_meters < 1000:
  398. return f"{max(0, round(distance_meters))} 米"
  399. return f"{distance_meters / 1000:.1f} 公里"
  400. async def _presence_sweeper() -> None:
  401. while True:
  402. try:
  403. await expire_teacher_presence()
  404. except Exception as exc:
  405. log.error(f"老师在线状态过期任务失败:{exc}")
  406. await asyncio.sleep(60)
  407. def start_presence_sweeper() -> None:
  408. global _presence_task
  409. if _presence_task and not _presence_task.done():
  410. return
  411. try:
  412. _presence_task = asyncio.create_task(
  413. _presence_sweeper(),
  414. name="teacher-presence-sweeper",
  415. )
  416. except RuntimeError:
  417. _presence_task = None
  418. async def stop_presence_sweeper() -> None:
  419. global _presence_task
  420. if _presence_task:
  421. _presence_task.cancel()
  422. with suppress(asyncio.CancelledError):
  423. await _presence_task
  424. _presence_task = None