Experimental Discord bot written in Python
選択できるのは25トピックまでです。 トピックは、先頭が英数字で、英数字とダッシュ('-')を使用した35文字以内のものにしてください。

kickstartercog.py 31KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900
  1. import heapq
  2. from asyncio import sleep
  3. from collections.abc import Awaitable, Callable, Iterable
  4. from enum import IntEnum
  5. from sqlite3 import Connection, connect
  6. from time import time as now_timestamp
  7. from typing import Optional, TypeVar
  8. from discord import Attachment, Guild, Interaction, Member, Role
  9. from discord.app_commands import Group, default_permissions
  10. from discord.errors import DiscordException, HTTPException
  11. from discord.ext.commands import Cog
  12. from discord.ui import FileUpload, Label, Modal
  13. from config import CONFIG
  14. from rocketbot.bot import Rocketbot
  15. from rocketbot.cogs.basecog import BaseCog
  16. from rocketbot.cogsetting import CogSetting
  17. from rocketbot.utils import (
  18. ADMIN_PERMISSIONS,
  19. MOD_PERMISSIONS,
  20. dump_stacktrace,
  21. is_discord_username,
  22. levenshtein,
  23. )
  24. _T = TypeVar('_T')
  25. class KickstarterCog(BaseCog):
  26. """
  27. Assigns Discord users a role if their username appears in an imported list,
  28. generally the results of a Kickstarter survey.
  29. Discord does not currently have a native Kickstarter integration, and doing
  30. a full API integration is a bit ambitious, so this is a stopgap solution.
  31. """
  32. shared: Optional['KickstarterCog'] = None
  33. SETTING_ENABLED = CogSetting(
  34. name='enabled',
  35. datatype=bool,
  36. default_value=False,
  37. brief='Kickstarter user linking',
  38. description='Whether this module is enabled for a guild.',
  39. )
  40. SETTING_ROLE = CogSetting(
  41. name='backer_role',
  42. datatype=int,
  43. default_value=0,
  44. brief='role to assign to Kickstarter backers',
  45. description=''
  46. )
  47. def __init__(self, bot: Rocketbot):
  48. super().__init__(
  49. bot,
  50. config_prefix='kickstarter',
  51. short_description='For linking Kickstarter backers to their Discord handles.',
  52. )
  53. Self = KickstarterCog
  54. self.add_setting(Self.SETTING_ENABLED)
  55. # SETTING_ROLE managed manually
  56. self.con: Connection = connect('kickstarter.sqlite3')
  57. Self.shared = self
  58. def __is_enabled(self, guild: Guild) -> bool:
  59. Self = KickstarterCog
  60. return self.get_guild_setting(guild, Self.SETTING_ENABLED)
  61. def __get_backer_role_id(self, guild: Guild) -> int | None:
  62. Self = KickstarterCog
  63. role_id = self.get_guild_setting(guild, Self.SETTING_ROLE)
  64. return role_id if role_id != 0 else None
  65. async def __fetch_backer_role(self, guild: Guild) -> Role:
  66. role_id = self.__get_backer_role_id(guild)
  67. if role_id is None:
  68. raise _NoBackerRoleException()
  69. ret_val = guild.get_role(role_id) or await guild.fetch_role(role_id)
  70. if ret_val is None:
  71. self.log(guild, f"Backer role with id {role_id} could not be retrieved. Role removed?")
  72. raise _NoBackerRoleException()
  73. return ret_val
  74. def __get_cached_backer_role(self, guild: Guild) -> Role | None:
  75. """Synchronously gets the configured backer role from the cache if possible."""
  76. role_id = self.__get_backer_role_id(guild)
  77. return guild.get_role(role_id) if role_id is not None else None
  78. def __set_backer_role(self, role: Role | None):
  79. Self = KickstarterCog
  80. self.set_guild_setting(role.guild, Self.SETTING_ROLE, role.id if role else None)
  81. async def __check_disabled(self, interaction: Interaction) -> bool:
  82. """Checks if this feature is enabled, and if not, sends an error
  83. response. Return value is whether caller should bail out."""
  84. if not self.__is_enabled(interaction.guild):
  85. text = f"{CONFIG['failure_emoji']} Kickstarter feature not enabled"
  86. await interaction.response.send_message(text, ephemeral=True)
  87. return True
  88. return False
  89. async def __check_configured(self, interaction: Interaction) -> bool:
  90. """Checks if the current guild is configured properly for this feature,
  91. and if not, sends an error response. Return value is whether caller
  92. should bail out."""
  93. if self.__get_backer_role_id(interaction.guild) is None:
  94. text = f"{CONFIG['failure_emoji']} Backer role is not configured"
  95. await interaction.response.send_message(text, ephemeral=True)
  96. return True
  97. return False
  98. # -- Admin commands -----
  99. kickstarter = Group(
  100. name='kickstarter',
  101. description='Manages roles for users identified as Kickstarter backers.',
  102. guild_only=True,
  103. default_permissions=MOD_PERMISSIONS
  104. )
  105. @kickstarter.command(
  106. description='Configures which role to give to Kickstarter backers.'
  107. )
  108. @default_permissions(ADMIN_PERMISSIONS)
  109. async def set_role(self, interaction: Interaction, role: Role):
  110. self.__set_backer_role(role)
  111. text = f"{CONFIG['info_emoji']} Backer role set to {role.name}"
  112. await interaction.response.send_message(text, ephemeral=True)
  113. @kickstarter.command(
  114. description='Assigns backer role to imported backer usernames.'
  115. )
  116. async def sync(self, interaction: Interaction):
  117. if await self.__check_disabled(interaction): return
  118. if await self.__check_configured(interaction): return
  119. guild = interaction.guild
  120. await interaction.response.defer(ephemeral=True, thinking=True)
  121. sync_result: _SyncUsernamesResult = await self.__sync_by_username(guild)
  122. text = f"{CONFIG['success_emoji']} Sync complete.\n" + \
  123. sync_result.summary_markdown()
  124. await interaction.followup.send(text, ephemeral=True)
  125. @kickstarter.command(
  126. description='Shows info about Kickstarter linked Discord members.'
  127. )
  128. async def info(self, interaction: Interaction):
  129. if await self.__check_disabled(interaction): return
  130. if await self.__check_configured(interaction): return
  131. await interaction.response.defer(ephemeral=True, thinking=True)
  132. guild = interaction.guild
  133. backer_role = await self.__fetch_backer_role(guild)
  134. stats: _Stats = self.__fetch_stats(guild.id)
  135. lines: list[str] = []
  136. if backer_role is None:
  137. lines.append("- No backer role configured yet (use " \
  138. f"`/{KickstarterCog.kickstarter.name} {KickstarterCog.set_role.name}`)")
  139. else:
  140. lines.append(f"- Backer role configured as `{backer_role.name}`")
  141. if stats.username_count > 0:
  142. lines.append(f"- **{stats.username_count:,}** Discord usernames imported")
  143. lines.append("- Last import with new records on " \
  144. f"<t:{stats.username_last_imported_at}:f>")
  145. lines.append(f"- **{stats.username_found_count:,}** members linked successfully")
  146. if stats.username_not_found_count > 0:
  147. lines.append(f"- **{stats.username_not_found_count:,}** imported " \
  148. "usernames not yet linked to Discord members")
  149. if stats.username_unprocessed_count > 0:
  150. lines.append(f"- **{stats.username_unprocessed_count:,}** " \
  151. "usernames added since last sync")
  152. else:
  153. lines.append("- No Discord usernames imported yet. Use " \
  154. f"`/{KickstarterCog.kickstarter.name} {KickstarterCog.upload.name}`")
  155. text = f"{CONFIG['info_emoji']} Kickstarter import stats\n\n"
  156. text += "\n".join(lines)
  157. await interaction.followup.send(text, ephemeral=True)
  158. @kickstarter.command(
  159. description='Shows Kickstarter link details about a specific Discord member.'
  160. )
  161. async def find_member(self, interaction: Interaction, member: Member):
  162. if await self.__check_disabled(interaction): return
  163. if await self.__check_configured(interaction): return
  164. await interaction.response.defer(ephemeral=True, thinking=True)
  165. guild = interaction.guild
  166. try:
  167. backer_role = await self.__fetch_backer_role(guild)
  168. except _NoBackerRoleException:
  169. text = f"{CONFIG['failure_emoji']} Backer role not yet configured!"
  170. await interaction.followup.send(text, ephemeral=True)
  171. return
  172. if backer_role in member.roles:
  173. text = f"{CONFIG['info_emoji']} Member `@{member.name}` has " \
  174. f"`{backer_role.name}` role."
  175. await interaction.followup.send(text, ephemeral=True)
  176. return
  177. if member.bot:
  178. text = f"{CONFIG['info_emoji']} Member `@{member.name}` is a robut! " \
  179. "Robuts cannot back Kickstarters!"
  180. await interaction.followup.send(text, ephemeral=True)
  181. return
  182. username_record = self.__fetch_kickstarter_username(guild.id, username=member.name)
  183. if username_record is not None:
  184. await member.add_roles(backer_role)
  185. username_record.lookup_status = _LookupStatus.member_found
  186. username_record.discord_member_id = member.id
  187. self.__update_kickstarter_username(username_record)
  188. text = f"{CONFIG['success_emoji']} Member's username found in " \
  189. f"Kickstarter export. The `{backer_role.name}` role has now " \
  190. "been assigned to them."
  191. await interaction.followup.send(text, ephemeral=True)
  192. return
  193. closest_matches = self.__find_nearest_usernames(guild.id, member.name,
  194. limit=8, max_distance=3)
  195. stats = self.__fetch_stats(guild.id)
  196. text = f"{CONFIG['info_emoji']} The username `@{member.name}` is not in " \
  197. "the latest Kickstarter import." \
  198. "\n" \
  199. "\n- Did they complete their survey yet?" \
  200. "\n- Did they fill in their Discord username in the survey?" \
  201. "\n- Did they complete their survey after the last import " \
  202. f"(<t:{stats.username_last_imported_at}:f>)? They might be in the next one."
  203. if len(closest_matches) > 0:
  204. text += "\n- Did they misspell their username in the survey? " \
  205. "Here are some similar ones."
  206. for similar in closest_matches:
  207. text += f"\n - `{similar.discord_username}`"
  208. if similar.lookup_status == _LookupStatus.misspelled_username:
  209. text += f" (manually linked to <@{similar.discord_member_id}> by mod)"
  210. text += f"\n-# Hint: Use `/{KickstarterCog.kickstarter.name} {KickstarterCog.link.name} " \
  211. f"@{member.name} surveyusername` to link the member to their " \
  212. "misspelled survey username"
  213. await interaction.followup.send(text, ephemeral=True)
  214. @kickstarter.command(
  215. description='Manually links a member to a misspelled survey username.'
  216. )
  217. async def link(self, interaction: Interaction, member: Member, username: str):
  218. if await self.__check_disabled(interaction): return
  219. if await self.__check_configured(interaction): return
  220. await interaction.response.defer(ephemeral=True, thinking=True)
  221. guild = interaction.guild
  222. try:
  223. backer_role = await self.__fetch_backer_role(guild)
  224. except _NoBackerRoleException:
  225. text = f"{CONFIG['failure_emoji']} Backer role not configured!"
  226. await interaction.followup.send(text, ephemeral=True)
  227. return
  228. normal_username = KickstarterCog.__normalize_discord_username(username)
  229. username = self.__fetch_kickstarter_username(guild.id, username=normal_username)
  230. if username is None:
  231. closest_matches = self.__find_nearest_usernames(guild.id, normal_username, 1)
  232. text = f"{CONFIG['failure_emoji']} No survey username found for `{username}`."
  233. if len(closest_matches) > 0:
  234. text += f" Did you mean `/{KickstarterCog.kickstarter.name} " \
  235. f"{KickstarterCog.link.name} @{member.name} " \
  236. f"{closest_matches[0].discord_username}`?"
  237. await interaction.followup.send(text, ephemeral=True)
  238. return
  239. if username.lookup_status == _LookupStatus.member_found:
  240. text = f"{CONFIG['failure_emoji']} That survey username is already " \
  241. f"linked to <@{username.discord_member_id}>. Cannot be linked to " \
  242. "another member.\n" \
  243. f"- `/{KickstarterCog.kickstarter.name} {KickstarterCog.reset.name} " \
  244. f"{username}` will unlink that Discord record\n" \
  245. "- If all else fails, you can just manually give them the backer role."
  246. await interaction.followup.send(text, ephemeral=True)
  247. return
  248. elif username.lookup_status == _LookupStatus.misspelled_username:
  249. text = f"{CONFIG['failure_emoji']} That survey username was already " \
  250. f"assigned to <@{username.discord_member_id}> by a mod using this " \
  251. "command.\n" \
  252. f"- `/{KickstarterCog.kickstarter.name} {KickstarterCog.reset.name} " \
  253. f"{username}` will unlink that Discord record\n" \
  254. "- If all else fails, you can just manually give them the backer role."
  255. await interaction.followup.send(text, ephemeral=True)
  256. return
  257. if backer_role not in member.roles:
  258. await member.add_roles(backer_role)
  259. username.discord_member_id = member.id
  260. username.lookup_status = _LookupStatus.misspelled_username
  261. self.__update_kickstarter_username(username)
  262. text = f"{CONFIG['success_emoji']} Member @{member.name} linked to " \
  263. f"survey username @{username.discord_username} and given " \
  264. f"{backer_role.name} role!"
  265. await interaction.followup.send(text, ephemeral=True)
  266. @kickstarter.command(
  267. description='Resets an imported Kickstarter record.'
  268. )
  269. async def reset(self, interaction: Interaction, username: str):
  270. if await self.__check_disabled(interaction): return
  271. guild = interaction.guild
  272. normal_username = KickstarterCog.__normalize_discord_username(username)
  273. record = self.__fetch_kickstarter_username(guild.id, username=normal_username)
  274. if record is None:
  275. text = f"{CONFIG['failure_emoji']} No record found for `@{normal_username}`."
  276. similar = self.__find_nearest_usernames(guild.id, normal_username, limit=5, max_distance=3)
  277. if len(similar) > 0:
  278. text += " Did you mean one of these?"
  279. for s in similar:
  280. text += f"\n- `@{s.discord_username}`"
  281. await interaction.response.send_message(text, ephemeral=True)
  282. return
  283. old_member_id = record.discord_member_id
  284. old_lookup_status = record.lookup_status
  285. record.discord_member_id = None
  286. record.lookup_status = _LookupStatus.unprocessed
  287. self.__update_kickstarter_username(record)
  288. text = f"{CONFIG['success_emoji']} Record for `@{normal_username}` reset. " \
  289. "(No roles removed from member.)"
  290. if old_member_id is not None:
  291. text += f"\n- `discord_member_id` changed from `{old_member_id}` to `{record.discord_member_id}`"
  292. if old_lookup_status is not None:
  293. text += f"\n- `lookup_status` changed from `{_describe_lookup_status(old_lookup_status)}`" \
  294. f" to `{_describe_lookup_status(record.lookup_status)}`"
  295. await interaction.response.send_message(text, ephemeral=True)
  296. @kickstarter.command(
  297. description='Uploads an export of Kickstarter users.'
  298. )
  299. async def upload(self, interaction: Interaction):
  300. if await self.__check_disabled(interaction): return
  301. if await self.__check_configured(interaction): return
  302. await interaction.response.send_modal(_UploadUsernamesModal())
  303. # @override
  304. async def interaction_check(self, interaction: Interaction) -> bool:
  305. if interaction.command is not None:
  306. self.__trace(interaction.guild, f"@{interaction.user.name} used " \
  307. f"/{interaction.command.qualified_name}")
  308. return True
  309. # -- Events --
  310. @Cog.listener()
  311. async def on_member_join(self, member: Member) -> None:
  312. guild = member.guild
  313. if not self.__is_enabled(guild): return
  314. if self.__get_backer_role_id(guild) is None: return
  315. try:
  316. backer_role = await self.__fetch_backer_role(guild)
  317. except _NoBackerRoleException:
  318. self.log(guild, "Backer role configured but can't be retrieved")
  319. return # Bad id? Role removed?
  320. if backer_role in member.roles:
  321. return # Already has role
  322. username = self.__fetch_kickstarter_username(guild.id, username=member.name)
  323. if username is None:
  324. return # Not on list
  325. if username.lookup_status == _LookupStatus.misspelled_username:
  326. self.log(guild, f"\u0007Member @{member.name} joined but their " \
  327. "username was already manually attached to user id " \
  328. f"{username.discord_member_id}")
  329. return
  330. self.__trace(guild, f"Member @{member.name} joined and is a backer. " \
  331. "Granting backer role.")
  332. await member.add_roles(backer_role)
  333. username.lookup_status = _LookupStatus.member_found
  334. username.discord_member_id = member.id
  335. self.__update_kickstarter_username(username)
  336. # -- UI callbacks -----
  337. async def on_username_upload_submit(self,
  338. interaction: Interaction,
  339. attachment: Attachment
  340. ):
  341. """Callback for username upload modal."""
  342. await interaction.response.defer(ephemeral=True, thinking=True)
  343. guild = interaction.guild
  344. import_result: _ImportResult = await self.__import_usernames(guild, attachment)
  345. try:
  346. sync_result: _SyncUsernamesResult = await self.__sync_by_username(guild)
  347. text = f"{CONFIG['success_emoji']} Username import complete.\n" + \
  348. import_result.summary_markdown() + "\n" + \
  349. sync_result.summary_markdown()
  350. except _NoBackerRoleException:
  351. text = f"{CONFIG['failure_emoji']} Backer role must be configured " \
  352. "before importing."
  353. await interaction.followup.send(text, ephemeral=True)
  354. # -- Operations -----
  355. async def __import_usernames(self,
  356. guild: Guild,
  357. attachment: Attachment
  358. ) -> '_ImportResult':
  359. """Imports Discord usernames from an upload attachment."""
  360. self.__trace(guild, f"Download - start - {attachment.filename} " \
  361. f"({attachment.size:,} bytes, {attachment.content_type})")
  362. file_bytes = await attachment.read()
  363. file_str = file_bytes.decode('utf-8')
  364. self.__trace(guild, "Download - complete")
  365. self.__trace(guild, "Parse - start")
  366. lines = file_str.splitlines(keepends=False)
  367. usernames = []
  368. malformed_count = 0
  369. malformed_usernames = []
  370. for line in lines:
  371. username = self.__normalize_discord_username(line)
  372. if (username.startswith('"') and username.endswith('"')) or \
  373. (username.startswith("'") and username.endswith("'")):
  374. # Remove quotes
  375. username = username[1:-1].strip()
  376. if username == '':
  377. continue
  378. if is_discord_username(username):
  379. usernames.append(username)
  380. else:
  381. self.__trace(guild, f"Not a Discord username: \"{line}\"")
  382. malformed_usernames.append(line)
  383. malformed_count += 1
  384. self.__trace(guild, f"Parse - complete - {len(lines):,} lines, {len(usernames):,} " \
  385. f"valid usernames, {malformed_count:,} malformed usernames")
  386. self.__trace(guild, f"Storing - start - {len(usernames):,} usernames")
  387. new_username_count = self.__store_kickstarter_usernames(guild.id, usernames)
  388. self.__trace(guild, f"Storing - complete - {new_username_count:,} unique " \
  389. "usernames stored")
  390. return _ImportResult(
  391. attachment.filename,
  392. attachment.size,
  393. len(lines),
  394. len(usernames),
  395. malformed_count,
  396. new_username_count,
  397. malformed_usernames
  398. )
  399. async def __sync_by_username(self, guild: Guild) -> '_SyncUsernamesResult':
  400. backer_role = await self.__fetch_backer_role(guild)
  401. complete_count = 0
  402. not_found_count = 0
  403. self.__trace(guild, "Fetch usernames - start")
  404. usernames: list[_KickstarterDiscordUser] = \
  405. self.__fetch_incomplete_kickstarter_usernames(guild.id,
  406. {
  407. _LookupStatus.unprocessed,
  408. _LookupStatus.username_not_found
  409. })
  410. self.__trace(guild, f"Fetch usernames - complete - found {len(usernames):,}")
  411. username_to_member_id: dict[str, int] = {}
  412. if self.bot.intents.members:
  413. self.__trace(guild, "Fetching guild members from API - start")
  414. async for member in guild.fetch_members(limit=None):
  415. username_to_member_id[member.name] = member.id
  416. self.__trace(guild, "Fetching guild members from API - complete - " \
  417. f"got {len(username_to_member_id):,}")
  418. async def username_loop_handler(username: _KickstarterDiscordUser):
  419. nonlocal complete_count
  420. nonlocal not_found_count
  421. member_id = username_to_member_id.get(username.discord_username)
  422. if member_id is not None:
  423. member = guild.get_member(member_id) or \
  424. await guild.fetch_member(member_id)
  425. else:
  426. member = guild.get_member_named(username.discord_username)
  427. if member is None:
  428. not_found_count += 1
  429. if username.lookup_status != _LookupStatus.username_not_found:
  430. username.lookup_status = _LookupStatus.username_not_found
  431. self.__update_kickstarter_username(username)
  432. return
  433. username.lookup_status = _LookupStatus.member_found
  434. username.discord_member_id = member.id
  435. self.__update_kickstarter_username(username)
  436. if backer_role not in member.roles:
  437. await member.add_roles(backer_role)
  438. complete_count += 1
  439. self.__trace(guild, "Sync loop - start")
  440. failure_count = await self.__throttled_loop(guild, usernames,
  441. username_loop_handler, update_seconds=10.0)
  442. self.__trace(guild, f"Sync loop - complete - {complete_count:,} completed, " \
  443. f"{not_found_count:,} not found")
  444. return _SyncUsernamesResult(complete_count, not_found_count, failure_count)
  445. async def __throttled_loop(self,
  446. guild: Guild,
  447. iter: Iterable[_T],
  448. callback: Callable[[_T], Awaitable[None]],
  449. update_seconds: float | None = None
  450. ) -> int:
  451. """Iterates a loop with automatically adjusting sleeps based on
  452. throttling exceptions."""
  453. failure_count = 0
  454. sleep_length = 0.0
  455. start_time = now_timestamp()
  456. last_update_time = start_time
  457. for iter_count, elem in enumerate(iter):
  458. complete = False
  459. for _ in range(5):
  460. try:
  461. await sleep(sleep_length)
  462. await callback(elem)
  463. complete = True
  464. break
  465. except HTTPException as ex:
  466. if ex.status == 429: # rate limited
  467. retry_header_value = ex.response.headers.get('retry_after')
  468. retry_after_millis = float(retry_header_value or '1000')
  469. self.__trace(guild, "Rate limited while processing. " \
  470. f"retry_after={retry_header_value}")
  471. await sleep(retry_after_millis / 1000.0)
  472. sleep_length = 1.0 if sleep_length == 0.0 else sleep_length * 2.0
  473. self.__trace(guild, f"Sleep increased to {sleep_length:,}s " \
  474. "due to rate limiting")
  475. else:
  476. dump_stacktrace(ex)
  477. except DiscordException as ex:
  478. dump_stacktrace(ex)
  479. if not complete:
  480. failure_count += 1
  481. if update_seconds is not None and now_timestamp() - last_update_time >= update_seconds:
  482. self.__trace(guild, f"Completed {iter_count + 1:,} iterations")
  483. last_update_time = now_timestamp()
  484. return failure_count
  485. # -- Database functions -----
  486. def __store_kickstarter_username(self, record: '_KickstarterDiscordUser'):
  487. cur = self.con.cursor()
  488. cur.execute("""
  489. INSERT OR IGNORE INTO kickstarter_discord_users (
  490. guild_id,
  491. discord_username,
  492. discord_member_id,
  493. lookup_status,
  494. imported_at
  495. ) VALUES (
  496. :guild_id,
  497. :discord_username,
  498. :discord_member_id,
  499. :lookup_status,
  500. :imported_at
  501. )
  502. """, {
  503. 'guild_id': record.guild_id,
  504. 'discord_username': self.__normalize_discord_username(record.discord_username),
  505. 'discord_member_id': record.discord_member_id,
  506. 'lookup_status': record.lookup_status,
  507. 'imported_at': record.imported_at
  508. })
  509. row_id = cur.lastrowid
  510. self.con.commit()
  511. cur.close()
  512. if row_id is not None and row_id != 0:
  513. record.pk = row_id
  514. def __store_kickstarter_usernames(self, guild_id: int, usernames: list[str]) -> int:
  515. imported_at: int = int(now_timestamp())
  516. cur = self.con.cursor()
  517. for username in usernames:
  518. cur.execute("""
  519. INSERT OR IGNORE INTO kickstarter_discord_users (
  520. guild_id,
  521. discord_username,
  522. imported_at
  523. ) VALUES (
  524. :guild_id,
  525. :discord_username,
  526. :imported_at
  527. )
  528. """, {
  529. 'guild_id': guild_id,
  530. 'discord_username': self.__normalize_discord_username(username),
  531. 'imported_at': imported_at
  532. })
  533. cur.execute("""
  534. SELECT COUNT(1)
  535. FROM kickstarter_discord_users
  536. WHERE imported_at = ?
  537. """, (imported_at, ))
  538. imported_count = cur.fetchone()[0]
  539. self.con.commit()
  540. cur.close()
  541. return imported_count
  542. def __fetch_kickstarter_username(self,
  543. guild_id: int,
  544. *,
  545. member_id: int | None = None,
  546. username: str | None = None
  547. ) -> Optional['_KickstarterDiscordUser']:
  548. """Fetches an imported Discord username by EITHER member id or username
  549. (must provide exactly one)"""
  550. cur = self.con.cursor()
  551. cur.row_factory = lambda c, r: _KickstarterDiscordUser(*r)
  552. if member_id is not None:
  553. cur.execute("""
  554. SELECT *
  555. FROM kickstarter_discord_users
  556. WHERE
  557. guild_id = :guild_id
  558. AND discord_member_id = :member_id
  559. """, { 'guild_id': guild_id, 'member_id': member_id })
  560. elif username is not None:
  561. cur.execute("""
  562. SELECT *
  563. FROM kickstarter_discord_users
  564. WHERE
  565. guild_id = :guild_id
  566. AND discord_username = :username
  567. """, { 'guild_id': guild_id, 'username': username })
  568. ret_val = cur.fetchone()
  569. cur.close()
  570. return ret_val
  571. def __fetch_incomplete_kickstarter_usernames(self,
  572. guild_id: int,
  573. statuses: set['_LookupStatus']
  574. ) -> list['_KickstarterDiscordUser']:
  575. cur = self.con.cursor()
  576. cur.row_factory = lambda c, r: _KickstarterDiscordUser(*r)
  577. params = [ guild_id ] + list(statuses)
  578. status_placeholders = ', '.join('?' * len(statuses))
  579. sql = f"""
  580. SELECT *
  581. FROM kickstarter_discord_users
  582. WHERE
  583. guild_id = ?
  584. AND lookup_status IN ({status_placeholders})
  585. """
  586. cur.execute(sql, params)
  587. ret_val = cur.fetchall()
  588. cur.close()
  589. return ret_val
  590. def __update_kickstarter_username(self, user: '_KickstarterDiscordUser'):
  591. cur = self.con.cursor()
  592. cur.execute("""
  593. UPDATE kickstarter_discord_users
  594. SET discord_member_id = :discord_member_id,
  595. lookup_status = :lookup_status
  596. WHERE pk = :pk
  597. """, {
  598. 'discord_member_id': user.discord_member_id,
  599. 'lookup_status': user.lookup_status,
  600. 'pk': user.pk
  601. })
  602. self.con.commit()
  603. cur.close()
  604. def __fetch_stats(self, guild_id: int) -> '_Stats':
  605. """Returns member link stats."""
  606. cur = self.con.cursor()
  607. cur.execute("""
  608. SELECT
  609. COUNT(1) AS total,
  610. SUM(IIF(lookup_status = 0, 1, 0)) AS unprocessed_count,
  611. SUM(IIF(lookup_status = 1, 1, 0)) AS not_found_count,
  612. SUM(IIF(lookup_status = 2, 1, 0)) AS found_count,
  613. MAX(imported_at) AS last_import
  614. FROM kickstarter_discord_users
  615. """)
  616. (
  617. username_count,
  618. username_unprocessed_count,
  619. username_not_found_count,
  620. username_found_count,
  621. username_last_imported_at
  622. ) = cur.fetchone()
  623. cur.close()
  624. return _Stats(
  625. username_count,
  626. username_unprocessed_count,
  627. username_not_found_count,
  628. username_found_count,
  629. username_last_imported_at
  630. )
  631. def __find_nearest_usernames(self,
  632. guild_id: int,
  633. username: str,
  634. limit: int = 10,
  635. max_distance: int = 999
  636. ) -> list['_KickstarterDiscordUser']:
  637. cur = self.con.cursor()
  638. cur.row_factory = lambda c, r: _KickstarterDiscordUser(*r)
  639. cur.execute("""
  640. SELECT *
  641. FROM kickstarter_discord_users
  642. WHERE guild_id = ?
  643. AND lookup_status IN (?, ?, ?)
  644. """, (
  645. guild_id,
  646. _LookupStatus.unprocessed,
  647. _LookupStatus.username_not_found,
  648. _LookupStatus.misspelled_username,
  649. ))
  650. records: list[_KickstarterDiscordUser] = cur.fetchall()
  651. cur.close()
  652. normal_username = KickstarterCog.__normalize_discord_username(username)
  653. in_bounds_count = 0
  654. def compute(r: _KickstarterDiscordUser) -> int:
  655. nonlocal in_bounds_count
  656. score = levenshtein(normal_username, r.discord_username)
  657. if score <= max_distance:
  658. in_bounds_count += 1
  659. return score
  660. closest = heapq.nsmallest(limit, records, key=compute)
  661. if len(closest) > in_bounds_count:
  662. closest = closest[:in_bounds_count]
  663. return closest
  664. # -- Utils -----
  665. def __trace(self, guild: Guild, message: str):
  666. self.log(guild, message)
  667. @staticmethod
  668. def __normalize_discord_username(username: str) -> str:
  669. norm = username.lower().strip()
  670. if norm.startswith('@'):
  671. norm = norm[1:].strip()
  672. return norm
  673. class _LookupStatus(IntEnum):
  674. # Username imported but has not been looked up yet
  675. unprocessed = 0
  676. # Username was searched for in guild but not found (user might not have joined yet)
  677. username_not_found = 1
  678. # Username found in guild and member id stored
  679. member_found = 2
  680. # Mod manually linked this imported username to a Discord member to correct a username typo
  681. misspelled_username = 3
  682. def _describe_lookup_status(status: _LookupStatus) -> str:
  683. if status == _LookupStatus.unprocessed:
  684. return 'unprocessed'
  685. if status == _LookupStatus.username_not_found:
  686. return 'member not in server'
  687. if status == _LookupStatus.member_found:
  688. return 'member found'
  689. if status == _LookupStatus.misspelled_username:
  690. return 'manually linked by mod'
  691. return '-unknown-'
  692. class _KickstarterDiscordUser:
  693. """kickstarter_discord_users table row"""
  694. def __init__(self,
  695. pk: int,
  696. guild_id: int,
  697. discord_username: str,
  698. discord_member_id: int | None,
  699. lookup_status: _LookupStatus,
  700. imported_at: int
  701. ):
  702. self.pk: int = pk
  703. self.guild_id: int = guild_id
  704. self.discord_username: str = discord_username
  705. """Discord username provided in Kickstarter survey."""
  706. self.discord_member_id: int | None = discord_member_id
  707. """Discord member ID when synced successfully or None when not yet linked."""
  708. self.lookup_status: _LookupStatus = lookup_status
  709. self.imported_at: int = imported_at
  710. """Unix timestamp when this record was first imported."""
  711. class _Stats:
  712. def __init__(self,
  713. username_count: int,
  714. username_unprocessed_count: int,
  715. username_not_found_count: int,
  716. username_found_count: int,
  717. username_last_imported_at: int
  718. ):
  719. self.username_count: int = username_count
  720. self.username_unprocessed_count: int = username_unprocessed_count
  721. self.username_not_found_count: int = username_not_found_count
  722. self.username_found_count: int = username_found_count
  723. self.username_last_imported_at: int = username_last_imported_at
  724. class _NoBackerRoleException(BaseException):
  725. pass
  726. class _ImportResult:
  727. def __init__(self,
  728. filename: str,
  729. file_bytes: int,
  730. line_count: int,
  731. valid_record_count: int,
  732. malformed_record_count: int,
  733. new_record_count: int,
  734. invalid_records: list[str]
  735. ):
  736. self.filename: str = filename
  737. """Filename of the uploaded file."""
  738. self.file_bytes: int = file_bytes
  739. """Size of uploaded file in bytes."""
  740. self.line_count: int = line_count
  741. """Number of lines in the uploaded text file."""
  742. self.valid_record_count: int = valid_record_count
  743. """How many records were valid."""
  744. self.malformed_record_count: int = malformed_record_count
  745. """How many records were skipped because they were malformed."""
  746. self.new_record_count: int = new_record_count
  747. """How many records were imported. May be less than valid_record_count
  748. if some were already imported."""
  749. self.invalid_records: list[str] = invalid_records
  750. """List of records that could not be imported. (May be partial if lots of failures.)"""
  751. def summary_markdown(self) -> str:
  752. lines: list[str] = []
  753. if self.valid_record_count > 0:
  754. lines.append(f"- Read {self.valid_record_count:,} valid records")
  755. else:
  756. lines.append("- Upload contained **no valid records**")
  757. if self.malformed_record_count > 0:
  758. lines.append(f"- Read **{self.malformed_record_count:,} malformed records**")
  759. if self.new_record_count > 0:
  760. lines.append(f"- Imported {self.new_record_count:,} new unique records")
  761. else:
  762. lines.append("- No new unique records (all previously imported)")
  763. return "\n".join(lines)
  764. class _SyncUsernamesResult:
  765. def __init__(self,
  766. complete_count: int,
  767. not_found_count: int,
  768. failure_count: int
  769. ):
  770. self.complete_count: int = complete_count
  771. self.not_found_count: int = not_found_count
  772. self.failure_count: int = failure_count
  773. def summary_markdown(self) -> str:
  774. lines: list[str] = []
  775. if self.complete_count > 0:
  776. lines.append(f"- Gave backer role to {self.complete_count:,} new members")
  777. else:
  778. lines.append("- No new members found")
  779. if self.not_found_count > 0:
  780. lines.append(f"- {self.not_found_count:,} members could not be found " \
  781. "by the provided username")
  782. if self.failure_count > 0:
  783. lines.append(f"- Failed to link {self.failure_count:,} members")
  784. return "\n".join(lines)
  785. class _UploadUsernamesModal(Modal):
  786. upload_label = Label(
  787. text='Discord user export',
  788. description='Upload a plain text file containing one backer Discord username per line.',
  789. component=FileUpload(
  790. required=True,
  791. min_values=1, max_values=1
  792. )
  793. )
  794. def __init__(self):
  795. super().__init__(title='Upload Discord Usernames', timeout=None)
  796. # @override
  797. async def on_submit(self, interaction: Interaction) -> None:
  798. # noinspection PyTypeChecker
  799. upload_input: FileUpload = self.upload_label.component
  800. if len(upload_input.values) < 1:
  801. text = f"{CONFIG['failure_emoji']} No export file included"
  802. await interaction.response.send(text, ephemeral=True)
  803. return
  804. attachment = upload_input.values[0]
  805. await KickstarterCog.shared.on_username_upload_submit(interaction, attachment)
  806. # @override
  807. async def on_error(self, interaction: Interaction, error: Exception) -> None:
  808. dump_stacktrace(error)
  809. try:
  810. await interaction.response.send_message(
  811. f'{CONFIG["failure_emoji"]} Upload failed :(',
  812. ephemeral=True,
  813. )
  814. except DiscordException as e:
  815. dump_stacktrace(e)