import base64 import logging import os import re import sys from pathlib import Path from types import TracebackType from typing import Any, Final, cast import dotenv from httpx import Client, HTTPStatusError, Response log = logging.getLogger("botstrap") # Note this instance will not have the .trace level # TODO: Remove once better error handling for constants.py is in place. if (dotenv.dotenv_values().get("BOT_TOKEN") or os.getenv("BOT_TOKEN")) is None: msg = ( "Couldn't find the `BOT_TOKEN` environment variable. " "Make sure to add it to your `.env` file like this: `BOT_TOKEN=value_of_your_bot_token`" ) log.fatal(msg) sys.exit(1) # Filter out the send typing monkeypatch logs from bot core when we import to get constants logging.getLogger("pydis_core").setLevel(logging.WARNING) # As a side effect, this also configures our logging styles from bot.constants import ( # noqa: E402 Bot as BotConstants, Guild as GuildConstants, Webhooks, _Categories, # pyright: ignore[reportPrivateUsage] _Channels, # pyright: ignore[reportPrivateUsage] _Emojis, # pyright: ignore[reportPrivateUsage] _Roles, # pyright: ignore[reportPrivateUsage] ) # Silence noisy loggers logging.getLogger("httpx").setLevel(logging.WARNING) logging.getLogger("httpcore").setLevel(logging.WARNING) ENV_FILE = Path(".env.server") COMMUNITY_FEATURE = "COMMUNITY" ANNOUNCEMENTS_CHANNEL_NAME = "announcements" RULES_CHANNEL_NAME = "rules" GUILD_CATEGORY_TYPE = 4 EMOJI_REGEX = re.compile(r"<:(\w+):(\d+)>") MINIMUM_FLAGS: int = ( 1 << 15 # guild_members_limited | 1 << 19 # message_content_limited ) if GuildConstants.id == type(GuildConstants).model_fields["id"].default: msg = ( "Couldn't find the `GUILD_ID` environment variable. " "Make sure to add it to your `.env` file like this: `GUILD_ID=value_of_your_discord_server_id`" ) log.error(msg) sys.exit(1) class SilencedDict[T](dict[str, T]): """A dictionary that silences KeyError exceptions upon subscription to non existent items.""" def __init__(self, name: str): self.name: str = name super().__init__() def __getitem__(self, item: str): try: return super().__getitem__(item) except KeyError: log.fatal("Couldn't find key: %s in dict: %s", item, self.name) log.warning( "Please follow our contribution guidelines " "https://pydis.com/contributing-bot " "to guarantee a successful run of botstrap " ) sys.exit(-1) class BotstrapError(Exception): """Raised when an error occurs during the botstrap process.""" class DiscordClient(Client): """An HTTP client to communicate with Discord's APIs.""" CDN_BASE_URL: Final[str] = "https://cdn.discordapp.com" def __init__(self, *, guild_id: int | str, bot_token: str): super().__init__( base_url="https://discord.com/api/v10", headers={"Authorization": f"Bot {bot_token}"}, event_hooks={"response": [self._raise_for_status]}, ) self.guild_id: int | str = guild_id self._app_info: dict[str, object] | None = None self._guild_info: dict[str, object] | None = None self._guild_channels: list[dict[str, object]] | None = None @staticmethod def _raise_for_status(response: Response) -> None: response.raise_for_status() def get_guild_info(self) -> dict[str, Any]: """Fetches the guild's information.""" if self._guild_info is None: response = self.get(f"/guilds/{self.guild_id}") self._guild_info = cast("dict[str, object]", response.json()) return self._guild_info def get_guild_channels(self) -> list[dict[str, Any]]: """Fetches the guild's channels.""" if self._guild_channels is None: response = self.get(f"/guilds/{self.guild_id}/channels") self._guild_channels = cast("list[dict[str, object]]", response.json()) return self._guild_channels def get_channel(self, id_: int | str) -> dict[str, object]: """Fetches a channel by its ID.""" for channel in self.get_guild_channels(): if channel["id"] == str(id_): return channel raise KeyError(f"Channel with ID {id_} not found.") def get_app_info(self) -> dict[str, Any]: """Fetches the application's information.""" if self._app_info is None: response = self.get("/applications/@me") self._app_info = cast("dict[str, object]", response.json()) return self._app_info def upgrade_application_flags_if_necessary(self) -> bool: """ Set the app's flags to allow the intents that we need. Returns a boolean defining whether changes were made. """ # Fetch first to modify, not overwrite current_flags = self.get_app_info().get("flags", 0) new_flags = current_flags | MINIMUM_FLAGS if new_flags != current_flags: resp = self.patch("/applications/@me", json={"flags": new_flags}) self._app_info = cast("dict[str, object]", resp.json()) return True return False def check_if_in_guild(self) -> bool: """Check if the bot is a member of the guild.""" try: self.get_guild_info() except HTTPStatusError as e: if e.response.status_code == 403 or e.response.status_code == 404: return False raise return True def upgrade_server_to_community_if_necessary( self, rules_channel_id: int | str, announcements_channel_id: int | str, ) -> bool: """Fetches server info & upgrades to COMMUNITY if necessary.""" payload = self.get_guild_info() if COMMUNITY_FEATURE not in payload["features"]: log.info("This server is currently not a community, upgrading.") payload["features"].append(COMMUNITY_FEATURE) payload["rules_channel_id"] = rules_channel_id payload["public_updates_channel_id"] = announcements_channel_id self._guild_info = self.patch(f"/guilds/{self.guild_id}", json=payload).json() log.info("Server %s has been successfully updated to a community.", self.guild_id) return True return False def get_all_roles(self) -> dict[str, int]: """Fetches all the roles in a guild.""" result = SilencedDict[int](name="Roles dictionary") roles = self.get_guild_info()["roles"] for role in roles: name = "_".join(part.lower() for part in role["name"].split(" ")).replace("-", "_") result[name] = role["id"] return result def get_all_channels_and_categories(self) -> tuple[dict[str, str], dict[str, str]]: """Fetches all the text channels & categories in a guild.""" off_topic_channel_name_regex = r"ot\d{1}(_.*)+" off_topic_count = 0 channels = SilencedDict[str](name="Channels dictionary") categories = SilencedDict[str](name="Categories dictionary") for channel in self.get_guild_channels(): channel_type = channel["type"] name = "_".join(part.lower() for part in channel["name"].split(" ")).replace("-", "_") if re.match(off_topic_channel_name_regex, name): name = f"off_topic_{off_topic_count}" off_topic_count += 1 if channel_type == GUILD_CATEGORY_TYPE: categories[name] = channel["id"] else: channels[name] = channel["id"] return channels, categories def get_all_guild_webhooks(self) -> list[dict[str, Any]]: """Lists all the webhooks for the guild.""" response = self.get(f"/guilds/{self.guild_id}/webhooks") return response.json() def create_webhook(self, name: str, channel_id: int | str) -> str: """Creates a new webhook for a particular channel.""" payload = {"name": name} response = self.post( f"/channels/{channel_id}/webhooks", json=payload, headers={"X-Audit-Log-Reason": "Creating webhook as part of PyDis botstrap"}, ) new_webhook = response.json() log.info("Creating webhook: %s has been successfully created.", name) return new_webhook["id"] def list_emojis(self) -> list[dict[str, Any]]: """Lists all the emojis for the guild.""" response = self.get(f"/guilds/{self.guild_id}/emojis") return response.json() def get_emoji_contents(self, id_: str | int) -> bytes | None: """Fetches the image data for an emoji by ID.""" # Emojis are located at https://cdn.discordapp.com/emojis/{emoji_id}.{ext} response = self.get(f"{self.CDN_BASE_URL}/emojis/{id_!s}.webp") return response.content def clone_emoji(self, *, new_name: str, original_emoji_id: str | int) -> str: """Creates a new emoji in the guild, cloned from another emoji by ID.""" emoji_data = self.get_emoji_contents(original_emoji_id) if not emoji_data: log.warning("Couldn't find emoji with ID %s.", original_emoji_id) return "" image_data = base64.b64encode(emoji_data).decode("utf-8") payload = { "name": new_name, "image": f"data:image/webp;base64,{image_data}", } response = self.post( f"/guilds/{self.guild_id}/emojis", json=payload, headers={"X-Audit-Log-Reason": "Creating emoji as part of PyDis botstrap"}, ) new_emoji = response.json() return new_emoji["id"] class BotStrapper: """Bootstrap the bot configuration for a given guild.""" def __init__( self, *, guild_id: int | str, env_file: Path, bot_token: str, ): self.guild_id: int | str = guild_id self.client: DiscordClient = DiscordClient(guild_id=guild_id, bot_token=bot_token) self.env_file: Path = env_file self.client_id = self.client.get_app_info()["id"] def __enter__(self): return self def __exit__( self, exc_type: type[BaseException] | None = None, exc_value: BaseException | None = None, traceback: TracebackType | None = None, ) -> None: self.client.__exit__(exc_type, exc_value, traceback) def _find_webhook_for_channel( self, *, webhook_name: str, webhook_id: int | str, channel_id: int | str, webhooks: list[dict[str, object]], ) -> dict[str, object] | None: """Find a usuable webhook by its name or ID from a list of existing webhooks.""" matches: list[dict[str, object]] = [] # This matches to a list to prefer webhooks created by this application first, # which may be encountered AFTER reaching a webhook with the same name. for webhook in webhooks: if not webhook.get("token"): continue # Webhook is unusable without a token if webhook["id"] == str(webhook_id): return webhook # Keep existing configuration if webhook["channel_id"] != str(channel_id): continue # Exclude webhooks from other channels if webhook["application_id"] == str(self.client_id): matches.insert(0, webhook) # Prefer webhooks created by this application elif webhook["name"] == webhook_name: if webhook in matches: return webhook # owned, and the name is the name, Use this one matches.append(webhook) # Fallback to matching by name return matches[0] if matches else None def upgrade_client(self) -> bool: """Upgrade the application's flags if necessary.""" if self.client.upgrade_application_flags_if_necessary(): log.info("Application flags upgraded successfully, and necessary intents are now enabled.") return True return False def check_guild_membership(self) -> None: """Check the bot is in the required guild.""" if not self.client.check_if_in_guild(): log.error("The bot is not a member of the configured guild with ID %s.", self.guild_id) log.warning( "Please invite with the following URL and rerun this script: " "https://discord.com/oauth2/authorize?client_id=%s&guild_id=%s&scope=bot+applications.commands&permissions=8", self.client_id, self.guild_id, ) raise BotstrapError("Bot is not a member of the configured guild.") def upgrade_guild(self, announcements_channel_id: str, rules_channel_id: str) -> bool: """Upgrade the guild to a community if necessary.""" return self.client.upgrade_server_to_community_if_necessary( rules_channel_id=rules_channel_id, announcements_channel_id=announcements_channel_id, ) def get_roles(self) -> dict[str, Any]: """Get a config map of all of the roles in the guild.""" log.debug("Syncing roles with bot configuration.") all_roles = self.client.get_all_roles() data: dict[str, int] = {} for role_name in _Roles.model_fields: role_id = all_roles.get(role_name, None) if not role_id: log.warning("Couldn't find the role %s in the guild, PyDis' default values will be used.", role_name) continue data[role_name] = role_id return data def get_channels(self) -> dict[str, Any]: """Get a config map of all of the channels in the guild.""" log.debug("Syncing channels with bot configuration.") all_channels, _categories = self.client.get_all_channels_and_categories() data: dict[str, str] = {} for channel_name in _Channels.model_fields: channel_id = all_channels.get(channel_name, None) if not channel_id: log.warning( "Couldn't find the channel %s in the guild, PyDis' default values will be used.", channel_name ) continue data[channel_name] = channel_id return data def get_categories(self) -> dict[str, str]: """Get a config map of all of the categories in guild.""" log.debug("Syncing categories with bot configuration.") _channels, all_categories = self.client.get_all_channels_and_categories() data: dict[str, str] = {} for category_name in _Categories.model_fields: category_id = all_categories.get(category_name, None) if not category_id: log.warning( "Couldn't find the category %s in the guild, PyDis' default values will be used.", category_name ) continue data[category_name] = category_id return data def sync_webhooks(self) -> dict[str, object]: """Get webhook config. Will create all webhooks that cannot be found.""" log.debug("Syncing webhooks with bot configuration.") all_channels, _categories = self.client.get_all_channels_and_categories() data: dict[str, object] = {} existing_webhooks = self.client.get_all_guild_webhooks() for webhook_name, configured_webhook in Webhooks.model_dump().items(): formatted_webhook_name = webhook_name.replace("_", " ").title() webhook_channel_id = all_channels[webhook_name] webhook = self._find_webhook_for_channel( webhook_name= formatted_webhook_name, webhook_id=configured_webhook["id"], channel_id=webhook_channel_id, webhooks=existing_webhooks, ) if not webhook: webhook_id = self.client.create_webhook(formatted_webhook_name, webhook_channel_id) else: webhook_id = webhook["id"] data[webhook_name + "__id"] = webhook_id return data def sync_emojis(self) -> dict[str, str]: """Get emoji config. Will create all emojis that cannot be found.""" existing_emojis = self.client.list_emojis() log.debug("Syncing emojis with bot configuration.") data: dict[str, str] = {} for emoji_config_name, emoji_config in _Emojis.model_fields.items(): if not (match := EMOJI_REGEX.fullmatch(emoji_config.default)): continue emoji_name = match.group(1) emoji_id: str = match.group(2) for emoji in existing_emojis: if emoji["name"] == emoji_name: emoji_id = emoji["id"] break else: log.info("Creating emoji %s", emoji_name) emoji_id = self.client.clone_emoji(new_name=emoji_name, original_emoji_id=emoji_id) data[emoji_config_name] = f"<:{emoji_name}:{emoji_id}>" return data def write_config_env(self, config: dict[str, dict[str, object]]) -> bool: """Write the configuration to the specified env_file.""" if not self.env_file.exists(): self.env_file.touch() with self.env_file.open("r+", encoding="utf-8") as file: before = file.read() file.seek(0) for num, (category, category_values) in enumerate(config.items()): # In order to support commented sections, we write the following file.write(f"# {category.capitalize()}\n") # Format the dictionary into .env style for key, value in category_values.items(): file.write(f"{category}_{key}={value}\n") if num < len(config) - 1: file.write("\n") file.truncate() file.seek(0) after = file.read() return before != after def run(self) -> bool: """Runs the botstrap process.""" # Track if any changes were made and exit with an error code if so. changes: bool = False config: dict[str, dict[str, object | Any]] = {} changes |= self.upgrade_client() self.check_guild_membership() channels = self.get_channels() # Ensure the guild is upgraded to a community if necessary. # This isn't strictly necessary for bot functionality, but # it prevents weird transients since PyDis is a community server. changes |= self.upgrade_guild(channels[ANNOUNCEMENTS_CHANNEL_NAME], channels[RULES_CHANNEL_NAME]) # Though sync_webhooks and sync_emojis DO make api calls that may modify server state, # those changes will be reflected in the config written to the .env file. # Therefore, we don't need to track if any emojis or webhooks are being changed within those settings. config = { "categories": self.get_categories(), "channels": channels, "roles": self.get_roles(), "webhooks": self.sync_webhooks(), "emojis": self.sync_emojis(), } changes |= self.write_config_env(config) return changes if __name__ == "__main__": botstrap = BotStrapper(guild_id=GuildConstants.id, env_file=ENV_FILE, bot_token=BotConstants.token) with botstrap: changes_made = botstrap.run() if changes_made: log.info("Botstrap completed successfully. Updated configuration has been written to %s", ENV_FILE) else: log.info("Botstrap completed successfully. No changes were necessary.") sys.exit(0)