tuxbot-bot/tuxbot/core/bot.py

205 lines
6.1 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]
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