dbpoints.py 29 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926
  1. from __future__ import annotations
  2. import asyncio
  3. import hashlib
  4. import re
  5. from collections import defaultdict
  6. from datetime import UTC, datetime, timedelta
  7. from secrets import token_hex
  8. from typing import Any
  9. from zoneinfo import ZoneInfo, ZoneInfoNotFoundError
  10. from pymongo import ASCENDING, DESCENDING
  11. from pymongo.errors import DuplicateKeyError
  12. from wbb import db
  13. SOURCE_CHECKIN = "checkin"
  14. SOURCE_ACTIVITY = "activity"
  15. SOURCE_UPVOTE = "upvote"
  16. SOURCE_ADMIN = "admin_adjustment"
  17. SOURCE_GIVEAWAY_ENTRY = "giveaway_entry"
  18. SOURCE_GIVEAWAY_REFUND = "giveaway_refund"
  19. SOURCE_GIVEAWAY_PARTICIPATION = "giveaway_participation"
  20. SOURCE_GIVEAWAY_WINNER = "giveaway_winner"
  21. DEFAULT_CHECKIN_SUCCESS_MESSAGE = (
  22. "签到成功:+{points} 积分,当前余额:{balance}。"
  23. )
  24. DEFAULT_CHECKIN_ALREADY_MESSAGE = (
  25. "你今天已经签到过了,当前余额:{balance}。"
  26. )
  27. CHECKIN_MESSAGE_MAX_LENGTH = 500
  28. POINT_SOURCES = {
  29. SOURCE_CHECKIN,
  30. SOURCE_ACTIVITY,
  31. SOURCE_UPVOTE,
  32. SOURCE_ADMIN,
  33. SOURCE_GIVEAWAY_ENTRY,
  34. SOURCE_GIVEAWAY_REFUND,
  35. SOURCE_GIVEAWAY_PARTICIPATION,
  36. SOURCE_GIVEAWAY_WINNER,
  37. }
  38. DEFAULT_POINT_RULES: dict[str, Any] = {
  39. "enabled": False,
  40. "timezone": "Asia/Shanghai",
  41. "checkin_enabled": True,
  42. "checkin_button_enabled": False,
  43. "checkin_points": 10,
  44. "checkin_success_message": DEFAULT_CHECKIN_SUCCESS_MESSAGE,
  45. "checkin_already_message": DEFAULT_CHECKIN_ALREADY_MESSAGE,
  46. "activity_enabled": True,
  47. "activity_points": 1,
  48. "activity_cooldown_seconds": 300,
  49. "activity_daily_cap": 10,
  50. "upvote_enabled": True,
  51. "upvote_points": 2,
  52. "upvote_pair_cooldown_seconds": 86400,
  53. "upvote_daily_cap": 10,
  54. }
  55. accountsdb = db.point_accounts
  56. transactionsdb = db.point_transactions
  57. rulesdb = db.point_rules
  58. rate_limitsdb = db.point_rate_limits
  59. _index_lock = asyncio.Lock()
  60. _indexes_ready = False
  61. _account_locks: defaultdict[tuple[int, int], asyncio.Lock] = defaultdict(asyncio.Lock)
  62. _rate_locks: defaultdict[str, asyncio.Lock] = defaultdict(asyncio.Lock)
  63. class PointsError(RuntimeError):
  64. pass
  65. class InsufficientPoints(PointsError):
  66. pass
  67. def utc_now() -> datetime:
  68. return datetime.now(UTC)
  69. def _as_utc(value: datetime) -> datetime:
  70. if value.tzinfo is None:
  71. return value.replace(tzinfo=UTC)
  72. return value.astimezone(UTC)
  73. def _local_day(now: datetime, timezone_name: str) -> str:
  74. try:
  75. zone = ZoneInfo(timezone_name)
  76. except ZoneInfoNotFoundError:
  77. zone = ZoneInfo("Asia/Shanghai")
  78. return _as_utc(now).astimezone(zone).date().isoformat()
  79. def _safe_int(value: Any, default: int, minimum: int, maximum: int) -> int:
  80. try:
  81. parsed = int(value)
  82. except (TypeError, ValueError):
  83. return default
  84. return max(minimum, min(parsed, maximum))
  85. def _safe_bool(value: Any, default: bool) -> bool:
  86. if isinstance(value, bool):
  87. return value
  88. if isinstance(value, (int, float)):
  89. return value != 0
  90. if isinstance(value, str):
  91. normalized = value.strip().lower()
  92. if normalized in {
  93. "true",
  94. "1",
  95. "yes",
  96. "on",
  97. "enable",
  98. "enabled",
  99. "开启",
  100. "打开",
  101. }:
  102. return True
  103. if normalized in {
  104. "false",
  105. "0",
  106. "no",
  107. "off",
  108. "disable",
  109. "disabled",
  110. "关闭",
  111. }:
  112. return False
  113. return default
  114. def _safe_text(value: Any, default: str, maximum: int) -> str:
  115. text = str(value or "").strip()
  116. if not text:
  117. return default
  118. return text[:maximum]
  119. def normalize_point_rules(raw: dict[str, Any] | None) -> dict[str, Any]:
  120. rules = dict(DEFAULT_POINT_RULES)
  121. if raw:
  122. rules.update({key: value for key, value in raw.items() if key in rules})
  123. for key in (
  124. "enabled",
  125. "checkin_enabled",
  126. "checkin_button_enabled",
  127. "activity_enabled",
  128. "upvote_enabled",
  129. ):
  130. rules[key] = _safe_bool(rules[key], DEFAULT_POINT_RULES[key])
  131. timezone_name = str(rules.get("timezone") or "Asia/Shanghai")
  132. try:
  133. ZoneInfo(timezone_name)
  134. except ZoneInfoNotFoundError:
  135. timezone_name = "Asia/Shanghai"
  136. rules["timezone"] = timezone_name
  137. rules["checkin_points"] = _safe_int(rules["checkin_points"], 10, 0, 100000)
  138. rules["checkin_success_message"] = _safe_text(
  139. rules["checkin_success_message"],
  140. DEFAULT_CHECKIN_SUCCESS_MESSAGE,
  141. CHECKIN_MESSAGE_MAX_LENGTH,
  142. )
  143. rules["checkin_already_message"] = _safe_text(
  144. rules["checkin_already_message"],
  145. DEFAULT_CHECKIN_ALREADY_MESSAGE,
  146. CHECKIN_MESSAGE_MAX_LENGTH,
  147. )
  148. rules["activity_points"] = _safe_int(rules["activity_points"], 1, 0, 100000)
  149. rules["activity_cooldown_seconds"] = _safe_int(
  150. rules["activity_cooldown_seconds"], 300, 10, 86400
  151. )
  152. rules["activity_daily_cap"] = _safe_int(
  153. rules["activity_daily_cap"], 10, 0, 1000000
  154. )
  155. rules["upvote_points"] = _safe_int(rules["upvote_points"], 2, 0, 100000)
  156. rules["upvote_pair_cooldown_seconds"] = _safe_int(
  157. rules["upvote_pair_cooldown_seconds"], 86400, 60, 2592000
  158. )
  159. rules["upvote_daily_cap"] = _safe_int(
  160. rules["upvote_daily_cap"], 10, 0, 1000000
  161. )
  162. return rules
  163. async def ensure_point_indexes() -> None:
  164. global _indexes_ready
  165. if _indexes_ready:
  166. return
  167. async with _index_lock:
  168. if _indexes_ready:
  169. return
  170. await accountsdb.create_index(
  171. [("chat_id", ASCENDING), ("user_id", ASCENDING)], unique=True
  172. )
  173. await accountsdb.create_index(
  174. [("chat_id", ASCENDING), ("balance", DESCENDING), ("user_id", ASCENDING)]
  175. )
  176. await transactionsdb.create_index([("idempotency_key", ASCENDING)], unique=True)
  177. await transactionsdb.create_index(
  178. [("chat_id", ASCENDING), ("user_id", ASCENDING), ("created_at", DESCENDING)]
  179. )
  180. await transactionsdb.create_index(
  181. [("chat_id", ASCENDING), ("source", ASCENDING), ("created_at", DESCENDING)]
  182. )
  183. await transactionsdb.create_index(
  184. [
  185. ("chat_id", ASCENDING),
  186. ("source", ASCENDING),
  187. ("idempotency_key", ASCENDING),
  188. ]
  189. )
  190. await rulesdb.create_index([("chat_id", ASCENDING)], unique=True)
  191. await rate_limitsdb.create_index([("rate_key", ASCENDING)], unique=True)
  192. await rate_limitsdb.create_index("expires_at", expireAfterSeconds=0)
  193. _indexes_ready = True
  194. async def get_point_rules(chat_id: int) -> dict[str, Any]:
  195. await ensure_point_indexes()
  196. doc = await rulesdb.find_one({"chat_id": chat_id})
  197. return normalize_point_rules(doc)
  198. async def set_point_rules(chat_id: int, values: dict[str, Any]) -> dict[str, Any]:
  199. await ensure_point_indexes()
  200. current = await get_point_rules(chat_id)
  201. current.update({key: value for key, value in values.items() if key in DEFAULT_POINT_RULES})
  202. normalized = normalize_point_rules(current)
  203. now = utc_now()
  204. await rulesdb.update_one(
  205. {"chat_id": chat_id},
  206. {"$set": {**normalized, "updated_at": now}, "$setOnInsert": {"created_at": now}},
  207. upsert=True,
  208. )
  209. return normalized
  210. async def list_checkin_keyboard_migrations(
  211. keyboard_version: int,
  212. ) -> list[int]:
  213. await ensure_point_indexes()
  214. cursor = rulesdb.find(
  215. {
  216. "enabled": True,
  217. "checkin_enabled": True,
  218. "checkin_button_enabled": True,
  219. "$or": [
  220. {"checkin_keyboard_version": {"$exists": False}},
  221. {"checkin_keyboard_version": {"$lt": int(keyboard_version)}},
  222. ],
  223. },
  224. {"chat_id": 1},
  225. )
  226. documents = await cursor.to_list(length=10000)
  227. return [int(document["chat_id"]) for document in documents]
  228. async def mark_checkin_keyboard_version(
  229. chat_id: int,
  230. keyboard_version: int,
  231. ) -> None:
  232. await ensure_point_indexes()
  233. await rulesdb.update_one(
  234. {"chat_id": int(chat_id)},
  235. {
  236. "$set": {
  237. "checkin_keyboard_version": int(keyboard_version),
  238. "updated_at": utc_now(),
  239. }
  240. },
  241. )
  242. async def _reconcile_account_locked(chat_id: int, user_id: int) -> dict[str, Any]:
  243. cached = await accountsdb.find_one({"chat_id": chat_id, "user_id": user_id})
  244. pipeline = [
  245. {"$match": {"chat_id": chat_id, "user_id": user_id}},
  246. {
  247. "$group": {
  248. "_id": None,
  249. "balance": {"$sum": "$delta"},
  250. "lifetime_earned": {
  251. "$sum": {"$cond": [{"$gt": ["$delta", 0]}, "$delta", 0]}
  252. },
  253. "lifetime_spent": {
  254. "$sum": {
  255. "$cond": [
  256. {"$lt": ["$delta", 0]},
  257. {"$multiply": ["$delta", -1]},
  258. 0,
  259. ]
  260. }
  261. },
  262. }
  263. },
  264. ]
  265. totals = [doc async for doc in transactionsdb.aggregate(pipeline)]
  266. latest = await transactionsdb.find_one(
  267. {"chat_id": chat_id, "user_id": user_id},
  268. sort=[("created_at", DESCENDING), ("_id", DESCENDING)],
  269. )
  270. total = totals[0] if totals else {}
  271. now = utc_now()
  272. account = {
  273. "chat_id": chat_id,
  274. "user_id": user_id,
  275. "balance": int(total.get("balance", 0)),
  276. "lifetime_earned": int(total.get("lifetime_earned", 0)),
  277. "lifetime_spent": int(total.get("lifetime_spent", 0)),
  278. "last_transaction_id": latest.get("transaction_id") if latest else None,
  279. "updated_at": now,
  280. }
  281. if latest:
  282. account["username"] = (
  283. cached.get("username") if cached and "username" in cached else latest.get("username")
  284. )
  285. account["first_name"] = (
  286. cached.get("first_name")
  287. if cached and "first_name" in cached
  288. else latest.get("first_name")
  289. )
  290. account["display_name"] = (
  291. cached.get("display_name")
  292. if cached and "display_name" in cached
  293. else latest.get("display_name") or latest.get("first_name")
  294. )
  295. await accountsdb.update_one(
  296. {"chat_id": chat_id, "user_id": user_id},
  297. {"$set": account, "$setOnInsert": {"created_at": now}},
  298. upsert=True,
  299. )
  300. return account
  301. async def reconcile_account(chat_id: int, user_id: int) -> dict[str, Any]:
  302. await ensure_point_indexes()
  303. async with _account_locks[(chat_id, user_id)]:
  304. return await _reconcile_account_locked(chat_id, user_id)
  305. async def get_point_account(chat_id: int, user_id: int) -> dict[str, Any]:
  306. await ensure_point_indexes()
  307. account = await accountsdb.find_one({"chat_id": chat_id, "user_id": user_id})
  308. latest = await transactionsdb.find_one(
  309. {"chat_id": chat_id, "user_id": user_id},
  310. projection={"transaction_id": 1},
  311. sort=[("created_at", DESCENDING), ("_id", DESCENDING)],
  312. )
  313. if latest and (
  314. not account or account.get("last_transaction_id") != latest.get("transaction_id")
  315. ):
  316. return await reconcile_account(chat_id, user_id)
  317. if account:
  318. return account
  319. return {
  320. "chat_id": chat_id,
  321. "user_id": user_id,
  322. "balance": 0,
  323. "lifetime_earned": 0,
  324. "lifetime_spent": 0,
  325. "last_transaction_id": None,
  326. }
  327. async def update_point_account_identity(
  328. *,
  329. chat_id: int,
  330. user_id: int,
  331. username: str | None,
  332. first_name: str | None,
  333. last_name: str | None,
  334. display_name: str,
  335. ) -> None:
  336. await ensure_point_indexes()
  337. await accountsdb.update_one(
  338. {"chat_id": int(chat_id), "user_id": int(user_id)},
  339. {
  340. "$set": {
  341. "username": username,
  342. "first_name": first_name,
  343. "last_name": last_name,
  344. "display_name": display_name or first_name,
  345. "identity_updated_at": utc_now(),
  346. }
  347. },
  348. )
  349. async def _write_points_locked(
  350. *,
  351. chat_id: int,
  352. user_id: int,
  353. source: str,
  354. idempotency_key: str,
  355. delta: int | None = None,
  356. target_balance: int | None = None,
  357. actor_id: int | str | None = None,
  358. reason: str = "",
  359. reference_id: str | None = None,
  360. username: str | None = None,
  361. first_name: str | None = None,
  362. display_name: str | None = None,
  363. ) -> tuple[dict[str, Any], bool]:
  364. existing = await transactionsdb.find_one({"idempotency_key": idempotency_key})
  365. if existing:
  366. account = await _reconcile_account_locked(chat_id, user_id)
  367. return account, False
  368. account = await accountsdb.find_one({"chat_id": chat_id, "user_id": user_id})
  369. latest = await transactionsdb.find_one(
  370. {"chat_id": chat_id, "user_id": user_id},
  371. projection={"transaction_id": 1},
  372. sort=[("created_at", DESCENDING), ("_id", DESCENDING)],
  373. )
  374. if not account or (
  375. latest and account.get("last_transaction_id") != latest.get("transaction_id")
  376. ):
  377. account = await _reconcile_account_locked(chat_id, user_id)
  378. balance_before = int(account.get("balance", 0))
  379. actual_delta = (
  380. int(target_balance) - balance_before
  381. if target_balance is not None
  382. else int(delta or 0)
  383. )
  384. balance_after = balance_before + actual_delta
  385. if balance_after < 0:
  386. raise InsufficientPoints("积分余额不足。")
  387. now = utc_now()
  388. transaction = {
  389. "transaction_id": token_hex(10),
  390. "idempotency_key": idempotency_key,
  391. "chat_id": chat_id,
  392. "user_id": user_id,
  393. "username": username,
  394. "first_name": first_name,
  395. "display_name": display_name or first_name,
  396. "delta": actual_delta,
  397. "balance_before": balance_before,
  398. "balance_after": balance_after,
  399. "source": source,
  400. "actor_id": actor_id,
  401. "reason": reason.strip(),
  402. "reference_id": reference_id,
  403. "created_at": now,
  404. }
  405. try:
  406. await transactionsdb.insert_one(transaction)
  407. except DuplicateKeyError:
  408. account = await _reconcile_account_locked(chat_id, user_id)
  409. return account, False
  410. update = {
  411. "$set": {
  412. "balance": balance_after,
  413. "last_transaction_id": transaction["transaction_id"],
  414. "username": username or account.get("username"),
  415. "first_name": first_name or account.get("first_name"),
  416. "display_name": display_name or first_name or account.get("display_name"),
  417. "updated_at": now,
  418. },
  419. "$setOnInsert": {"created_at": now},
  420. "$inc": {
  421. "lifetime_earned": max(actual_delta, 0),
  422. "lifetime_spent": max(-actual_delta, 0),
  423. },
  424. }
  425. await accountsdb.update_one(
  426. {"chat_id": chat_id, "user_id": user_id}, update, upsert=True
  427. )
  428. account = await accountsdb.find_one({"chat_id": chat_id, "user_id": user_id})
  429. return account or {"chat_id": chat_id, "user_id": user_id, "balance": balance_after}, True
  430. async def adjust_points(
  431. *,
  432. chat_id: int,
  433. user_id: int,
  434. delta: int,
  435. source: str,
  436. idempotency_key: str,
  437. actor_id: int | str | None = None,
  438. reason: str = "",
  439. reference_id: str | None = None,
  440. username: str | None = None,
  441. first_name: str | None = None,
  442. display_name: str | None = None,
  443. ) -> tuple[dict[str, Any], bool]:
  444. await ensure_point_indexes()
  445. if source not in POINT_SOURCES:
  446. raise PointsError(f"不支持的积分来源:{source}")
  447. if not idempotency_key.strip():
  448. raise PointsError("缺少请求幂等键。")
  449. delta = int(delta)
  450. if delta == 0:
  451. account = await get_point_account(chat_id, user_id)
  452. return account, False
  453. async with _account_locks[(chat_id, user_id)]:
  454. return await _write_points_locked(
  455. chat_id=chat_id,
  456. user_id=user_id,
  457. delta=delta,
  458. source=source,
  459. idempotency_key=idempotency_key,
  460. actor_id=actor_id,
  461. reason=reason,
  462. reference_id=reference_id,
  463. username=username,
  464. first_name=first_name,
  465. display_name=display_name,
  466. )
  467. async def set_points(
  468. *,
  469. chat_id: int,
  470. user_id: int,
  471. balance: int,
  472. actor_id: int | str,
  473. reason: str,
  474. idempotency_key: str,
  475. username: str | None = None,
  476. first_name: str | None = None,
  477. display_name: str | None = None,
  478. ) -> tuple[dict[str, Any], bool]:
  479. await ensure_point_indexes()
  480. if balance < 0:
  481. raise PointsError("积分余额不能小于 0。")
  482. if not idempotency_key.strip():
  483. raise PointsError("缺少请求幂等键。")
  484. async with _account_locks[(chat_id, user_id)]:
  485. return await _write_points_locked(
  486. chat_id=chat_id,
  487. user_id=user_id,
  488. target_balance=balance,
  489. source=SOURCE_ADMIN,
  490. idempotency_key=idempotency_key,
  491. actor_id=actor_id,
  492. reason=reason,
  493. username=username,
  494. first_name=first_name,
  495. display_name=display_name,
  496. )
  497. async def get_point_transaction_by_key(idempotency_key: str) -> dict[str, Any] | None:
  498. await ensure_point_indexes()
  499. return await transactionsdb.find_one({"idempotency_key": idempotency_key})
  500. async def list_point_accounts(
  501. *,
  502. chat_id: int | None = None,
  503. query: str = "",
  504. page: int = 1,
  505. page_size: int = 20,
  506. ) -> tuple[list[dict[str, Any]], int]:
  507. await ensure_point_indexes()
  508. filters: dict[str, Any] = {}
  509. if chat_id is not None:
  510. filters["chat_id"] = chat_id
  511. query = query.strip()
  512. if query:
  513. if query.lstrip("-").isdigit():
  514. filters["user_id"] = int(query)
  515. else:
  516. value = re.escape(query.lstrip("@"))
  517. filters["$or"] = [
  518. {"username": {"$regex": value, "$options": "i"}},
  519. {"first_name": {"$regex": value, "$options": "i"}},
  520. {"display_name": {"$regex": value, "$options": "i"}},
  521. ]
  522. page = max(1, page)
  523. page_size = max(1, min(page_size, 100))
  524. total = await accountsdb.count_documents(filters)
  525. cursor = (
  526. accountsdb.find(filters)
  527. .sort([("balance", DESCENDING), ("user_id", ASCENDING)])
  528. .skip((page - 1) * page_size)
  529. .limit(page_size)
  530. )
  531. return [doc async for doc in cursor], total
  532. async def list_point_transactions(
  533. *,
  534. chat_id: int | None = None,
  535. user_id: int | None = None,
  536. source: str | None = None,
  537. created_from: datetime | None = None,
  538. created_to: datetime | None = None,
  539. page: int = 1,
  540. page_size: int = 20,
  541. ) -> tuple[list[dict[str, Any]], int]:
  542. await ensure_point_indexes()
  543. filters: dict[str, Any] = {}
  544. if chat_id is not None:
  545. filters["chat_id"] = chat_id
  546. if user_id is not None:
  547. filters["user_id"] = user_id
  548. if source:
  549. if source not in POINT_SOURCES:
  550. return [], 0
  551. filters["source"] = source
  552. if created_from or created_to:
  553. created_filter: dict[str, datetime] = {}
  554. if created_from:
  555. created_filter["$gte"] = _as_utc(created_from)
  556. if created_to:
  557. created_filter["$lte"] = _as_utc(created_to)
  558. filters["created_at"] = created_filter
  559. page = max(1, page)
  560. page_size = max(1, min(page_size, 100))
  561. total = await transactionsdb.count_documents(filters)
  562. cursor = (
  563. transactionsdb.find(filters)
  564. .sort([("created_at", DESCENDING), ("_id", DESCENDING)])
  565. .skip((page - 1) * page_size)
  566. .limit(page_size)
  567. )
  568. return [doc async for doc in cursor], total
  569. async def reconcile_all_accounts(
  570. *, chat_id: int | None = None, limit: int = 1000
  571. ) -> dict[str, int]:
  572. """Rebuild cached accounts from the immutable ledger."""
  573. await ensure_point_indexes()
  574. match: dict[str, Any] = {}
  575. if chat_id is not None:
  576. match["chat_id"] = int(chat_id)
  577. pipeline: list[dict[str, Any]] = []
  578. if match:
  579. pipeline.append({"$match": match})
  580. pipeline.extend(
  581. [
  582. {"$group": {"_id": {"chat_id": "$chat_id", "user_id": "$user_id"}}},
  583. {"$limit": max(1, min(int(limit), 10000))},
  584. ]
  585. )
  586. repaired = 0
  587. async for item in transactionsdb.aggregate(pipeline):
  588. await reconcile_account(int(item["_id"]["chat_id"]), int(item["_id"]["user_id"]))
  589. repaired += 1
  590. return {"repaired": repaired}
  591. async def award_checkin(
  592. *,
  593. chat_id: int,
  594. user_id: int,
  595. username: str | None,
  596. first_name: str | None,
  597. display_name: str | None = None,
  598. now: datetime | None = None,
  599. ) -> tuple[dict[str, Any], bool]:
  600. rules = await get_point_rules(chat_id)
  601. if not rules["enabled"] or not rules["checkin_enabled"]:
  602. raise PointsError("本群尚未开启积分签到。")
  603. now = _as_utc(now or utc_now())
  604. day = _local_day(now, rules["timezone"])
  605. return await adjust_points(
  606. chat_id=chat_id,
  607. user_id=user_id,
  608. delta=rules["checkin_points"],
  609. source=SOURCE_CHECKIN,
  610. idempotency_key=f"checkin:{chat_id}:{user_id}:{day}",
  611. reason=f"每日签到({day})",
  612. username=username,
  613. first_name=first_name,
  614. display_name=display_name,
  615. )
  616. async def get_checkin_metrics(
  617. *,
  618. chat_id: int,
  619. user_id: int,
  620. now: datetime | None = None,
  621. ) -> dict[str, int]:
  622. await ensure_point_indexes()
  623. rules = await get_point_rules(chat_id)
  624. now = _as_utc(now or utc_now())
  625. try:
  626. zone = ZoneInfo(rules["timezone"])
  627. except ZoneInfoNotFoundError:
  628. zone = ZoneInfo("Asia/Shanghai")
  629. day = now.astimezone(zone).date().isoformat()
  630. daily_key_pattern = (
  631. rf"^checkin:{re.escape(str(int(chat_id)))}:\d+:{re.escape(day)}$"
  632. )
  633. daily_cursor = transactionsdb.find(
  634. {
  635. "chat_id": int(chat_id),
  636. "source": SOURCE_CHECKIN,
  637. "idempotency_key": {"$regex": daily_key_pattern},
  638. },
  639. {"user_id": 1},
  640. ).sort([("created_at", ASCENDING), ("_id", ASCENDING)])
  641. daily_entries = [item async for item in daily_cursor]
  642. rank = next(
  643. (
  644. index
  645. for index, item in enumerate(daily_entries, 1)
  646. if int(item.get("user_id", 0)) == int(user_id)
  647. ),
  648. 0,
  649. )
  650. user_key_prefix = (
  651. f"checkin:{int(chat_id)}:{int(user_id)}:"
  652. )
  653. month_key_pattern = (
  654. rf"^{re.escape(user_key_prefix + day[:7])}-\d{{2}}$"
  655. )
  656. month_days = await transactionsdb.count_documents(
  657. {
  658. "chat_id": int(chat_id),
  659. "user_id": int(user_id),
  660. "source": SOURCE_CHECKIN,
  661. "idempotency_key": {"$regex": month_key_pattern},
  662. }
  663. )
  664. history_key_pattern = (
  665. rf"^{re.escape(user_key_prefix)}\d{{4}}-\d{{2}}-\d{{2}}$"
  666. )
  667. history_cursor = transactionsdb.find(
  668. {
  669. "chat_id": int(chat_id),
  670. "user_id": int(user_id),
  671. "source": SOURCE_CHECKIN,
  672. "idempotency_key": {"$regex": history_key_pattern},
  673. },
  674. {"idempotency_key": 1},
  675. )
  676. checkin_days = set()
  677. async for item in history_cursor:
  678. try:
  679. checkin_days.add(
  680. datetime.strptime(
  681. str(item["idempotency_key"]).rsplit(":", 1)[-1],
  682. "%Y-%m-%d",
  683. ).date()
  684. )
  685. except (KeyError, TypeError, ValueError):
  686. continue
  687. expected_day = datetime.strptime(day, "%Y-%m-%d").date()
  688. streak = 0
  689. while expected_day in checkin_days:
  690. streak += 1
  691. expected_day -= timedelta(days=1)
  692. return {"rank": rank, "streak": streak, "month_days": int(month_days)}
  693. async def award_activity(
  694. *,
  695. chat_id: int,
  696. user_id: int,
  697. message_id: int,
  698. content: str,
  699. username: str | None,
  700. first_name: str | None,
  701. display_name: str | None = None,
  702. now: datetime | None = None,
  703. ) -> tuple[dict[str, Any] | None, bool]:
  704. rules = await get_point_rules(chat_id)
  705. if not rules["enabled"] or not rules["activity_enabled"]:
  706. return None, False
  707. points = int(rules["activity_points"])
  708. cap = int(rules["activity_daily_cap"])
  709. if points <= 0 or cap <= 0:
  710. return None, False
  711. now = _as_utc(now or utc_now())
  712. day = _local_day(now, rules["timezone"])
  713. rate_key = f"activity:{chat_id}:{user_id}:{day}"
  714. content_hash = hashlib.sha256(content.strip().lower().encode("utf-8")).hexdigest()
  715. async with _rate_locks[rate_key]:
  716. state = await rate_limitsdb.find_one({"rate_key": rate_key}) or {}
  717. last_awarded = state.get("last_awarded_at")
  718. if last_awarded and (
  719. now - _as_utc(last_awarded)
  720. ).total_seconds() < rules["activity_cooldown_seconds"]:
  721. return None, False
  722. if int(state.get("awarded_points", 0)) >= cap:
  723. return None, False
  724. recent_hashes = [
  725. item
  726. for item in state.get("recent_hashes", [])
  727. if now - _as_utc(item["created_at"]) < timedelta(hours=1)
  728. ]
  729. if any(item["hash"] == content_hash for item in recent_hashes):
  730. return None, False
  731. award = min(points, cap - int(state.get("awarded_points", 0)))
  732. account, created = await adjust_points(
  733. chat_id=chat_id,
  734. user_id=user_id,
  735. delta=award,
  736. source=SOURCE_ACTIVITY,
  737. idempotency_key=f"activity:{chat_id}:{message_id}:{user_id}",
  738. reference_id=str(message_id),
  739. reason="活跃消息奖励",
  740. username=username,
  741. first_name=first_name,
  742. display_name=display_name,
  743. )
  744. if created:
  745. recent_hashes.append({"hash": content_hash, "created_at": now})
  746. await rate_limitsdb.update_one(
  747. {"rate_key": rate_key},
  748. {
  749. "$set": {
  750. "chat_id": chat_id,
  751. "user_id": user_id,
  752. "rule": SOURCE_ACTIVITY,
  753. "last_awarded_at": now,
  754. "recent_hashes": recent_hashes[-20:],
  755. "expires_at": now + timedelta(days=3),
  756. },
  757. "$inc": {"awarded_points": award},
  758. "$setOnInsert": {"created_at": now},
  759. },
  760. upsert=True,
  761. )
  762. return account, created
  763. async def award_upvote(
  764. *,
  765. chat_id: int,
  766. voter_id: int,
  767. target_id: int,
  768. message_id: int,
  769. username: str | None,
  770. first_name: str | None,
  771. display_name: str | None = None,
  772. now: datetime | None = None,
  773. ) -> tuple[dict[str, Any] | None, bool]:
  774. if voter_id == target_id:
  775. return None, False
  776. rules = await get_point_rules(chat_id)
  777. if not rules["enabled"] or not rules["upvote_enabled"]:
  778. return None, False
  779. points = int(rules["upvote_points"])
  780. cap = int(rules["upvote_daily_cap"])
  781. if points <= 0 or cap <= 0:
  782. return None, False
  783. now = _as_utc(now or utc_now())
  784. day = _local_day(now, rules["timezone"])
  785. daily_key = f"upvote-target:{chat_id}:{target_id}:{day}"
  786. pair_key = f"upvote-pair:{chat_id}:{voter_id}:{target_id}"
  787. async with _rate_locks[pair_key]:
  788. async with _rate_locks[daily_key]:
  789. pair = await rate_limitsdb.find_one({"rate_key": pair_key}) or {}
  790. last_awarded = pair.get("last_awarded_at")
  791. if last_awarded and (
  792. now - _as_utc(last_awarded)
  793. ).total_seconds() < rules["upvote_pair_cooldown_seconds"]:
  794. return None, False
  795. daily = await rate_limitsdb.find_one({"rate_key": daily_key}) or {}
  796. already = int(daily.get("awarded_points", 0))
  797. if already >= cap:
  798. return None, False
  799. award = min(points, cap - already)
  800. account, created = await adjust_points(
  801. chat_id=chat_id,
  802. user_id=target_id,
  803. delta=award,
  804. source=SOURCE_UPVOTE,
  805. idempotency_key=f"upvote:{chat_id}:{message_id}:{voter_id}:{target_id}",
  806. actor_id=voter_id,
  807. reference_id=str(message_id),
  808. reason="收到有效点赞",
  809. username=username,
  810. first_name=first_name,
  811. display_name=display_name,
  812. )
  813. if created:
  814. expires_at = now + timedelta(days=35)
  815. await rate_limitsdb.update_one(
  816. {"rate_key": pair_key},
  817. {
  818. "$set": {
  819. "chat_id": chat_id,
  820. "user_id": target_id,
  821. "voter_id": voter_id,
  822. "rule": "upvote_pair",
  823. "last_awarded_at": now,
  824. "expires_at": expires_at,
  825. },
  826. "$setOnInsert": {"created_at": now},
  827. },
  828. upsert=True,
  829. )
  830. await rate_limitsdb.update_one(
  831. {"rate_key": daily_key},
  832. {
  833. "$set": {
  834. "chat_id": chat_id,
  835. "user_id": target_id,
  836. "rule": SOURCE_UPVOTE,
  837. "last_awarded_at": now,
  838. "expires_at": now + timedelta(days=3),
  839. },
  840. "$inc": {"awarded_points": award},
  841. "$setOnInsert": {"created_at": now},
  842. },
  843. upsert=True,
  844. )
  845. return account, created