feat(bot): Everything.
This commit is contained in:
Nyeki.py
+370
-3
@@ -1,6 +1,373 @@
|
||||
from discord.ext import bridge
|
||||
from typing import Literal
|
||||
|
||||
from datetime import datetime, timezone
|
||||
|
||||
import discord
|
||||
from discord import abc
|
||||
from discord.ext import commands
|
||||
from discord.ext.commands import errors
|
||||
|
||||
from tortoise import expressions, transactions
|
||||
|
||||
import redis.asyncio as aredis
|
||||
|
||||
from payphone.db import setup, models
|
||||
from payphone.core import config, constants, exceptions, formatting, help_command
|
||||
|
||||
|
||||
# TODO: Make it an AutoShardedBot on production.
|
||||
class Payphone(bridge.Bot):
|
||||
pass
|
||||
class Payphone(commands.AutoShardedBot):
|
||||
redis: aredis.Redis
|
||||
|
||||
def __init__(self) -> None:
|
||||
intents = discord.Intents.default()
|
||||
intents.message_content = True
|
||||
|
||||
super().__init__(
|
||||
intents=intents,
|
||||
command_prefix=self.get_prefix,
|
||||
help_command=help_command.HelpCommand(),
|
||||
)
|
||||
|
||||
async def on_ready(self) -> None:
|
||||
"""Initializes the bot.
|
||||
|
||||
Initializes the bot's DB and connects to Redis, and then prints a ready message to the
|
||||
console.
|
||||
"""
|
||||
|
||||
# Import from inside the function to prevent circular import errors.
|
||||
from payphone.cogs.mods import views as mod_views
|
||||
|
||||
await setup.init()
|
||||
|
||||
self.redis = aredis.Redis(host="localhost" if config.IS_DEBUG else "redis")
|
||||
await self.redis.ping()
|
||||
|
||||
await self.redis.delete(
|
||||
constants.CHANNELS_ONCALL_SET, constants.CHANNELS_SEARCH_QUEUE
|
||||
)
|
||||
|
||||
self.add_view(mod_views.BanNotificationView(self, models.Ban()))
|
||||
|
||||
print(f"Ready! Logged in as {self.user.display_name}.")
|
||||
|
||||
async def on_shard_ready(self, shard_id: int) -> None:
|
||||
"""Sets the activity for the shard.
|
||||
|
||||
Args:
|
||||
shard_id (int): The ID of the initialized shard.
|
||||
"""
|
||||
|
||||
await self.change_presence(
|
||||
activity=discord.CustomActivity(
|
||||
name=f"Listening to p.help | Shard {shard_id + 1}"
|
||||
),
|
||||
status=discord.Status.online,
|
||||
shard_id=shard_id,
|
||||
)
|
||||
|
||||
print(f"Shard is ready! Shard ID: {shard_id}")
|
||||
|
||||
async def on_shard_resumed(self, shard_id: int) -> None:
|
||||
"""Sets the activity for the shard.
|
||||
|
||||
Args:
|
||||
shard_id (int): The ID of the resumed shard.
|
||||
"""
|
||||
|
||||
await self.change_presence(
|
||||
activity=discord.CustomActivity(
|
||||
name=f"Listening to p.help | Shard {shard_id + 1}"
|
||||
),
|
||||
status=discord.Status.online,
|
||||
shard_id=shard_id,
|
||||
)
|
||||
|
||||
print(f"Shard was resumed! Shard ID: {shard_id}")
|
||||
|
||||
async def get_or_fetch_channel(
|
||||
self, channel_id: int
|
||||
) -> abc.GuildChannel | abc.PrivateChannel | None:
|
||||
"""Get's a channel from the bot's cache, or fetches it from Discord if it's not available.
|
||||
|
||||
Args:
|
||||
channel_id (int): The channel's ID.
|
||||
|
||||
Returns:
|
||||
(abc.GuildChannel | abc.PrivateChannel | None): The channel, or None if not found.
|
||||
"""
|
||||
|
||||
channel = self.get_channel(channel_id)
|
||||
|
||||
if channel is None:
|
||||
return await self.fetch_channel(channel_id)
|
||||
|
||||
return channel
|
||||
|
||||
async def is_banned(
|
||||
self, user: discord.User | discord.Member | models.User | int
|
||||
) -> bool:
|
||||
"""Checks whether a user is banned.
|
||||
|
||||
Checks whether a user is banned, and caches the result in Redis. The cache expires in 5
|
||||
minutes.
|
||||
|
||||
Args:
|
||||
user (discord.User | discord.Member | models.User | int): The user to check or their
|
||||
ID.
|
||||
|
||||
Returns:
|
||||
bool: Whether the user is banned or not.
|
||||
"""
|
||||
|
||||
user_id = user if isinstance(user, int) else user.id
|
||||
key = constants.BAN_TEMPLATE % user_id
|
||||
|
||||
is_banned = await self.redis.get(key)
|
||||
|
||||
if is_banned is not None:
|
||||
return bool(int(is_banned))
|
||||
|
||||
is_banned = await models.Ban.filter(
|
||||
expressions.Q(user_id=user_id)
|
||||
& expressions.Q(
|
||||
expires_at=None,
|
||||
expires_at__lt=datetime.now(timezone.utc),
|
||||
join_type="OR",
|
||||
)
|
||||
).exists()
|
||||
|
||||
await self.redis.set(key, int(is_banned), ex=300)
|
||||
|
||||
return is_banned
|
||||
|
||||
@transactions.atomic()
|
||||
async def ban(
|
||||
self,
|
||||
user: discord.User | discord.Member | models.User | int,
|
||||
until: datetime | None,
|
||||
reason: str,
|
||||
moderator: discord.User | discord.Member | models.User | int | None = None,
|
||||
) -> models.Ban:
|
||||
"""Bans a user from Payphone.
|
||||
|
||||
Args:
|
||||
user (discord.User | discord.Member | models.User | int): The user to ban or their ID.
|
||||
until (datetime | None): The expiry date, or None if it does not expire.
|
||||
reason (str): The reason of the ban.
|
||||
moderator (discord.User | discord.Member | models.User | int | None, optional): The
|
||||
moderator who performed the action.
|
||||
|
||||
Returns:
|
||||
models.Ban: The newly created ban.
|
||||
"""
|
||||
|
||||
user_id = user if isinstance(user, int) else user.id
|
||||
key = constants.BAN_TEMPLATE % user_id
|
||||
|
||||
mod_id = (
|
||||
moderator.id
|
||||
if isinstance(moderator, (discord.User, discord.Member, models.User))
|
||||
else moderator
|
||||
)
|
||||
|
||||
row = await models.Ban.create(
|
||||
user_id=user_id,
|
||||
reason=reason,
|
||||
expires_at=until,
|
||||
banned_by_id=mod_id,
|
||||
)
|
||||
|
||||
await self.redis.set(key, "1", ex=300)
|
||||
|
||||
return row
|
||||
|
||||
@transactions.atomic()
|
||||
async def unban(
|
||||
self, user: discord.User | discord.Member | models.User | int
|
||||
) -> None:
|
||||
"""Unbans a user from Payphone.
|
||||
|
||||
Args:
|
||||
user (discord.User | discord.Member | models.User | int): The user to unban or their
|
||||
ID.
|
||||
"""
|
||||
|
||||
user_id = user if isinstance(user, int) else user.id
|
||||
key = constants.BAN_TEMPLATE % user_id
|
||||
|
||||
await models.Ban.filter(
|
||||
expressions.Q(
|
||||
expires_at=None,
|
||||
expires_at__gt=datetime.now(timezone.utc),
|
||||
join_type="OR",
|
||||
),
|
||||
user_id=user_id,
|
||||
).update(expires_at=datetime.now(timezone.utc))
|
||||
|
||||
await self.redis.set(key, "0", ex=300)
|
||||
|
||||
async def get_prefix(self, message: discord.Message) -> str | list[str]:
|
||||
"""Returns the prefix for the message.
|
||||
|
||||
The prefix is case-insensitive, meaning that it'll adapt to the message to be valid not
|
||||
depending of the case. In case that the prefix is invalid, the default or set prefix in its
|
||||
original case will be returned.
|
||||
|
||||
Args:
|
||||
message (discord.Message): The message to check the prefix for.
|
||||
|
||||
Returns:
|
||||
str: The prefix for the message.
|
||||
"""
|
||||
|
||||
if message.guild is None:
|
||||
return ["p.", "P."]
|
||||
|
||||
server_id = message.guild.id
|
||||
|
||||
key = constants.PREFIX_TEMPLATE % server_id
|
||||
|
||||
prefix: bytes | None = await self.redis.get(key)
|
||||
|
||||
if prefix is not None:
|
||||
prefix = prefix.decode()
|
||||
|
||||
else:
|
||||
server, _ = await models.Server.get_or_create(id=server_id)
|
||||
|
||||
await self.redis.set(key, server.prefix, ex=60 * 60)
|
||||
|
||||
prefix = server.prefix
|
||||
|
||||
if message.content.lower().startswith(prefix.lower()):
|
||||
return message.content[: len(prefix)]
|
||||
|
||||
return prefix
|
||||
|
||||
async def set_prefix(
|
||||
self, server: discord.Guild | models.Server | int, prefix: str = "p."
|
||||
) -> None:
|
||||
"""Update a server's prefix.
|
||||
|
||||
Args:
|
||||
server (discord.Guild | models.Server | int): The server to update or it's ID.
|
||||
prefix (str, optional): The new prefix. It cannot be longer than 3 characters. Defaults
|
||||
to "p.".
|
||||
|
||||
Raises:
|
||||
ValueError: The prefix provided is invalid.
|
||||
"""
|
||||
|
||||
if len(prefix) > 3 or len(prefix) == 0:
|
||||
raise ValueError(
|
||||
"The prefix must not be longer than 3 characters and must not be an empty string."
|
||||
)
|
||||
|
||||
server_id = server if isinstance(server, int) else server.id
|
||||
|
||||
await models.Server.filter(id=server_id).update(prefix=prefix)
|
||||
|
||||
await self.redis.set(constants.PREFIX_TEMPLATE % server_id, prefix, ex=60 * 60)
|
||||
|
||||
async def on_command_error(
|
||||
self, context: commands.Context, exception: commands.CommandError
|
||||
) -> None:
|
||||
"""Handles each of the errors raised by commands.
|
||||
|
||||
Args:
|
||||
context (commands.Context): The context of the command.
|
||||
exception (commands.CommandError): The exception raised.
|
||||
"""
|
||||
|
||||
if isinstance(exception, exceptions.ChannelDisabledError):
|
||||
await context.send(
|
||||
"❌ You cannot run commands in this channel.", delete_after=20.0
|
||||
)
|
||||
|
||||
elif isinstance(exception, exceptions.BannedError):
|
||||
await context.send("❌ You have been banned from Payphone.")
|
||||
|
||||
elif isinstance(exception, exceptions.NotOnCallError):
|
||||
await context.send("❌ You need to be on a call to run this command.")
|
||||
|
||||
elif isinstance(exception, exceptions.NotAModeratorError):
|
||||
await context.send("❌ You're not a bot moderator.")
|
||||
|
||||
elif isinstance(exception, errors.MissingRequiredArgument):
|
||||
await context.send(
|
||||
f"❌ `{exception.param.name}` is missing. Use `{context.clean_prefix}{context.command.name}` to learn how to use it."
|
||||
)
|
||||
|
||||
elif isinstance(exception, errors.BotMissingPermissions):
|
||||
await context.send(
|
||||
f"❌ Payphone's missing {formatting.text_join([f"`{perm.replace("_", " ").title()}`" for perm in exception.missing_permissions])} permissions to run this command."
|
||||
)
|
||||
|
||||
elif isinstance(exception, errors.NoPrivateMessage):
|
||||
await context.send("❌ This command can only be run in a server.")
|
||||
|
||||
elif isinstance(exception, errors.MissingPermissions):
|
||||
await context.send("❌ You don't have permission to run this command.")
|
||||
|
||||
elif isinstance(exception, errors.CommandOnCooldown):
|
||||
await context.send(
|
||||
f"❌ You're going too fast! Wait `{int(exception.retry_after)}` seconds before trying again."
|
||||
)
|
||||
|
||||
elif isinstance(exception, errors.CommandNotFound):
|
||||
return
|
||||
|
||||
else:
|
||||
return await super().on_command_error(context, exception)
|
||||
|
||||
@property
|
||||
def embed_author(self) -> discord.EmbedAuthor:
|
||||
"""The embed author for embeds sent by Payphone.
|
||||
|
||||
Returns:
|
||||
discord.EmbedAuthor: The embed author object.
|
||||
"""
|
||||
|
||||
return discord.EmbedAuthor(
|
||||
name=self.user.display_name, icon_url=self.user.display_avatar.url
|
||||
)
|
||||
|
||||
|
||||
bot = Payphone()
|
||||
|
||||
bot.load_extension("payphone.cogs.admin.cog")
|
||||
bot.load_extension("payphone.cogs.calls.cog")
|
||||
bot.load_extension("payphone.cogs.mods.cog")
|
||||
bot.load_extension("payphone.cogs.users.cog")
|
||||
bot.load_extension("payphone.cogs.utility.cog")
|
||||
|
||||
|
||||
@bot.check
|
||||
async def is_channel_enabled(ctx: commands.Context) -> Literal[True]:
|
||||
"""Checks whether the current channel is enabled to run a command or not.
|
||||
|
||||
Args:
|
||||
ctx (commands.Context): The current context.
|
||||
|
||||
Raises:
|
||||
exceptions.ChannelDisabledError: The current channel is disabled.
|
||||
"""
|
||||
|
||||
if ctx.guild is None:
|
||||
return True
|
||||
|
||||
server_id = ctx.guild.id
|
||||
channel_id = ctx.channel.id
|
||||
|
||||
await models.Server.get_or_create(id=server_id)
|
||||
channel, _ = await models.Channel.get_or_create(id=channel_id, server_id=server_id)
|
||||
|
||||
if channel.enabled:
|
||||
return True
|
||||
|
||||
if await models.Channel.filter(server_id=server_id, enabled=True).exists():
|
||||
raise exceptions.ChannelDisabledError()
|
||||
|
||||
return True
|
||||
|
||||
Reference in New Issue
Block a user