import discord from discord import default_permissions, guild_only from discord.ext import commands from config import debug, con_data, curs_data, con_premium, curs_premium, ctx_to_guid class NoPrivateMessages(commands.CheckFailure): pass def dms_only(): async def predicate(ctx): if ctx.guild is not None: raise NoPrivateMessages('Hey no private messages!') return True return commands.check(predicate) class Setup(discord.Cog): def __init__(self, bot: discord.Bot): super().__init__() self.bot = bot @discord.slash_command(name="setup", description="Setup the bot") @discord.option(name="channel_id", description="The channel id", required=True) @discord.option(name="api_key", description="The api key", required=True) @default_permissions(administrator=True) @guild_only() async def setup( self, ctx: discord.ApplicationContext, channel: discord.TextChannel, api_key: str, ): if channel is None: await ctx.respond("Invalid channel id", ephemeral=True) return try: curs_data.execute("SELECT * FROM data WHERE guild_id = ?", (ctx.guild.id,)) data = curs_data.fetchone() if data[3] == None: data = None except: data = None if data != None: curs_data.execute( "UPDATE data SET channel_id = ?, api_key = ? WHERE guild_id = ?", (channel.id, api_key, ctx.guild.id), ) # c.execute("UPDATE data SET is_active = ?, max_tokens = ?, temperature = ?, frequency_penalty = ?, presence_penalty = ?, prompt_size = ? WHERE guild_id = ?", (False, 64, 0.9, 0.0, 0.0, 5, ctx.guild.id)) con_data.commit() await ctx.respond( "The channel id and the api key have been updated", ephemeral=True ) else: curs_data.execute( "INSERT INTO data VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", ( ctx.guild.id, channel.id, api_key, False, 64, 0.9, 0.0, 0.0, 0, 5, "", False, "", False, ), ) con_data.commit() await ctx.respond( "The channel id and the api key have been added", ephemeral=True ) @discord.slash_command(name="setup_dms", description="Setup the bot in dms") @discord.option(name="api_key", description="The api key", required=True) @default_permissions(administrator=True) @dms_only() async def setup_dms( self, ctx: discord.ApplicationContext, api_key: str, ): channel = ctx.channel if channel is None: await ctx.respond("Invalid channel id", ephemeral=True) return try: curs_data.execute("SELECT * FROM data WHERE guild_id = ?", (ctx.user.id,)) data = curs_data.fetchone() if data[3] == None: data = None except: data = None if data != None: curs_data.execute( "UPDATE data SET channel_id = ?, api_key = ? WHERE guild_id = ?", (channel.id, api_key, ctx.user.id), ) con_data.commit() await ctx.respond( "The channel id and the api key have been updated", ephemeral=True ) else: curs_data.execute( "INSERT INTO data VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", ( ctx.user.id, channel.id, api_key, False, 64, 0.9, 0.0, 0.0, 0, 5, "", False, "", False, ), ) con_data.commit() await ctx.respond( "The api key has been added", ephemeral=True ) @discord.slash_command( name="delete", description="Delete the information about this server" ) @default_permissions(administrator=True) async def delete(self, ctx: discord.ApplicationContext): # check if the guild is in the database curs_data.execute("SELECT * FROM data WHERE guild_id = ?", (ctx_to_guid(ctx),)) if curs_data.fetchone() is None: await ctx.respond("This server is not setup", ephemeral=True) return # delete the guild from the database, except the guild id and the uses_count_today curs_data.execute( "UPDATE data SET api_key = ?, channel_id = ?, is_active = ?, max_tokens = ?, temperature = ?, frequency_penalty = ?, presence_penalty = ?, prompt_size = ? WHERE guild_id = ?", (None, None, False, 50, 0.9, 0.0, 0.0, 0, ctx_to_guid(ctx)), ) con_data.commit() await ctx.respond("Deleted", ephemeral=True) # create a command called "enable" that only admins can use @discord.slash_command(name="enable", description="Enable the bot") @default_permissions(administrator=True) async def enable(self, ctx: discord.ApplicationContext): curs_data.execute("SELECT * FROM data WHERE guild_id = ?", (ctx_to_guid(ctx),)) if curs_data.fetchone() is None: await ctx.respond("This server is not setup", ephemeral=True) return # enable the guild curs_data.execute( "UPDATE data SET is_active = ? WHERE guild_id = ?", (True, ctx_to_guid(ctx)) ) con_data.commit() await ctx.respond("Enabled", ephemeral=True) # create a command called "disable" that only admins can use @discord.slash_command(name="disable", description="Disable the bot") @default_permissions(administrator=True) async def disable(self, ctx: discord.ApplicationContext): # check if the guild is in the database curs_data.execute("SELECT * FROM data WHERE guild_id = ?", (ctx_to_guid(ctx),)) if curs_data.fetchone() is None: await ctx.respond("This server is not setup", ephemeral=True) return # disable the guild curs_data.execute( "UPDATE data SET is_active = ? WHERE guild_id = ?", (False, ctx_to_guid(ctx)) ) con_data.commit() await ctx.respond("Disabled", ephemeral=True) # create a command calles "add channel" that can only be used in premium servers @discord.slash_command( name="add_channel", description="Add a channel to the list of channels. Premium only.", ) @discord.option( name="channel", description="The channel to add", type=discord.TextChannel, required=False, ) @default_permissions(administrator=True) @guild_only() async def add_channel( self, ctx: discord.ApplicationContext, channel: discord.TextChannel = None ): # check if the guild is in the database curs_data.execute("SELECT * FROM data WHERE guild_id = ?", (ctx.guild.id,)) if curs_data.fetchone() is None: await ctx.respond("This server is not setup", ephemeral=True) return con_premium.execute( "SELECT premium FROM data WHERE guild_id = ?", (ctx.guild.id,) ) premium = con_premium.fetchone()[0] if not premium: await ctx.respond("This server is not premium", ephemeral=True) return if channel is None: channel = ctx.channel # check if the channel is already in the list curs_data.execute( "SELECT channel_id FROM data WHERE guild_id = ?", (ctx.guild.id,) ) if str(channel.id) == curs_data.fetchone()[0]: await ctx.respond( "This channel is already set as the main channel", ephemeral=True ) return con_premium.execute( "SELECT * FROM channels WHERE guild_id = ?", (ctx.guild.id,) ) guild_channels = con_premium.fetchone() if guild_channels is None: # if the channel is not in the list, add it con_premium.execute( "INSERT INTO channels VALUES (?, ?, ?, ?, ?, ?)", (ctx.guild.id, channel.id, None, None, None, None), ) con_premium.commit() await ctx.respond(f"Added channel **{channel.name}**", ephemeral=True) return channels = guild_channels[1:] if str(channel.id) in channels: await ctx.respond("This channel is already added", ephemeral=True) return for i in range(5): if channels[i] == None: con_premium.execute( f"UPDATE channels SET channel{i} = ? WHERE guild_id = ?", (channel.id, ctx.guild.id), ) con_premium.commit() await ctx.respond(f"Added channel **{channel.name}**", ephemeral=True) return await ctx.respond("You can only add 5 channels", ephemeral=True) # create a command called "remove channel" that can only be used in premium servers @discord.slash_command( name="remove_channel", description="Remove a channel from the list of channels. Premium only.", ) @discord.option( name="channel", description="The channel to remove", type=discord.TextChannel, required=False, ) @default_permissions(administrator=True) @guild_only() async def remove_channel( self, ctx: discord.ApplicationContext, channel: discord.TextChannel = None ): # check if the guild is in the database curs_data.execute("SELECT * FROM data WHERE guild_id = ?", (ctx.guild.id,)) if curs_data.fetchone() is None: await ctx.respond("This server is not setup", ephemeral=True) return # check if the guild is premium try: con_premium.execute( "SELECT premium FROM data WHERE guild_id = ?", (ctx.guild.id,) ) premium = con_premium.fetchone()[0] except: premium = 0 if not premium: await ctx.respond("This server is not premium", ephemeral=True) return if channel is None: channel = ctx.channel # check if the channel is in the list con_premium.execute( "SELECT * FROM channels WHERE guild_id = ?", (ctx.guild.id,) ) guild_channels = con_premium.fetchone() curs_data.execute( "SELECT channel_id FROM data WHERE guild_id = ?", (ctx.guild.id,) ) if str(channel.id) == curs_data.fetchone()[0]: await ctx.respond( "This channel is set as the main channel and therefore cannot be removed. Type /setup to change the main channel.", ephemeral=True, ) return if guild_channels is None: await ctx.respond( "This channel was not added. Nothing changed", ephemeral=True ) return channels = guild_channels[1:] if str(channel.id) not in channels: await ctx.respond( "This channel was not added. Nothing changed", ephemeral=True ) return # remove the channel from the list for i in range(5): if channels[i] == str(channel.id): con_premium.execute( f"UPDATE channels SET channel{i} = ? WHERE guild_id = ?", (None, ctx.guild.id), ) con_premium.commit() await ctx.respond(f"Removed channel **{channel.name}**", ephemeral=True) return