Repository: Cog-Creators/Red-DiscordBot

Issue:
Trivia upload command
Add a command for uploading trivia files.

Upload should be rejected if invalid
Owner only command
Should utilize yaml safeload


Discussion:
I will be looking into this. Should have a PR shortly. 

Resolving patch:
diff --git a/redbot/cogs/trivia/trivia.py b/redbot/cogs/trivia/trivia.py
--- a/redbot/cogs/trivia/trivia.py
+++ b/redbot/cogs/trivia/trivia.py
@@ -1,17 +1,19 @@
 """Module for Trivia cog."""
+import asyncio
 import math
 import pathlib
 from collections import Counter
 from typing import List
-
+import io
 import yaml
 import discord
-
 from redbot.core import commands
 from redbot.core import Config, checks
 from redbot.core.data_manager import cog_data_path
 from redbot.core.i18n import Translator, cog_i18n
 from redbot.core.utils.chat_formatting import box, pagify, bold
+from redbot.core.utils.predicates import MessagePredicate, ReactionPredicate
+from redbot.core.utils.menus import start_adding_reactions
 from redbot.cogs.bank import check_global_setting_admin
 from .log import LOG
 from .session import TriviaSession
@@ -183,6 +185,86 @@ async def triviaset_payout_multiplier(self, ctx: commands.Context, multiplier: f
         else:
             await ctx.send(_("Done. I will no longer reward the winner with a payout."))
 
+    @triviaset.group(name="custom")
+    @commands.is_owner()
+    async def triviaset_custom(self, ctx: commands.Context):
+        """Manage Custom Trivia lists."""
+        pass
+
+    @triviaset_custom.command(name="list")
+    async def custom_trivia_list(self, ctx: commands.Context):
+        """List uploaded custom trivia."""
+        personal_lists = sorted([p.resolve().stem for p in cog_data_path(self).glob("*.yaml")])
+        no_lists_uploaded = _("No custom Trivia lists uploaded.")
+
+        if not personal_lists:
+            if await ctx.embed_requested():
+                await ctx.send(
+                    embed=discord.Embed(
+                        colour=await ctx.embed_colour(), description=no_lists_uploaded
+                    )
+                )
+            else:
+                await ctx.send(no_lists_uploaded)
+            return
+
+        if await ctx.embed_requested():
+            await ctx.send(
+                embed=discord.Embed(
+                    title=_("Uploaded trivia lists"),
+                    colour=await ctx.embed_colour(),
+                    description=", ".join(sorted(personal_lists)),
+                )
+            )
+        else:
+            msg = box(
+                bold(_("Uploaded trivia lists")) + "\n\n" + ", ".join(sorted(personal_lists))
+            )
+            if len(msg) > 1000:
+                await ctx.author.send(msg)
+            else:
+                await ctx.send(msg)
+
+    @commands.is_owner()
+    @triviaset_custom.command(name="upload", aliases=["add"])
+    async def trivia_upload(self, ctx: commands.Context):
+        """Upload a trivia file."""
+        if not ctx.message.attachments:
+            await ctx.send(_("Supply a file with next message or type anything to cancel."))
+            try:
+                message = await ctx.bot.wait_for(
+                    "message", check=MessagePredicate.same_context(ctx), timeout=30
+                )
+            except asyncio.TimeoutError:
+                await ctx.send(_("You took too long to upload a list."))
+                return
+            if not message.attachments:
+                await ctx.send(_("You have cancelled the upload process."))
+                return
+            parsedfile = message.attachments[0]
+        else:
+            parsedfile = ctx.message.attachments[0]
+        try:
+            await self._save_trivia_list(ctx=ctx, attachment=parsedfile)
+        except yaml.error.MarkedYAMLError as exc:
+            await ctx.send(_("Invalid syntax: ") + str(exc))
+        except yaml.error.YAMLError:
+            await ctx.send(
+                _("There was an error parsing the trivia list. See logs for more info.")
+            )
+            LOG.exception("Custom Trivia file %s failed to upload", parsedfile.filename)
+
+    @commands.is_owner()
+    @triviaset_custom.command(name="delete", aliases=["remove"])
+    async def trivia_delete(self, ctx: commands.Context, name: str):
+        """Delete a trivia file."""
+        filepath = cog_data_path(self) / f"{name}.yaml"
+        if filepath.exists():
+            filepath.unlink()
+            await ctx.send(_("Trivia {filename} was deleted.").format(filename=filepath.stem))
+        else:
+            await ctx.send(_("Trivia file was not found."))
+
     @commands.group(invoke_without_command=True)
     @commands.guild_only()
     async def trivia(self, ctx: commands.Context, *categories: str):
@@ -525,6 +607,73 @@ def get_trivia_list(self, category: str) -> dict:
             else:
                 return dict_
 
+    async def _save_trivia_list(
+        self, ctx: commands.Context, attachment: discord.Attachment
+    ) -> None:
+        """Checks and saves a trivia list to data folder.
+
+        Parameters
+        ----------
+        file : discord.Attachment
+            A discord message attachment.
+
+        Returns
+        -------
+        None
+        """
+        filename = attachment.filename.rsplit(".", 1)[0]
+
+        # Check if trivia filename exists in core files or if it is a command
+        if filename in self.trivia.all_commands or any(
+            filename == item.stem for item in get_core_lists()
+        ):
+            await ctx.send(
+                _(
+                    "{filename} is a reserved trivia name and cannot be replaced.\n"
+                    "Choose another name."
+                ).format(filename=filename)
+            )
+            return
+
+        file = cog_data_path(self) / f"{filename}.yaml"
+        if file.exists():
+            overwrite_message = _("{filename} already exists. Do you wish to overwrite?").format(
+                filename=filename
+            )
+
+            can_react = ctx.channel.permissions_for(ctx.me).add_reactions
+            if not can_react:
+                overwrite_message += " (y/n)"
+
+            overwrite_message_object: discord.Message = await ctx.send(overwrite_message)
+            if can_react:
+                # noinspection PyAsyncCall
+                start_adding_reactions(
+                    overwrite_message_object, ReactionPredicate.YES_OR_NO_EMOJIS
+                )
+                pred = ReactionPredicate.yes_or_no(overwrite_message_object, ctx.author)
+                event = "reaction_add"
+            else:
+                pred = MessagePredicate.yes_or_no(ctx=ctx)
+                event = "message"
+            try:
+                await ctx.bot.wait_for(event, check=pred, timeout=30)
+            except asyncio.TimeoutError:
+                await ctx.send(_("You took too long answering."))
+                return
+
+            if pred.result is False:
+                await ctx.send(_("I am not replacing the existing file."))
+                return
+
+        buffer = io.BytesIO(await attachment.read())
+        yaml.safe_load(buffer)
+        buffer.seek(0)
+
+        with file.open("wb") as fp:
+            fp.write(buffer.read())
+        await ctx.send(_("Saved Trivia list as {filename}.").format(filename=filename))
+
     def _get_trivia_session(self, channel: discord.TextChannel) -> TriviaSession:
         return next(
             (session for session in self.trivia_sessions if session.ctx.channel == channel), None


Test patch: