tuxbot-bot/tuxbot/core/bot.py

200 lines
6 KiB
Python

import asyncio
import datetime
import logging
import sys
from typing import List, Union
import discord
from colorama import Fore, Style, init
from discord.ext import commands
from . import Config
from .data_manager import logs_data_path
from .utils.functions.cli import bordered
from . import __version__
from .utils.functions.extra import ContextPlus
log = logging.getLogger("tuxbot")
init()
NAME = r"""
_____ _ _ _ _
|_ _| ___ _| |__ ___ | |_ | |__ ___ | |_
| || | | \ \/ / '_ \ / _ \| __|____| '_ \ / _ \| __|
| || |_| |> <| |_) | (_) | ||_____| |_) | (_) | |_
|_| \__,_/_/\_\_.__/ \___/ \__| |_.__/ \___/ \__|
"""
packages: List[str] = ["jishaku", "tuxbot.cogs.warnings"]
class Tux(commands.AutoShardedBot):
_loading: asyncio.Task
def __init__(self, *args, cli_flags=None, **kwargs):
# by default, if the bot shutdown without any intervention,
# it's a crash
self.shutdown_code = ExitCodes.CRITICAL
self.cli_flags = cli_flags
self.instance_name = self.cli_flags.instance_name
self.last_exception = None
self.logs = logs_data_path(self.instance_name)
self.config = Config(self.instance_name)
async def _prefixes(bot, message) -> List[str]:
prefixes = self.config("core").get("prefixes")
prefixes.extend(self.config.get_prefixes(message.guild))
if self.config("core").get("mentionable"):
return commands.when_mentioned_or(*prefixes)(bot, message)
return prefixes
if "command_prefix" not in kwargs:
kwargs["command_prefix"] = _prefixes
if "owner_ids" in kwargs:
kwargs["owner_ids"] = set(kwargs["owner_ids"])
else:
kwargs["owner_ids"] = self.config.owners_id()
message_cache_size = 100_000
kwargs["max_messages"] = message_cache_size
self.max_messages = message_cache_size
self.uptime = None
self._app_owners_fetched = False # to prevent abusive API calls
super().__init__(*args, help_command=None, **kwargs)
async def load_packages(self):
if packages:
print("Loading packages...")
for package in packages:
try:
self.load_extension(package)
except Exception as e:
print(
Fore.RED
+ f"Failed to load package {package}"
+ Style.RESET_ALL
+ f" check "
f"{str((self.logs / 'tuxbot.log').resolve())} "
f"for more details"
)
log.exception(f"Failed to load package {package}", exc_info=e)
async def on_ready(self):
self.uptime = datetime.datetime.now()
INFO = {
"title": "INFO",
"rows": [
str(self.user),
f"Prefixes: {', '.join(self.config('core').get('prefixes'))}",
f"Language: {self.config('core').get('locale')}",
f"Tuxbot Version: {__version__}",
f"Discord.py Version: {discord.__version__}",
"Python Version: " + sys.version.replace("\n", ""),
f"Shards: {self.shard_count}",
f"Servers: {len(self.guilds)}",
f"Users: {len(self.users)}",
],
}
COGS = {"title": "COGS", "rows": []}
for extension in packages:
COGS["rows"].append(
f"[{'X' if extension in self.extensions else ' '}] {extension}"
)
print(Fore.LIGHTBLUE_EX + NAME)
print(Style.RESET_ALL)
print(bordered(INFO, COGS))
print(f"\n{'=' * 118}\n\n")
async def is_owner(self, user: Union[discord.User, discord.Member]) -> bool:
"""Determines if the user is a bot owner.
Parameters
----------
user: Union[discord.User, discord.Member]
Returns
-------
bool
"""
if user.id in self.config.owners_id():
return True
owner = False
if not self._app_owners_fetched:
app = await self.application_info()
if app.team:
ids = [m.id for m in app.team.members]
await self.config.update("core", "owners_id", ids)
owner = user.id in ids
self._app_owners_fetched = True
return owner
async def get_context(self, message: discord.Message, *, cls=None):
return await super().get_context(message, cls=ContextPlus)
async def process_commands(self, message: discord.Message):
"""Check for blacklists.
"""
if message.author.bot:
return
if (
message.guild.id in self.config.get_blacklist("guild")
or message.channel.id in self.config.get_blacklist("channel")
or message.author.id in self.config.get_blacklist("user")
):
return
ctx = await self.get_context(message)
if ctx is None or ctx.valid is False:
self.dispatch("message_without_command", message)
else:
await self.invoke(ctx)
async def on_message(self, message: discord.Message):
await self.process_commands(message)
async def logout(self):
"""Disconnect from Discord and closes all actives connections.
Todo: add postgresql logout here
"""
await super().logout()
async def shutdown(self, *, restart: bool = False):
"""Gracefully quit.
Parameters
----------
restart:bool
If `True`, systemd or the launcher gonna see custom exit code
and reboot.
"""
if not restart:
self.shutdown_code = ExitCodes.SHUTDOWN
else:
self.shutdown_code = ExitCodes.RESTART
await self.logout()
sys.exit(self.shutdown_code)
class ExitCodes:
CRITICAL = 1
SHUTDOWN = 0
RESTART = 42