jarvis-bot/jarvis/cogs/starboard.py

import aiohttp
import jarvis
from datetime import datetime
import discord
from discord import TextChannel, Message
from discord.ext import commands
from discord.utils import find
from discord_slash import cog_ext, SlashContext
from discord_slash.utils.manage_commands import create_option
from jarvis.config import get_config
from jarvis.utils import build_embed
from jarvis.utils.db import DBManager
from jarvis.utils.field import Field
from jarvis.utils.permissions import admin_or_permissions

supported_images = [
    "image/png",
    "image/gif",
    "image/jpeg",
    "image/webp",
    "image/svg",
]


class StarboardCog(commands.Cog):
    def __init__(self, bot):
        self.bot = bot
        self.db = DBManager(get_config().mongo).mongo

    @cog_ext.cog_subcommand(
        base="starboard",
        name="list",
        description="Lists all Starboards",
        guild_ids=[418094694325813248, 578757004059738142],
    )
    @commands.has_permissions(administrator=True)
    async def _list(self, ctx):
        starboards = [
            x for x in self.db.jarvis.starboard.find({"guild": ctx.guild.id})
        ]
        if starboards != []:
            message = "Available Starboards:\n"
            for s in starboards:
                message += f"<#{s['target']}>\n"
            await ctx.send(message)
        else:
            await ctx.send("No Starboards available.")

    @cog_ext.cog_subcommand(
        base="starboard",
        name="create",
        description="Create a starboard",
        guild_ids=[418094694325813248, 578757004059738142],
        options=[
            create_option(
                name="target",
                description="Target channel",
                option_type=7,
                required=True,
            ),
        ],
    )
    @commands.has_permissions(administrator=True)
    async def _create(self, ctx, target: TextChannel):
        if target not in ctx.guild.channels:
            await ctx.send(
                "Target channel not in guild. Choose an existing channel."
            )
            return
        exists = self.db.jarvis.starboard.find_one(
            {"target": target.id, "guild": ctx.guild.id}
        )
        if exists:
            await ctx.send(f"Starboard already exists at {target.mention}.")
            return

        self.db.jarvis.starboard.insert_one(
            {
                "guild": ctx.guild.id,
                "target": target.id,
                "admin": ctx.author.id,
                "time": datetime.now(),
            }
        )
        await ctx.send(f"Starboard created. Check it out at {target.mention}.")

    @cog_ext.cog_subcommand(
        base="starboard",
        name="delete",
        description="Delete a starboard",
        guild_ids=[418094694325813248, 578757004059738142],
        options=[
            create_option(
                name="channel",
                description="Starboard channel",
                option_type=7,
                required=True,
            ),
        ],
    )
    @commands.has_permissions(administrator=True)
    async def _delete(self, ctx, target: TextChannel):
        deleted = self.db.jarvis.starboard.delete_one(
            {
                "target": target.id,
                "guild": ctx.guild.id,
            }
        )
        if deleted:
            self.db.jarvis.stars.delete_many({"starboard": target.id})
            await ctx.send(f"Starboard deleted from {target.mention}.")
        else:
            await ctx.send(f"Starboard not found in {target.mention}.")

    @cog_ext.cog_subcommand(
        base="star",
        name="add",
        description="Star a message",
        guild_ids=[418094694325813248, 578757004059738142],
        options=[
            create_option(
                name="message",
                description="Message to star",
                option_type=3,
                required=True,
            ),
            create_option(
                name="starboard",
                description="Starboard to send message to",
                option_type=7,
                required=True,
            ),
            create_option(
                name="channel",
                description="Channel that has the message, "
                + "required if different than command message",
                option_type=7,
                required=False,
            ),
        ],
    )
    @commands.has_permissions(administrator=True)
    async def _star_add(
        self,
        ctx: SlashContext,
        message: str,
        starboard: TextChannel,
        channel: TextChannel = None,
    ):
        if not channel:
            channel = ctx.channel
        exists = self.db.jarvis.starboard.find_one(
            {"target": starboard.id, "guild": ctx.guild.id}
        )
        if not exists:
            await ctx.send(
                f"Starboard does not exist in {starboard.mention}. "
                + "Please create it first"
            )
            return

        message = await channel.fetch_message(int(message))

        exists = self.db.jarvis.stars.find_one(
            {
                "message": message.id,
                "channel": message.channel.id,
                "guild": message.guild.id,
                "starboard": starboard.id,
            }
        )

        if exists:
            await ctx.send(
                f"Message already sent to Starboard {starboard.mention}"
            )
            return

        content = message.content

        attachments = message.attachments
        image_url = None
        if attachments:
            for attachment in attachments:
                if attachment.content_type in supported_images:
                    image_url = attachment.url
                    break
        if not content and image_url:
            content = "\u200b"

        embed = build_embed(
            title="Click Here to view context",
            description=content,
            fields=[],
            url=message.jump_url,
            timestamp=message.created_at,
        )
        embed.set_author(
            name=message.author.name,
            url=message.jump_url,
            icon_url=message.author.avatar_url,
        )
        embed.set_footer(
            text=message.guild.name + " | " + message.channel.name
        )
        if image_url:
            embed.set_image(url=image_url)

        star = await starboard.send(embed=embed)

        self.db.jarvis.stars.insert_one(
            {
                "message": message.id,
                "channel": message.channel.id,
                "guild": message.guild.id,
                "starboard": starboard.id,
                "admin": ctx.author.id,
                "time": datetime.now(),
                "star": star.id,
            }
        )

        await ctx.send(
            "Message saved to Starboard.\n" + f"See it in {starboard.mention}"
        )


def setup(bot):
    bot.add_cog(StarboardCog(bot))