First commit.

ThatRedox · ThatRedox · commit 72ac01b72d77 · 2022-03-23T19:16:39.000-07:00
diff --git a/config.ini b/config.ini
@@ -0,0 +1,8 @@
+[IMAGE_ONLY]
+# Test channel
+956319519725469706 = You need to provide a direct link to your render or upload it as an attachment!
+# Chunky #renders channel
+549680988989423631 = You need to provide a direct link to your render or upload it as an attachment!
+
+[GITHUB]
+repository = chunky-dev/chunky
diff --git a/main.py b/main.py
@@ -0,0 +1,186 @@
+import argparse
+import configparser
+import logging
+import re
+from typing import *
+
+import discord
+import discord_slash
+import github
+import github.Repository
+
+import utils
+
+REMOVE_EMOJI = discord.PartialEmoji(name="❌")
+
+
+LOG_LEVEL_MAP = {
+    "ALL": logging.NOTSET,
+    "DEBUG": logging.DEBUG,
+    "INFO": logging.INFO,
+    "WARN": logging.WARNING,
+    "ERROR": logging.ERROR,
+    "FATAL": logging.FATAL
+}
+
+
+class Bot(discord.Client):
+    GH_REGEX = re.compile(r"#\d+")
+
+    def __init__(self, repo, image_only: List[Tuple[int, str]], *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._repo = repo
+        self._image_only = image_only
+        self._image_only_channels = set(i[0] for i in image_only)
+        self._logger = logging.getLogger("bot")
+
+    async def on_message(self, message: discord.Message):
+        # Check if the message is from ourselves
+        if message.author.id == self.user.id:
+            return
+
+        # Check if we are in the renderers channel
+        if message.channel.id in self._image_only_channels:
+            if not utils.is_image(message):
+                self._logger.info(f"Removing message {message.id} in {message.channel.id} for not having an image.")
+                warn = await message.reply(
+                    content="You need to provide a direct link to your render or upload it as an attachment!",
+                    mention_author=True
+                )
+                await message.delete()
+                await warn.delete(delay=10)
+            return
+
+        # Look for GitHub issues / pull requests
+        numbers = self.GH_REGEX.findall(message.content)
+        numbers = [int(number[1:]) for number in numbers]
+
+        # Create the embed
+        embed = None
+        if len(numbers) == 1:
+            self._logger.info(f"Message {message.id} with one GitHub number.")
+            embed = utils.generate_gh_embed(numbers[0], self._repo)
+        elif len(numbers) > 1:
+            self._logger.info(f"Message {message.id} with {len(numbers)} GitHub numbers.")
+            embed = discord.Embed(title="Issues / pull requests")
+            for number in numbers:
+                utils.generate_gh_embed_snippet(embed, number, self._repo)
+
+        # Send the message
+        if embed is not None:
+            embed.set_footer(text=f"React with {REMOVE_EMOJI} to remove.\n{message.author.id}")
+            m = await message.reply(
+                embed=embed,
+                mention_author=False
+            )
+            await m.add_reaction(REMOVE_EMOJI)
+
+    async def on_raw_reaction_add(self, payload: discord.RawReactionActionEvent):
+        if payload.user_id == self.user.id:
+            return  # Event from us
+        if payload.emoji != REMOVE_EMOJI:
+            return  # Incorrect emoji
+        if payload.event_type != "REACTION_ADD":
+            return  # Emoji was not added
+
+        # Get the channel
+        channel = self.get_channel(payload.channel_id)
+        message: discord.Message = await channel.fetch_message(payload.message_id)
+
+        if message.author.id != self.user.id:
+            return  # Original author was not us
+        if len(message.embeds) != 1:
+            return  # Original message had an incorrect number of embeds
+
+        # Try and get the user id from the embed
+        embed = message.embeds[0]
+        if not isinstance(embed.footer.text, str):
+            return  # Invalid embed
+        user = embed.footer.text.split("\n")[-1]
+        try:
+            user = int(user)
+        except ValueError:
+            return  # Invalid user id
+        if user != payload.user_id:
+            return  # User does not have permission to remove this
+
+        self._logger.info(f"React-deleting our message {message.id}")
+        await message.delete()
+
+
+class Slash(discord_slash.SlashCommand):
+    def __init__(self, repo, image_only: List[Tuple[int, str]], *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._repo = repo
+        self._logger = logging.getLogger("bot-slash")
+        self._image_only_channels = set(i[0] for i in image_only)
+
+        self.add_slash_command(self.gh, name="gh", description="Get a Github pull / issue from its number.")
+
+    async def gh(self, ctx, number: int):
+        if ctx.channel_id in self._image_only_channels:
+            self._logger.info(f"Attempted slash command in protected channel {ctx.channel_id}.")
+            await ctx.send(content=f"Cannot send text messages in this channel.", hidden=True)
+            return
+
+        embed = utils.generate_gh_embed(number, self._repo)
+        if embed is not None:
+            self._logger.info(f"Slash command with valid GitHub number #{number}.")
+            embed.set_footer(text=f"React with {REMOVE_EMOJI} to remove.\n{ctx.author_id}")
+            m = await ctx.send(embed=embed, hidden=False)
+            await m.add_reaction(REMOVE_EMOJI)
+        else:
+            self._logger.info(f"Slash command with invalid GitHub number #{number}.")
+            await ctx.send(content=f"Invalid pull / issue number: #{number}", hidden=True)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Chunky Discord Bot")
+    parser.add_argument("discord", help="Discord API key.")
+    parser.add_argument("--github", help="Github API key.", default=None)
+    parser.add_argument("--log-level", help="Log level (default INFO).", default="INFO")
+    parser.add_argument("--config", help="Path to the config file.", default="config.ini")
+    parser.add_argument("--debug-guild", help="Debug guild id.", default=None)
+    args = parser.parse_args()
+
+    # Setup logging
+    if args.log_level not in LOG_LEVEL_MAP.keys():
+        print("Log level must be one of: ALL, DEBUG, INFO, WARN, ERROR, FATAL")
+        return
+    logging.basicConfig(level=LOG_LEVEL_MAP.get(args.log_level))
+
+    # Load config
+    config = configparser.ConfigParser()
+    config.read(args.config)
+
+    # Setup GitHub
+    if "GITHUB" not in config:
+        print("Config must have [GITHUB] section.")
+        return
+    if "repository" not in config["GITHUB"]:
+        print("Config must have \"repository\" under [GITHUB] section.")
+        return
+    gh = github.Github(login_or_token=args.github)
+    repo = gh.get_repo(config["GITHUB"]["repository"])
+
+    # Image only channels
+    image_only: List[Tuple[int, str]] = []
+    if "IMAGE_ONLY" in config:
+        for key, value in config["IMAGE_ONLY"].items():
+            try:
+                image_only.append((int(key), value,))
+            except ValueError:
+                logging.getLogger("bot").error(f"Invalid [IMAGE_ONLY] channel {key}.")
+    else:
+        logging.getLogger("bot").warning("Config does not contain an [IMAGE_ONLY] section. Bot will not filter any channels.")
+
+    bot = Bot(repo, image_only)
+    slash = Slash(repo, image_only, client=bot, debug_guild=args.debug_guild, sync_commands=True)
+
+    # OAUTH2 must have `bot` and `applications.commands` scopes
+    # Bot permissions: 274877982784
+    bot.run(args.discord)
+
+
+if __name__ == '__main__':
+    main()
diff --git a/requirements.txt b/requirements.txt
@@ -0,0 +1,3 @@
+discord.py~=1.7.3
+discord-py-slash-command~=3.0.3
+PyGithub~=1.55
diff --git a/utils.py b/utils.py
@@ -0,0 +1,100 @@
+from typing import *
+import logging
+import re
+import urllib.parse
+
+import discord
+import github
+import github.Repository
+
+IMAGE_SUFFIXES = [
+    "jpg", "jpeg", "png", "tif", "tiff", "webp", "gif", "mp4"
+]
+
+URL_REGEX = re.compile(r"http\S*")
+
+
+def _match_fname(filename: str) -> bool:
+    fname = filename.lower()
+    for suffix in IMAGE_SUFFIXES:
+        if fname.endswith(suffix):
+            return True
+    return False
+
+
+def is_image(message: discord.Message) -> bool:
+    # Image(s) were uploaded
+    if len(message.attachments) > 0:
+        for attachment in message.attachments:
+            assert isinstance(attachment.filename, str)
+            if _match_fname(attachment.filename):
+                return True
+
+    # Check for an image URL
+    urls = URL_REGEX.findall(message.content)
+    for url in urls:
+        try:
+            url = urllib.parse.urlparse(url)
+            if _match_fname(url.path):
+                return True
+        except ValueError:
+            pass
+
+    return False
+
+
+def generate_gh_embed(number: int, repo: github.Repository.Repository) -> Optional[discord.Embed]:
+    try:
+        issue = repo.get_issue(number)
+        embed = discord.Embed(
+            title=issue.html_url,
+            url=issue.html_url,
+            type="rich",
+            description=issue.title,
+        )
+        embed.add_field(
+            name="By",
+            value=issue.user.login,
+            inline=True
+        )
+        embed.add_field(
+            name="Status",
+            value=issue.state,
+            inline=True
+        )
+        embed.add_field(
+            name="Description",
+            value=issue.body,
+            inline=False
+        )
+        return embed
+    except github.GithubException as e:
+        logging.getLogger("github").warning(f"Failed to fetch object number {number}. {e}")
+        return None
+
+
+def generate_gh_embed_snippet(embed: discord.Embed, number: id, repo: github.Repository.Repository):
+    try:
+        issue = repo.get_issue(number)
+        embed.add_field(
+            name="Link",
+            value=issue.html_url,
+            inline=False
+        )
+        embed.add_field(
+            name="Title",
+            value=issue.title,
+            inline=True
+        )
+        embed.add_field(
+            name="By",
+            value=issue.user.login,
+            inline=True
+        )
+        embed.add_field(
+            name="Status",
+            value=issue.state,
+            inline=True
+        )
+    except github.GithubException as e:
+        logging.getLogger("github").warning(f"Failed to fetch object number {number}. {e}")

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	`+discord.py~=1.7.3`
	`2`	`+discord-py-slash-command~=3.0.3`
	`3`	`+PyGithub~=1.55`