145 lines
5.7 KiB
Python
145 lines
5.7 KiB
Python
import asyncio
|
|
import json
|
|
import logging
|
|
import os
|
|
import time
|
|
from datetime import datetime
|
|
|
|
import discord
|
|
|
|
from pluralkit import db
|
|
from pluralkit.bot import channel_logger, commands, proxy
|
|
from pluralkit.stats import InfluxStatCollector, NullStatCollector
|
|
|
|
logging.basicConfig(level=logging.INFO, format="[%(asctime)s] [%(name)s] [%(levelname)s] %(message)s")
|
|
# logging.getLogger("pluralkit").setLevel(logging.DEBUG)
|
|
|
|
class PluralKitBot:
|
|
def __init__(self, token):
|
|
self.token = token
|
|
self.logger = logging.getLogger("pluralkit.bot")
|
|
|
|
self.client = discord.Client()
|
|
self.client.event(self.on_error)
|
|
self.client.event(self.on_ready)
|
|
self.client.event(self.on_message)
|
|
self.client.event(self.on_socket_raw_receive)
|
|
|
|
self.stats = NullStatCollector()
|
|
|
|
self.channel_logger = channel_logger.ChannelLogger(self.client)
|
|
self.proxy = proxy.Proxy(self.client, token, self.channel_logger, self.stats)
|
|
|
|
async def on_error(self, evt, *args, **kwargs):
|
|
self.logger.exception("Error while handling event {} with arguments {}:".format(evt, args))
|
|
|
|
async def on_ready(self):
|
|
self.logger.info("Connected to Discord.")
|
|
self.logger.info("- Account: {}#{}".format(self.client.user.name, self.client.user.discriminator))
|
|
self.logger.info("- User ID: {}".format(self.client.user.id))
|
|
self.logger.info("- {} servers".format(len(self.client.servers)))
|
|
|
|
async def on_message(self, message):
|
|
# Ignore bot messages
|
|
if message.author.bot:
|
|
return
|
|
|
|
if await self.handle_command_dispatch(message):
|
|
return
|
|
|
|
if await self.handle_proxy_dispatch(message):
|
|
return
|
|
|
|
async def on_socket_raw_receive(self, msg):
|
|
# Since on_reaction_add is buggy (only works for messages the bot's already cached, ie. no old messages)
|
|
# we parse socket data manually for the reaction add event
|
|
if isinstance(msg, str):
|
|
try:
|
|
msg_data = json.loads(msg)
|
|
if msg_data.get("t") == "MESSAGE_REACTION_ADD":
|
|
evt_data = msg_data.get("d")
|
|
if evt_data:
|
|
user_id = evt_data["user_id"]
|
|
message_id = evt_data["message_id"]
|
|
emoji = evt_data["emoji"]["name"]
|
|
|
|
async with self.pool.acquire() as conn:
|
|
await self.proxy.handle_reaction(conn, user_id, message_id, emoji)
|
|
elif msg_data.get("t") == "MESSAGE_DELETE":
|
|
evt_data = msg_data.get("d")
|
|
if evt_data:
|
|
message_id = evt_data["id"]
|
|
async with self.pool.acquire() as conn:
|
|
await self.proxy.handle_deletion(conn, message_id)
|
|
except ValueError:
|
|
pass
|
|
|
|
async def handle_command_dispatch(self, message):
|
|
command_items = commands.command_list.items()
|
|
command_items = sorted(command_items, key=lambda x: len(x[0]), reverse=True)
|
|
|
|
prefix = "pk;"
|
|
for command_name, command in command_items:
|
|
if message.content.lower().startswith(prefix + command_name):
|
|
args_str = message.content[len(prefix + command_name):].strip()
|
|
args = args_str.split(" ")
|
|
|
|
# Splitting on empty string yields one-element array, remove that
|
|
if len(args) == 1 and not args[0]:
|
|
args = []
|
|
|
|
async with self.pool.acquire() as conn:
|
|
time_before = time.perf_counter()
|
|
await command.function(self.client, conn, message, args)
|
|
time_after = time.perf_counter()
|
|
|
|
# Report command time stats
|
|
execution_time = time_after - time_before
|
|
response_time = (datetime.now() - message.timestamp).total_seconds()
|
|
await self.stats.report_command(command_name, execution_time, response_time)
|
|
|
|
return True
|
|
|
|
async def handle_proxy_dispatch(self, message):
|
|
# Try doing proxy parsing
|
|
async with self.pool.acquire() as conn:
|
|
return await self.proxy.try_proxy_message(conn, message)
|
|
|
|
async def periodical_stat_timer(self, pool):
|
|
async with pool.acquire() as conn:
|
|
while True:
|
|
await self.stats.report_periodical_stats(conn)
|
|
await asyncio.sleep(30)
|
|
|
|
async def run(self):
|
|
try:
|
|
self.logger.info("Connecting to database...")
|
|
self.pool = await db.connect(
|
|
os.environ["DATABASE_USER"],
|
|
os.environ["DATABASE_PASS"],
|
|
os.environ["DATABASE_NAME"],
|
|
os.environ["DATABASE_HOST"],
|
|
int(os.environ["DATABASE_PORT"])
|
|
)
|
|
|
|
self.logger.info("Attempting to create tables...")
|
|
async with self.pool.acquire() as conn:
|
|
await db.create_tables(conn)
|
|
|
|
if "INFLUX_HOST" in os.environ:
|
|
self.logger.info("Connecting to InfluxDB...")
|
|
self.stats = await InfluxStatCollector.connect(
|
|
os.environ["INFLUX_HOST"],
|
|
os.environ["INFLUX_PORT"],
|
|
os.environ["INFLUX_DB"]
|
|
)
|
|
|
|
self.logger.info("Starting periodical stat reporting...")
|
|
asyncio.get_event_loop().create_task(self.periodical_stat_timer(self.pool))
|
|
|
|
self.logger.info("Connecting to Discord...")
|
|
await self.client.start(self.token)
|
|
finally:
|
|
self.logger.info("Logging out from Discord...")
|
|
await self.client.logout()
|