tuxbot-bot/tuxbot/core/bot.py

361 lines
11 KiB
Python
Raw Normal View History

2020-06-05 00:29:14 +02:00
import asyncio
import datetime
import importlib
2020-06-04 19:16:51 +02:00
import logging
2020-10-22 00:00:48 +02:00
from collections import Counter
from typing import List, Tuple
2020-10-22 00:00:48 +02:00
import aiohttp
2020-06-04 19:16:51 +02:00
import discord
from discord.ext import commands
from rich import box
from rich.columns import Columns
from rich.panel import Panel
from rich.progress import Progress
from rich.table import Table
from tortoise import Tortoise
from tuxbot import version_info
from tuxbot.core.utils.data_manager import (
logs_data_path,
config_file,
)
from tuxbot.core.utils.functions.extra import ContextPlus
from tuxbot.core.utils.functions.prefix import get_prefixes
from tuxbot.core.utils.console import console
from tuxbot.core.config import (
2020-10-19 22:17:19 +02:00
Config,
ConfigFile,
search_for,
)
2020-08-28 01:06:57 +02:00
from . import __version__, ExitCodes
from . import exceptions
2020-06-04 19:16:51 +02:00
log = logging.getLogger("tuxbot")
2021-04-21 18:28:09 +02:00
packages: Tuple = (
"jishaku",
2020-11-09 01:18:55 +01:00
"tuxbot.cogs.Admin",
"tuxbot.cogs.Logs",
"tuxbot.cogs.Dev",
2020-11-09 01:18:55 +01:00
"tuxbot.cogs.Utils",
"tuxbot.cogs.Polls",
"tuxbot.cogs.Custom",
"tuxbot.cogs.Network",
"tuxbot.cogs.Linux",
2021-05-13 16:45:49 +02:00
"tuxbot.cogs.Mod",
2021-05-16 00:32:14 +02:00
"tuxbot.cogs.Tags",
2021-04-21 18:28:09 +02:00
)
class Tux(commands.AutoShardedBot):
2020-06-05 00:29:14 +02:00
_loading: asyncio.Task
_progress = {"tasks": {}, "main": Progress()}
2020-06-05 00:29:14 +02:00
def __init__(self, *args, cli_flags=None, **kwargs):
2020-06-03 19:41:30 +02:00
# by default, if the bot shutdown without any intervention,
# it's a crash
2020-06-05 00:29:14 +02:00
self.shutdown_code = ExitCodes.CRITICAL
self.cli_flags = cli_flags
self.last_exception = None
self.logs = logs_data_path()
2020-06-03 19:41:30 +02:00
2020-10-22 00:00:48 +02:00
self.console = console
self.stats = {"commands": Counter(), "socket": Counter()}
self.config: Config = ConfigFile(config_file, Config).config
self.instance_name = self.config.Core.instance_name
2020-06-04 19:16:51 +02:00
async def _prefixes(bot, message) -> List[str]:
prefixes = self.config.Core.prefixes
2020-06-05 00:29:14 +02:00
prefixes.extend(get_prefixes(self, message.guild))
2020-06-04 19:16:51 +02:00
if self.config.Core.mentionable:
2020-06-04 19:16:51 +02:00
return commands.when_mentioned_or(*prefixes)(bot, message)
return prefixes
if "command_prefix" not in kwargs:
kwargs["command_prefix"] = _prefixes
2020-06-03 19:41:30 +02:00
if "owner_ids" in kwargs:
kwargs["owner_ids"] = set(kwargs["owner_ids"])
else:
kwargs["owner_ids"] = self.config.Core.owners_id
2020-06-03 19:41:30 +02:00
2021-04-21 18:28:09 +02:00
message_cache_size = 10_000
2020-06-03 19:41:30 +02:00
kwargs["max_messages"] = message_cache_size
self.max_messages = message_cache_size
2020-06-03 19:41:30 +02:00
self.uptime = None
2021-02-11 18:11:19 +01:00
self.last_on_ready = None
2020-06-05 00:29:14 +02:00
self._app_owners_fetched = False # to prevent abusive API calls
self.loop = asyncio.get_event_loop()
2020-11-09 01:27:19 +01:00
super().__init__(
*args,
intents=discord.Intents.all(),
loop=self.loop,
**kwargs,
2020-11-09 01:27:19 +01:00
)
2020-10-22 00:00:48 +02:00
self.session = aiohttp.ClientSession(loop=self.loop)
async def _is_blacklisted(self, message: discord.Message) -> bool:
"""Check for blacklists."""
if message.author.bot:
return True
if (
search_for(self.config.Users, message.author.id, "blacklisted")
or search_for(
self.config.Channels, message.channel.id, "blacklisted"
)
or (
message.guild
and search_for(
self.config.Servers, message.guild.id, "blacklisted"
)
)
):
return True
return False
2020-06-05 00:29:14 +02:00
async def load_packages(self):
if packages:
with Progress() as progress:
task = progress.add_task(
"Loading packages...", total=len(packages)
)
for package in packages:
try:
self.load_extension(package)
progress.console.print(f"{package} loaded")
except Exception as e:
log.exception(
"Failed to load package %s", package, exc_info=e
)
progress.console.print(
f"[red]Failed to load package {package} "
f"[i](see "
f"{str((self.logs / 'tuxbot.log').resolve())} "
f"for more details)[/i]"
)
progress.advance(task)
2020-06-06 18:51:47 +02:00
2020-06-04 19:16:51 +02:00
async def on_ready(self):
2021-02-11 18:11:19 +01:00
if self.uptime is not None:
self.last_on_ready = datetime.datetime.now()
return
2020-06-05 00:29:14 +02:00
self.uptime = datetime.datetime.now()
2021-02-11 18:11:19 +01:00
self.last_on_ready = self.uptime
2020-10-19 22:17:19 +02:00
with self._progress["main"] as progress:
progress.stop_task(self._progress["tasks"]["discord_connecting"])
progress.remove_task(self._progress["tasks"]["discord_connecting"])
self._progress["tasks"].pop("discord_connecting")
self.console.clear()
self.console.print(
Panel(f"[bold blue]Tuxbot V{version_info.major}", style="blue"),
justify="center",
)
self.console.print()
columns = Columns(align="center", expand=True)
table = Table(style="dim", border_style="not dim", box=box.HEAVY_HEAD)
table.add_column(
"INFO",
)
table.add_row(str(self.user))
table.add_row(f"Prefixes: {', '.join(self.config.Core.prefixes)}")
table.add_row(f"Language: {self.config.Core.locale}")
table.add_row(f"Tuxbot Version: {__version__}")
table.add_row(f"Discord.py Version: {discord.__version__}")
table.add_row(f"Instance name: {self.instance_name}")
table.add_row(f"Shards: {self.shard_count}")
table.add_row(f"Servers: {len(self.guilds)}")
table.add_row(f"Users: {len(self.users)}")
columns.add_renderable(table)
table = Table(style="dim", border_style="not dim", box=box.HEAVY_HEAD)
table.add_column(
"COGS",
)
2020-06-05 00:29:14 +02:00
for extension in packages:
if extension in self.extensions:
status = f"[green]:heavy_check_mark: {extension} "
else:
status = f"[red]:heavy_multiplication_x: {extension} "
2020-06-04 19:16:51 +02:00
table.add_row(status)
columns.add_renderable(table)
2020-06-04 19:16:51 +02:00
self.console.print(columns)
self.console.print()
2020-06-05 00:29:14 +02:00
async def is_owner(
self, user: discord.User | discord.Member | discord.Object
) -> bool:
2020-06-05 00:29:14 +02:00
"""Determines if the user is a bot owner.
Parameters
----------
user: discord.User | discord.Member
2020-06-05 00:29:14 +02:00
Returns
-------
bool
"""
if user.id in self.config.Core.owners_id:
2020-06-05 00:29:14 +02:00
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.Core.owners_id = ids
2020-06-05 00:29:14 +02:00
owner = user.id in ids
self._app_owners_fetched = True
return owner
# pylint: disable=unused-argument
2020-06-05 00:29:14 +02:00
async def get_context(self, message: discord.Message, *, cls=None):
ctx: ContextPlus = await super().get_context(message, cls=ContextPlus)
2020-06-05 00:29:14 +02:00
if (ctx is None or not ctx.valid) and (
user_aliases := search_for(
self.config.Users, message.author.id, "aliases"
)
):
# noinspection PyUnboundLocalVariable
for alias, command in user_aliases.items():
back_content = message.content
message.content = message.content.replace(alias, command, 1)
if (
ctx := await super().get_context(message, cls=ContextPlus)
) is None or not ctx.valid:
message.content = back_content
else:
break
return ctx
async def process_commands(self, message: discord.Message):
ctx: ContextPlus = await self.get_context(message)
if ctx is not None and ctx.valid:
2021-01-30 16:43:42 +01:00
if message.guild and ctx.command in search_for(
2021-04-21 18:28:09 +02:00
self.config.Servers, message.guild.id, "disabled_command", ()
):
raise exceptions.DisabledCommandByServerOwner
if ctx.command in self.config.Core.disabled_command:
raise exceptions.DisabledCommandByBotOwner
2020-06-05 00:29:14 +02:00
await self.invoke(ctx)
else:
self.dispatch("message_without_command", message)
2020-06-05 00:29:14 +02:00
async def on_message(self, message: discord.Message):
if not await self._is_blacklisted(message):
await self.process_commands(message)
2020-06-05 00:29:14 +02:00
async def start(self, token): # pylint: disable=arguments-differ
"""Connect to Discord and start all connections."""
with Progress() as progress:
task = progress.add_task(
"Connecting to PostgreSQL...", total=len(self.extensions)
)
models = []
for extension, _ in self.extensions.items():
if extension == "jishaku":
progress.advance(task)
continue
if importlib.import_module(extension).HAS_MODELS:
models.append(f"{extension}.models.__init__")
progress.advance(task)
await Tortoise.init(
db_url="postgres://{}:{}@{}:{}/{}".format(
self.config.Core.Database.username,
self.config.Core.Database.password,
self.config.Core.Database.domain,
self.config.Core.Database.port,
self.config.Core.Database.db_name,
),
modules={"models": models},
)
await Tortoise.generate_schemas()
with self._progress["main"] as progress:
task_id = self._progress["tasks"][
"discord_connecting"
] = progress.add_task(
"Connecting to Discord...",
task_name="discord_connecting",
start=False,
)
progress.update(task_id)
await super().start(token)
2020-06-05 00:29:14 +02:00
async def logout(self):
"""Disconnect from Discord and closes all actives connections.
Todo: add postgresql logout here
"""
with self._progress["main"] as progress:
for task in self._progress["tasks"]:
progress.log("Shutting down", task)
progress.stop_task(self._progress["tasks"][task])
progress.remove_task(self._progress["tasks"][task])
progress.stop()
pending = [
t for t in asyncio.all_tasks() if t is not asyncio.current_task()
]
for task in pending:
self.console.log(
"Canceling", task.get_name(), f"({task.get_coro()})"
)
task.cancel()
2020-09-02 00:08:06 +02:00
await asyncio.gather(*pending, return_exceptions=False)
await super().close()
2020-06-05 00:29:14 +02:00
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.
"""
2021-02-11 23:18:12 +01:00
self.shutdown_code = (
ExitCodes.RESTART if restart else ExitCodes.SHUTDOWN
)
2020-06-05 00:29:14 +02:00
await self.logout()
2020-09-02 00:08:06 +02:00
sys_e = SystemExit()
sys_e.code = self.shutdown_code
raise sys_e