Added admin settings, join/leave messages.

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(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))
+
     
-@listen()
-async def on_message_create(event: MessageCreate):
-    print(f"{event.message.author.username}: {event.message.content}")
 
 
 @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")

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")

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

database.py

@@ -1,7 +1,10 @@
 import datetime
+import os
 
+from dotenv import load_dotenv
 from peewee import (
     SqliteDatabase,
+    PostgresqlDatabase,
     Model,
     BigIntegerField,
     CharField,
@@ -14,6 +17,11 @@ from peewee import (
     ForeignKeyField,
 )
 
+
+db = None
+if (postgres_url := os.getenv("HEIMDALLR_POSTGRES_URL")) is not None:
+    db = PostgresqlDatabase(postgres_url)
+else:
     db = SqliteDatabase("heimdallr.db")
 
 
@@ -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: