Heimdallr/heimdallr/Heimdallr.py

import logging
from os import getenv

from naff import (
    Client,
    Intents,
    listen,
    slash_command,
    InteractionContext,
    AuditLogEventType,
    AuditLogEntry,
    RoleSelectMenu,
)
from naff.api import events
from naff.models.discord.embed import (
    Embed,
    EmbedAuthor,
    EmbedField,
    EmbedFooter,
)

from dotenv import load_dotenv
from database import GuildSettings as GuildSettingsModel, JoinLeave

class HeimdallrClient(Client):
    @listen()
    async def on_ready(self):
        print("------------------------------------")
        print(f"Bot '{bot.user.username}' is ready!")
        print(f"This bot is owned by {bot.owner}")
        print("------------------------------------")

        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(events.MemberAdd)
    async def on_member_join(self, 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(self, event: events.MemberRemove):
        joinleave_q: JoinLeave
        joinleave_q, _ = JoinLeave.get_or_create(guild_id=event.guild.id)

        guildsettings_q: GuildSettingsModel
        guildsettings_q, _ = GuildSettingsModel.get_or_create(guild_id=event.guild.id)

        entry: AuditLogEntry
        async for entry in event.guild.audit_log_history(
            action_type=AuditLogEventType.MEMBER_KICK,
            limit=10
        ):
            if (entry.target_id != event.member.id or
                guildsettings_q.admin_channel is None):
                continue

            channel = await bot.fetch_channel(guildsettings_q.admin_channel)
            await channel.send(f"{event.member.mention} was kicked with reason: {entry.reason}")
            break
                


        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")
    async def ping_command(self, ctx: InteractionContext):
        ctx.ephemeral = True
        await ctx.send("Pong!",
            components=RoleSelectMenu(placeholder="HONK")
        )


    @slash_command(name="bot-info", description="Get info about the bot")
    async def bot_info_command(self, ctx: InteractionContext):
        await ctx.send(
            ephemeral=True,
            embed=Embed(
                title=f"{bot.user.username}",
                description=f"This bot is owned by {bot.owner}",
                color=0x00FF00,
                timestamp=bot.user.created_at,
                url=bot.user.avatar.as_url(),
                author=EmbedAuthor(
                    name=bot.user.username,
                    icon_url=bot.user.avatar.as_url(),
                ),
                footer=EmbedFooter(
                    text=f"Created at {bot.user.created_at}",
                    icon_url=bot.user.avatar.as_url(),
                ),
                fields=[
                    EmbedField(
                        name="**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)]),
                    ),
                ],
            ),
        )

bot = HeimdallrClient(intents=Intents.ALL, debug_scope=387153131378835456, 
    sync_interactions=True,
    fetch_members=True,
    
)



def set_loglevel(level: str):
    loglevel = logging.WARNING

    match str(level).lower().strip():
        case "d" | "debug":
            loglevel = logging.DEBUG
        case "i" | "info" | "information":
            loglevel = logging.INFO

        case "w" | "warn" | "warning":
            loglevel = logging.WARNING

        case "e" | "error":
            loglevel = logging.ERROR

        case "c" | "critical":
            loglevel = logging.CRITICAL

        case _:
            loglevel = logging.WARNING

    logging.basicConfig(level=loglevel)


def main():
    load_dotenv() # Load environment variables from .env file
    set_loglevel(getenv("HEIMDALLR_LOGLEVEL"))

    # Create basic tables
    GuildSettingsModel.create_table()
    JoinLeave.create_table()

    # Load extensions
    bot.load_extension("heimdallr.commands.admin")
    bot.load_extension("heimdallr.commands.gatekeep")
    bot.load_extension("heimdallr.commands.quote")
    bot.load_extension("heimdallr.commands.infractions")
    bot.load_extension("heimdallr.commands.self_roles")
    bot.load_extension("heimdallr.commands.polls")
    bot.load_extension("heimdallr.commands.bot_messages")
    bot.load_extension("heimdallr.commands.modmail")
    bot.start(getenv("DISCORD_TOKEN"))

if __name__ == "__main__":
    main()