diff --git a/Heimdallr.py b/Heimdallr.py index 144f5c1..912927d 100644 --- a/Heimdallr.py +++ b/Heimdallr.py @@ -1,3 +1,4 @@ +import logging from os import getenv from naff import ( @@ -7,7 +8,7 @@ from naff import ( slash_command, InteractionContext, ) -from naff.api.events.discord import MessageCreate +from naff.api import events from naff.models.discord.embed import ( Embed, EmbedAuthor, @@ -15,6 +16,9 @@ from naff.models.discord.embed import ( EmbedFooter, ) + +from database import GuildSettings as GuildSettingsModel, JoinLeave + from dotenv import load_dotenv bot = Client(intents=Intents.ALL, debug_scope=387153131378835456) @@ -25,10 +29,59 @@ async def on_ready(): print(f"Bot '{bot.user.username}' is ready!") print(f"This bot is owned by {bot.owner}") + for guild in bot.guilds: + guild_q = GuildSettingsModel.get_or_none(GuildSettingsModel.guild_id == guild.id) + if guild_q is None: + guild_q = GuildSettingsModel(guild_id=guild.id) + guild_q.save() + -@listen() -async def on_message_create(event: MessageCreate): - print(f"{event.message.author.username}: {event.message.content}") + +# @listen() +# async def on_message_create(event: MessageCreate): +# print(f"{event.message.author.username}: {event.message.content}") + +@listen(events.MemberAdd) +async def on_member_join(event: events.MemberAdd): + joinleave_q: JoinLeave + joinleave_q, _ = JoinLeave.get_or_create(guild_id=event.guild.id) + + if joinleave_q.message_channel is None: + return + channel = await bot.fetch_channel(joinleave_q.message_channel) + + if not joinleave_q.join_message_enabled: + return + + if joinleave_q.join_message is None or joinleave_q.join_message == "": + await channel.send( + f"{event.member.mention} has joined the server!" + ) + return + + await channel.send(str(joinleave_q.join_message).format(member=event.member, guild=event.guild)) + +@listen(events.MemberRemove) +async def on_member_leave(event: events.MemberRemove): + joinleave_q: JoinLeave + joinleave_q, _ = JoinLeave.get_or_create(guild_id=event.guild.id) + + if joinleave_q.message_channel is None: + return + channel = await bot.fetch_channel(joinleave_q.message_channel) + + if not joinleave_q.leave_message_enabled: + return + + if joinleave_q.leave_message is None or joinleave_q.leave_message == "": + await channel.send( + f"{event.member.mention} has left the server!" + ) + return + + await channel.send(str(joinleave_q.leave_message).format(member=event.member, guild=event.guild)) + + @slash_command(name="ping", description="Ping the bot") @@ -58,15 +111,24 @@ async def bot_info_command(ctx: InteractionContext): fields=[ EmbedField( name="**Extensions**", - value=f"{len(bot.extensions)}", - ) + value=f"{', '.join(bot.ext.keys())}", + ), + EmbedField( + name="**Guilds**", + value="\n".join([g.name for g in sorted(bot.guilds, key=lambda g: g.name)]), + ), ], ), ) if __name__ == "__main__": + logging.basicConfig(level=logging.INFO) + GuildSettingsModel.create_table() + JoinLeave.create_table() + load_dotenv() + bot.load_extension("commands.admin") bot.load_extension("commands.quote") bot.load_extension("commands.infractions") bot.load_extension("commands.self_roles") diff --git a/commands/admin.py b/commands/admin.py new file mode 100644 index 0000000..349caf8 --- /dev/null +++ b/commands/admin.py @@ -0,0 +1,276 @@ +import logging +from typing import Optional +from naff import ( + Client, + Extension, + slash_command, + slash_option, + OptionTypes, + Permissions, + InteractionContext, + Embed, + EmbedField, + GuildChannel, + ChannelTypes, +) + +from database import GuildSettings as GuildSettingsModel, JoinLeave as JoinLeaveModel + + +class Admin(Extension): + def __init__(self, client: Client) -> None: + self.client = client + + @slash_command( + name="adm", + sub_cmd_name="list", + sub_cmd_description="List all settings for this guild", + dm_permission=False, + default_member_permissions=Permissions.MANAGE_GUILD, + ) + async def adm_list(self, ctx: InteractionContext) -> None: + guild_settings: Optional[GuildSettingsModel] + guild_settings, _ = GuildSettingsModel.get_or_create(guild_id=ctx.guild_id) + joinleave_settings: Optional[GuildSettingsModel] + joinleave_settings, _ = JoinLeaveModel.get_or_create(guild_id=ctx.guild_id) + embed = Embed( + title="Settings for {}".format(ctx.guild.name), + fields=[ + EmbedField( + name="**Admin channel**", + value=( + await ctx.guild.fetch_channel(guild_settings.admin_channel) + ).mention + if guild_settings.admin_channel is not None + else "Not set", + ), + EmbedField( + name="**Name filter**", + value=str(guild_settings.use_name_filter) + if guild_settings.use_name_filter is not None + else "Not set (disabled)", + ), + EmbedField( + name="**Gatekeep**", + value=str(guild_settings.use_gatekeep) + if guild_settings.use_gatekeep is not None + else "Not set (disabled)", + ), + EmbedField( + name="**Join/leave messages channel**", + value=( + await ctx.guild.fetch_channel( + joinleave_settings.message_channel + ) + ).mention + if joinleave_settings.message_channel is not None + else "Not set", + ), + EmbedField( + name="**Join messages enabled**", + value=str(joinleave_settings.join_message_enabled) + if joinleave_settings.join_message_enabled is not None + else "Not set (disabled)", + ), + EmbedField( + name="**Leave messages enabled**", + value=str(joinleave_settings.leave_message_enabled) + if joinleave_settings.leave_message_enabled is not None + else "Not set (disabled)", + ), + EmbedField( + name="**Join message**", + value=joinleave_settings.join_message + if joinleave_settings.join_message is not None + else "Not set", + ), + EmbedField( + name="**Leave message**", + value=joinleave_settings.leave_message + if joinleave_settings.leave_message is not None + else "Not set", + ), + ], + ) + + await ctx.send(embed=embed) + + @slash_command( + name="adm", + group_name="set", + group_description="Set settings for this guild", + sub_cmd_name="admin-channel", + sub_cmd_description="Set the admin channel", + dm_permission=False, + default_member_permissions=Permissions.MANAGE_GUILD, + ) + @slash_option( + name="channel", + description="Channel to set as admin channel", + required=True, + opt_type=OptionTypes.CHANNEL, + channel_types=[ChannelTypes.GUILD_TEXT], + ) + async def adm_set_admin_channel( + self, ctx: InteractionContext, channel: GuildChannel + ) -> None: + guild_settings: Optional[GuildSettingsModel] + guild_settings, _ = GuildSettingsModel.get_or_create(guild_id=ctx.guild_id) + guild_settings.admin_channel = channel.id + guild_settings.save() + await ctx.send( + "Admin channel set to {}".format(channel.mention), ephemeral=True + ) + + + @slash_command( + name="adm", + group_name="set", + group_description="Set settings for this guild", + sub_cmd_name="use-name-filter", + sub_cmd_description="Set whether or not to use the name filter", + dm_permission=False, + default_member_permissions=Permissions.MANAGE_GUILD, + ) + @slash_option( + name="enabled", + description="Whether or not to use the name filter", + required=True, + opt_type=OptionTypes.BOOLEAN, + ) + async def adm_set_use_name_filter(self, ctx: InteractionContext, enabled: bool) -> None: + guild_settings: Optional[GuildSettingsModel] + guild_settings, _ = GuildSettingsModel.get_or_create(guild_id=ctx.guild_id) + guild_settings.use_name_filter = enabled + guild_settings.save() + await ctx.send("Name filter set to {}".format(enabled), ephemeral=True) + + + @slash_command( + name="adm", + group_name="set", + group_description="Set settings for this guild", + sub_cmd_name="use-gatekeep", + sub_cmd_description="Set whether or not to use gatekeep", + dm_permission=False, + default_member_permissions=Permissions.MANAGE_GUILD, + ) + @slash_option( + name="enabled", + description="Whether or not to use gatekeep", + required=True, + opt_type=OptionTypes.BOOLEAN, + ) + async def adm_set_use_gatekeep(self, ctx: InteractionContext, enabled: bool) -> None: + guild_settings: Optional[GuildSettingsModel] + guild_settings, _ = GuildSettingsModel.get_or_create(guild_id=ctx.guild_id) + guild_settings.use_gatekeep = enabled + guild_settings.save() + await ctx.send("Gatekeep set to {}".format(enabled), ephemeral=True) + + + @slash_command( + name="adm", + group_name="set", + group_description="Set settings for this guild", + sub_cmd_name="join-leave-channel", + sub_cmd_description="Set the join/leave messages channel", + dm_permission=False, + default_member_permissions=Permissions.MANAGE_GUILD, + ) + @slash_option( + name="channel", + description="Channel to set as join/leave messages channel", + required=True, + opt_type=OptionTypes.CHANNEL, + channel_types=[ChannelTypes.GUILD_TEXT], + ) + async def adm_set_join_leave_channel( + self, ctx: InteractionContext, channel: GuildChannel + ) -> None: + joinleave_settings: Optional[JoinLeaveModel] + joinleave_settings, _ = JoinLeaveModel.get_or_create(guild_id=ctx.guild_id) + joinleave_settings.message_channel = channel.id + joinleave_settings.save() + await ctx.send( + "Join/leave messages channel set to {}".format(channel.mention), ephemeral=True + ) + + + @slash_command( + name="adm", + group_name="set", + group_description="Set settings for this guild", + sub_cmd_name="join-message", + sub_cmd_description="Set the join message and toggle its enabled state", + dm_permission=False, + default_member_permissions=Permissions.MANAGE_GUILD, + ) + @slash_option( + name="enabled", + description="Whether or not to enable the join message", + required=True, + opt_type=OptionTypes.BOOLEAN, + ) + @slash_option( + name="message", + description="The join message", + required=False, + opt_type=OptionTypes.STRING, + ) + async def adm_set_join_message( + self, ctx: InteractionContext, enabled: bool, message: str = None + ) -> None: + joinleave_settings: Optional[JoinLeaveModel] + joinleave_settings, _ = JoinLeaveModel.get_or_create(guild_id=ctx.guild_id) + joinleave_settings.join_message_enabled = enabled + joinleave_settings.join_message = message + joinleave_settings.save() + if enabled: + await ctx.send( + "Join message enabled and set to {}".format(message), ephemeral=True + ) + else: + await ctx.send("Join message disabled", ephemeral=True) + + + @slash_command( + name="adm", + group_name="set", + group_description="Set settings for this guild", + sub_cmd_name="leave-message", + sub_cmd_description="Set the leave message and toggle its enabled state", + dm_permission=False, + default_member_permissions=Permissions.MANAGE_GUILD, + ) + @slash_option( + name="enabled", + description="Whether or not to enable the leave message", + required=True, + opt_type=OptionTypes.BOOLEAN, + ) + @slash_option( + name="message", + description="The leave message", + required=False, + opt_type=OptionTypes.STRING, + ) + async def adm_set_leave_message( + self, ctx: InteractionContext, enabled: bool, message: str = None + ) -> None: + joinleave_settings: Optional[JoinLeaveModel] + joinleave_settings, _ = JoinLeaveModel.get_or_create(guild_id=ctx.guild_id) + joinleave_settings.leave_message_enabled = enabled + joinleave_settings.leave_message = message + joinleave_settings.save() + if enabled: + await ctx.send( + "Leave message enabled and set to {}".format(message), ephemeral=True + ) + else: + await ctx.send("Leave message disabled", ephemeral=True) + + +def setup(client: Client) -> None: + Admin(client) + logging.info("Admin extension loaded") diff --git a/commands/infractions.py b/commands/infractions.py index e15ff18..2bbb1d9 100644 --- a/commands/infractions.py +++ b/commands/infractions.py @@ -121,8 +121,6 @@ class Infractions(Extension): ) embeds.append(embed) - print(infractions.count()) - print(infractions_weight.count()) await ctx.send( content=f"{user.mention} has {infractions_weight[0].total_infractions} infractions, for a total weight of {infractions_weight[0].total_weight:.2f}" if infractions_weight.count() > 0 diff --git a/database.py b/database.py index e71bf88..557ad58 100644 --- a/database.py +++ b/database.py @@ -1,7 +1,10 @@ import datetime +import os +from dotenv import load_dotenv from peewee import ( SqliteDatabase, + PostgresqlDatabase, Model, BigIntegerField, CharField, @@ -14,7 +17,12 @@ from peewee import ( ForeignKeyField, ) -db = SqliteDatabase("heimdallr.db") + +db = None +if (postgres_url := os.getenv("HEIMDALLR_POSTGRES_URL")) is not None: + db = PostgresqlDatabase(postgres_url) +else: + db = SqliteDatabase("heimdallr.db") class Infractions(Model): @@ -37,7 +45,6 @@ class GuildSettings(Model): admin_channel = BigIntegerField(null=True) use_name_filter = BooleanField(default=False) use_gatekeep = BooleanField(default=False) - use_joinleave = BooleanField(default=True) use_logging = BooleanField(default=False) class Meta: