|
| 1 | +import asyncio |
| 2 | +import datetime |
| 3 | +import logging |
| 4 | +from queue import Empty |
| 5 | +from typing import Union |
| 6 | + |
| 7 | +import discord |
| 8 | +from discord.ext import commands, tasks |
| 9 | + |
| 10 | +from tgtg_scanner.errors import DiscordConfigurationError, MaskConfigurationError |
| 11 | +from tgtg_scanner.models import Config, Favorites, Item, Reservations |
| 12 | +from tgtg_scanner.models.reservations import Reservation |
| 13 | +from tgtg_scanner.notifiers.base import Notifier |
| 14 | + |
| 15 | +log = logging.getLogger("tgtg") |
| 16 | + |
| 17 | + |
| 18 | +class Discord(Notifier): |
| 19 | + """Notifier for Discord""" |
| 20 | + |
| 21 | + def __init__(self, config: Config, reservations: Reservations, favorites: Favorites): |
| 22 | + super().__init__(config, reservations, favorites) |
| 23 | + self.enabled = config.discord.enabled |
| 24 | + self.prefix = config.discord.prefix |
| 25 | + self.token = config.discord.token |
| 26 | + self.channel = config.discord.channel |
| 27 | + self.body = config.discord.body |
| 28 | + self.disable_commands = config.discord.disable_commands |
| 29 | + self.cron = config.discord.cron |
| 30 | + self.mute: Union[datetime.datetime, None] = None |
| 31 | + self.bot_id = None |
| 32 | + self.channel_id = None |
| 33 | + self.server_id = None |
| 34 | + |
| 35 | + if self.enabled: |
| 36 | + if self.token is None or self.channel == 0: |
| 37 | + raise DiscordConfigurationError() |
| 38 | + try: |
| 39 | + Item.check_mask(self.body) |
| 40 | + except MaskConfigurationError as exc: |
| 41 | + raise DiscordConfigurationError(exc.message) from exc |
| 42 | + |
| 43 | + try: |
| 44 | + # Setting event loop explicitly for python 3.9 compatibility |
| 45 | + loop = asyncio.new_event_loop() |
| 46 | + asyncio.set_event_loop(loop) |
| 47 | + self.bot = commands.Bot(command_prefix=self.prefix, intents=discord.Intents.all()) |
| 48 | + except MaskConfigurationError as exc: |
| 49 | + raise DiscordConfigurationError(exc.message) from exc |
| 50 | + |
| 51 | + async def _send(self, item: Union[Item, Reservation]) -> None: # type: ignore[override] |
| 52 | + """Sends item information using Discord bot""" |
| 53 | + if self.mute and self.mute > datetime.datetime.now(): |
| 54 | + return |
| 55 | + if self.mute: |
| 56 | + log.info("Reactivated Discord Notifications") |
| 57 | + self.mute = None |
| 58 | + if isinstance(item, Item): |
| 59 | + message = item.unmask(self.body) |
| 60 | + self.bot.dispatch("send_notification", message) |
| 61 | + |
| 62 | + @tasks.loop(seconds=1) |
| 63 | + async def _listen_for_items(self): |
| 64 | + """Method for polling notifications every second""" |
| 65 | + try: |
| 66 | + item = self.queue.get(block=False) |
| 67 | + if item is None: |
| 68 | + self.bot.dispatch("close") |
| 69 | + return |
| 70 | + log.debug("Sending %s Notification", self.name) |
| 71 | + await self._send(item) |
| 72 | + except Empty: |
| 73 | + pass |
| 74 | + except Exception as exc: |
| 75 | + log.error("Failed sending %s: %s", self.name, exc) |
| 76 | + |
| 77 | + def _run(self): |
| 78 | + async def _start_bot() -> None: |
| 79 | + await self.bot.start(self.token) |
| 80 | + |
| 81 | + # Events include methods for post-init, shutting down, and notification sending |
| 82 | + self._setup_events() |
| 83 | + if not self.disable_commands: |
| 84 | + # Commands are handled separately, in case commands are not enabled |
| 85 | + self._setup_commands() |
| 86 | + |
| 87 | + # Setting event loop explicitly for python 3.9 compatibility |
| 88 | + loop = asyncio.new_event_loop() |
| 89 | + asyncio.set_event_loop(loop) |
| 90 | + self.config.set_locale() |
| 91 | + asyncio.run(_start_bot()) |
| 92 | + |
| 93 | + def _setup_events(self): |
| 94 | + @self.bot.event |
| 95 | + async def on_ready(): |
| 96 | + """Callback after successful login (only explicitly used in test_notifiers.py)""" |
| 97 | + self.bot_id = self.bot.user.id |
| 98 | + self.channel_id = self.channel |
| 99 | + self.server_id = self.bot.guilds[0].id if len(self.bot.guilds) > 0 else 0 |
| 100 | + self._listen_for_items.start() |
| 101 | + |
| 102 | + @self.bot.event |
| 103 | + async def on_send_notification(message): |
| 104 | + """Callback for item notification""" |
| 105 | + channel = self.bot.get_channel(self.channel) or await self.bot.fetch_channel(self.channel) |
| 106 | + if channel: |
| 107 | + await channel.send(message) |
| 108 | + |
| 109 | + @self.bot.event |
| 110 | + async def on_close(): |
| 111 | + """Logout from Discord (only explicitly used in test_notifiers.py)""" |
| 112 | + await self.bot.close() |
| 113 | + |
| 114 | + def _setup_commands(self): |
| 115 | + @self.bot.command(name="mute") |
| 116 | + async def _mute(ctx, *args): |
| 117 | + """Deactivates Discord Notifications for x days""" |
| 118 | + days = int(args[0]) if len(args) > 0 and args[0].isnumeric() else 1 |
| 119 | + self.mute = datetime.datetime.now() + datetime.timedelta(days=days) |
| 120 | + log.info("Deactivated Discord Notifications for %s day(s)", days) |
| 121 | + log.info("Reactivation at %s", self.mute) |
| 122 | + await ctx.send( |
| 123 | + f"Deactivated Discord notifications for {days} days.\nReactivating at {self.mute} or use `{self.prefix}unmute`." |
| 124 | + ) |
| 125 | + |
| 126 | + @self.bot.command(name="unmute") |
| 127 | + async def _unmute(ctx): |
| 128 | + """Reactivate Discord notifications""" |
| 129 | + self.mute = None |
| 130 | + log.info("Reactivated Discord notifications") |
| 131 | + await ctx.send("Reactivated Discord notifications") |
| 132 | + |
| 133 | + @self.bot.command(name="listfavorites") |
| 134 | + async def _list_favorites(ctx): |
| 135 | + """List favorites using display name""" |
| 136 | + favorites = self.favorites.get_favorites() |
| 137 | + if not favorites: |
| 138 | + await ctx.send("You currently don't have any favorites.") |
| 139 | + else: |
| 140 | + await ctx.send("\n".join([f"• {item.item_id} - {item.display_name}" for item in favorites])) |
| 141 | + |
| 142 | + @self.bot.command(name="listfavoriteids") |
| 143 | + async def _list_favorite_ids(ctx): |
| 144 | + """List favorites using id""" |
| 145 | + favorites = self.favorites.get_favorites() |
| 146 | + if not favorites: |
| 147 | + await ctx.send("You currently don't have any favorites.") |
| 148 | + else: |
| 149 | + await ctx.send(" ".join([item.item_id for item in favorites])) |
| 150 | + |
| 151 | + @self.bot.command(name="addfavorites") |
| 152 | + async def _add_favorites(ctx, *args): |
| 153 | + """Add favorite(s)""" |
| 154 | + item_ids = list( |
| 155 | + filter( |
| 156 | + lambda x: x.isdigit() and int(x) != 0, |
| 157 | + map( |
| 158 | + str.strip, |
| 159 | + [split_args for arg in args for split_args in arg.split(",")], |
| 160 | + ), |
| 161 | + ) |
| 162 | + ) |
| 163 | + if not item_ids: |
| 164 | + await ctx.channel.send( |
| 165 | + "Please supply item ids in one of the following ways: " |
| 166 | + f"'{self.prefix}addfavorites 12345 23456 34567' or " |
| 167 | + f"'{self.prefix}addfavorites 12345,23456,34567'" |
| 168 | + ) |
| 169 | + return |
| 170 | + |
| 171 | + self.favorites.add_favorites(item_ids) |
| 172 | + await ctx.send(f"Added the following item ids to favorites: {' '.join(item_ids)}") |
| 173 | + log.debug('Added the following item ids to favorites: "%s"', item_ids) |
| 174 | + |
| 175 | + @self.bot.command(name="removefavorites") |
| 176 | + async def _remove_favorites(ctx, *args): |
| 177 | + """Remove favorite(s)""" |
| 178 | + item_ids = list( |
| 179 | + filter( |
| 180 | + lambda x: x.isdigit() and int(x) != 0, |
| 181 | + map( |
| 182 | + str.strip, |
| 183 | + [split_args for arg in args for split_args in arg.split(",")], |
| 184 | + ), |
| 185 | + ) |
| 186 | + ) |
| 187 | + if not item_ids: |
| 188 | + await ctx.channel.send( |
| 189 | + "Please supply item ids in one of the following ways: " |
| 190 | + f"'{self.prefix}removefavorites 12345 23456 34567' or " |
| 191 | + f"'{self.prefix}removefavorites 12345,23456,34567'" |
| 192 | + ) |
| 193 | + return |
| 194 | + |
| 195 | + self.favorites.remove_favorite(item_ids) |
| 196 | + await ctx.send(f"Removed the following item ids from favorites: {' '.join(item_ids)}") |
| 197 | + log.debug('Removed the following item ids from favorites: "%s"', item_ids) |
| 198 | + |
| 199 | + @self.bot.command(name="gettoken") |
| 200 | + async def _get_token(ctx): |
| 201 | + """Display token used to login (without needing to manually check in config.ini)""" |
| 202 | + await ctx.send(f"Token in use: {self.token}") |
| 203 | + |
| 204 | + @self.bot.command(name="getinfo") |
| 205 | + async def _get_info(ctx): |
| 206 | + """Display basic info about connection""" |
| 207 | + bot_id = ctx.me.id |
| 208 | + bot_name = ctx.me.display_name |
| 209 | + bot_mention = ctx.me.mention |
| 210 | + joined_at = ctx.me.joined_at |
| 211 | + channel_id = ctx.channel.id |
| 212 | + channel_name = ctx.channel.name |
| 213 | + guild_id = ctx.guild.id |
| 214 | + guild_name = ctx.guild.name |
| 215 | + |
| 216 | + response = ( |
| 217 | + f"Hi! I'm {bot_mention}, the TGTG Bot on this server. I joined at {joined_at}\n" |
| 218 | + f"```Bot (ID): {bot_name} ({bot_id})\n" |
| 219 | + f"Channel (ID): {channel_name} ({channel_id})\n" |
| 220 | + f"Server (ID): {guild_name} ({guild_id})```" |
| 221 | + ) |
| 222 | + |
| 223 | + await ctx.send(response) |
| 224 | + |
| 225 | + def __repr__(self) -> str: |
| 226 | + return f"Discord: Channel ID {self.channel}" |
0 commit comments