204 lines
6 KiB
Python
204 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__}",
|
|
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
|