490 lines
15 KiB
Python
490 lines
15 KiB
Python
from typing import Literal
|
|
|
|
from datetime import datetime, timezone, timedelta
|
|
|
|
from multiprocessing import managers
|
|
|
|
import sys
|
|
|
|
import discord
|
|
from discord import abc
|
|
from discord.ext import commands
|
|
from discord.ext.commands import errors
|
|
|
|
from tortoise import expressions, transactions, exceptions as tortoise_exc
|
|
|
|
import redis.asyncio as aredis
|
|
|
|
import asyncpg.exceptions
|
|
|
|
from payphone.db import models
|
|
from payphone.core import (
|
|
constants,
|
|
exceptions,
|
|
formatting,
|
|
help_command,
|
|
config,
|
|
logging,
|
|
)
|
|
|
|
|
|
class Payphone(commands.AutoShardedBot):
|
|
log = logging.logger
|
|
|
|
redis: aredis.Redis = aredis.Redis(
|
|
host="localhost" if config.IS_DEBUG else "redis.default.svc.cluster.local"
|
|
)
|
|
|
|
def __init__(self) -> None:
|
|
intents = discord.Intents.default()
|
|
intents.message_content = True
|
|
intents.guilds = True
|
|
|
|
if config.IS_DEBUG:
|
|
intents.members = True
|
|
|
|
super().__init__(
|
|
intents=intents,
|
|
command_prefix=self.get_prefix,
|
|
help_command=help_command.HelpCommand(),
|
|
strip_after_prefix=True,
|
|
case_insensitive=True,
|
|
owner_ids=[898569996949676052, 777338793803513886],
|
|
)
|
|
|
|
async def on_ready(self) -> None:
|
|
"""Runs when the bot has finished initializing."""
|
|
|
|
self.log.info("Ready! Logged in as %s.", self.user.display_name)
|
|
|
|
async def set_presence(self, shard_id: int) -> None:
|
|
"""Sets the presence of the bot back to online.
|
|
|
|
Args:
|
|
shard_id (int): The ID of the shard to set the status for.
|
|
"""
|
|
|
|
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,
|
|
)
|
|
|
|
async def on_shard_connect(self, shard_id: int) -> None:
|
|
"""Runs when the shard connects.
|
|
|
|
Args:
|
|
shard_id (int): The ID of the shard that connected.
|
|
"""
|
|
|
|
self.log.info("Shard has connected.")
|
|
|
|
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.
|
|
"""
|
|
|
|
self.log.info("Shard is ready.")
|
|
|
|
await self.set_presence(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.
|
|
"""
|
|
|
|
self.log.info("Shard has resumed.")
|
|
|
|
await self.set_presence(shard_id)
|
|
|
|
async def get_or_fetch_channel(
|
|
self, channel_id: int
|
|
) -> abc.GuildChannel | abc.PrivateChannel | discord.Thread | 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 | discord.Thread | 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__gt=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,
|
|
duration: timedelta | None,
|
|
reason: str,
|
|
moderator: discord.User | discord.Member | models.User | int | None = None,
|
|
report_id: 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.
|
|
duration (timedelta | 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.
|
|
report_id (int | None, optional): The ID of the report the ban is related to.
|
|
|
|
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
|
|
)
|
|
|
|
until = datetime.now(timezone.utc) + duration if duration is not None else None
|
|
|
|
if duration is not None:
|
|
last_ban = (
|
|
await models.Ban.filter(
|
|
expressions.Q(
|
|
expires_at=None,
|
|
expires_at__gt=datetime.now(timezone.utc),
|
|
join_type="OR",
|
|
),
|
|
user_id=user_id,
|
|
)
|
|
.order_by("-created_at")
|
|
.limit(1)
|
|
.first()
|
|
)
|
|
|
|
if last_ban is not None:
|
|
until = (
|
|
duration + last_ban.expires_at
|
|
if last_ban.expires_at is not None
|
|
else None
|
|
)
|
|
|
|
row = await models.Ban.create(
|
|
user_id=user_id,
|
|
reason=reason,
|
|
expires_at=until,
|
|
banned_by_id=mod_id,
|
|
report_id=report_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,
|
|
unbanned_by: discord.User | discord.Member | models.User | int | None = None,
|
|
) -> None:
|
|
"""Unbans a user from Payphone.
|
|
|
|
Args:
|
|
user (discord.User | discord.Member | models.User | int): The user to unban or their
|
|
ID.
|
|
unbanned_by (discord.User | discord.Member | models.User | int | None, optional): The
|
|
moderator who unbanned the user or their ID.
|
|
"""
|
|
|
|
user_id = user if isinstance(user, int) else user.id
|
|
unbanned_by_id = unbanned_by if isinstance(unbanned_by, int) else unbanned_by.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), unbanned_by_id=unbanned_by_id)
|
|
|
|
await self.redis.set(key, "0", ex=300)
|
|
|
|
async def get_prefix(
|
|
self, message: discord.Message | None = None, server_id: int | None = None
|
|
) -> 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.
|
|
|
|
This method can also be used on a server's raw ID. Either `message` or `server_id` must be
|
|
set, but not both.
|
|
|
|
Args:
|
|
message (discord.Message | None): The message to check the prefix for, or None.
|
|
server_id (int | None): The server's ID to get the prefix from.
|
|
|
|
Returns:
|
|
str: The prefix for the message.
|
|
|
|
Raises:
|
|
ValueError: Either `message` or `server_id` must be set, but not both.
|
|
"""
|
|
|
|
default_prefixes = ["p.", "P.", self.user.mention]
|
|
|
|
if message is not None and server_id is not None:
|
|
raise ValueError(
|
|
"Either `message` or `server_id` must be set, but not both."
|
|
)
|
|
|
|
if message is None and server_id is None:
|
|
raise ValueError("Either `message` or `server_id` must be set.")
|
|
|
|
if message is not None and message.guild is None:
|
|
return default_prefixes
|
|
|
|
if message is not None and message.content is None:
|
|
return default_prefixes
|
|
|
|
if (
|
|
message is not None
|
|
and message.content is not None
|
|
and message.content.startswith(self.user.mention)
|
|
):
|
|
return default_prefixes
|
|
|
|
server_id = server_id if server_id is not None else 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.update_or_create(
|
|
(
|
|
{"name": message.guild.name}
|
|
if message is not None and message.guild is not None
|
|
else None
|
|
),
|
|
id=server_id,
|
|
)
|
|
|
|
await self.redis.set(key, server.prefix, ex=60 * 60)
|
|
|
|
prefix = server.prefix
|
|
|
|
if message is not None and 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 5 characters. Defaults
|
|
to "p.".
|
|
|
|
Raises:
|
|
ValueError: The prefix provided is invalid.
|
|
"""
|
|
|
|
if len(prefix) > 5 or len(prefix) == 0:
|
|
raise ValueError(
|
|
"The prefix must not be longer than 5 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}help {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.NotOwner):
|
|
await context.send("❌ You're not the owner of this bot.")
|
|
|
|
elif isinstance(exception, errors.CommandNotFound):
|
|
return
|
|
|
|
elif isinstance(exception, tortoise_exc.TransactionManagementError):
|
|
# Tortoise ORM bug. This may need to be logged for debugging, but for now they just
|
|
# clutter the logs.
|
|
return
|
|
|
|
else:
|
|
return await super().on_command_error(context, exception)
|
|
|
|
async def on_error(
|
|
self, event_method: str, *args: discord.Any, **kwargs: discord.Any
|
|
) -> None:
|
|
_, exception, _ = sys.exc_info()
|
|
|
|
if isinstance(
|
|
exception,
|
|
(asyncpg.exceptions.InterfaceError, asyncpg.exceptions.PostgresError),
|
|
):
|
|
# IMPORTANT! These are just spam at the moment, but may cause trouble in the future as
|
|
# the bot grows.
|
|
return
|
|
|
|
return await super().on_error(event_method, *args, **kwargs)
|
|
|
|
@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.easter.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
|
|
|
|
if ctx.command.qualified_name in ["config enable", "config disable"]:
|
|
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 channel.enabled is False:
|
|
raise exceptions.ChannelDisabledError()
|
|
|
|
if await models.Channel.filter(server_id=server_id, enabled=True).exists():
|
|
raise exceptions.ChannelDisabledError()
|
|
|
|
return True
|