277 lines
9.8 KiB
Python
277 lines
9.8 KiB
Python
from __future__ import annotations
|
|
|
|
import discord
|
|
from discord import app_commands
|
|
from discord.ext import tasks, commands
|
|
|
|
import logging
|
|
import xml.etree.ElementTree as et
|
|
from aiohttp import ClientSession
|
|
import asyncpg
|
|
from typing import TYPE_CHECKING, Optional
|
|
|
|
import constants
|
|
|
|
log = logging.getLogger("lina.main")
|
|
|
|
if TYPE_CHECKING:
|
|
from cogs import PlayerTrack
|
|
from cogs import Online
|
|
|
|
extensions = (
|
|
"cogs.online",
|
|
"cogs.playertrack",
|
|
"cogs.core",
|
|
"cogs.misc",
|
|
"cogs.pokemap",
|
|
"cogs.games"
|
|
)
|
|
|
|
|
|
class STKRequestError(Exception):
|
|
"""Raised when an error occurs upon performing a request to STK servers."""
|
|
pass
|
|
|
|
|
|
class STKAddonsAPIError(Exception):
|
|
pass
|
|
|
|
|
|
class Lina(commands.Bot):
|
|
"""
|
|
Class representing lina herself.
|
|
"""
|
|
|
|
pool: asyncpg.Pool
|
|
|
|
def __init__(self):
|
|
|
|
intents = discord.Intents.default()
|
|
intents.message_content = True
|
|
intents.members = True
|
|
|
|
allowed_mentions = discord.AllowedMentions(everyone=False, roles=False)
|
|
|
|
super().__init__(intents=intents,
|
|
activity=discord.Game(name="SuperTuxKart"),
|
|
allowed_mentions=allowed_mentions,
|
|
command_prefix=constants.PREFIX)
|
|
|
|
self.accent_color = constants.ACCENT_COLOR
|
|
self.stk_userid: int = None
|
|
self.stk_token: str = None
|
|
self.stk_sessid: str = None
|
|
self.version = "1.1.3"
|
|
|
|
async def stkPostReq(self, target, args):
|
|
"""Helper function to send a POST request to STK servers."""
|
|
assert self.session is not None
|
|
log.debug(
|
|
"Sending %s to %s",
|
|
args.replace(str(self.stk_token), "[REDACTED]").replace(constants.STK_PASSWORD, "[REDACTED]"),
|
|
str(self.session._base_url) + target
|
|
)
|
|
async with self.session.post(
|
|
target, data=args,
|
|
headers={
|
|
**self.session.headers,
|
|
"Content-Type": "application/x-www-form-urlencoded"
|
|
}
|
|
) as r:
|
|
r.raise_for_status()
|
|
|
|
data = et.fromstring(await r.text())
|
|
|
|
if data.attrib["success"] == "no":
|
|
raise STKRequestError(data.attrib["info"])
|
|
else:
|
|
return data
|
|
|
|
async def stkGetReq(self, target):
|
|
"""Helper function to send a GET request to STK servers."""
|
|
assert self.session is not None
|
|
async with self.session.get(target) as r:
|
|
r.raise_for_status()
|
|
return et.fromstring(await r.text())
|
|
|
|
async def authSTK(self):
|
|
"""Authenticate to STK"""
|
|
log.info(f"Trying to authenticate STK account {constants.STK_USERNAME}")
|
|
loginPayload = await self.stkPostReq(
|
|
"/api/v2/user/connect",
|
|
f"username={constants.STK_USERNAME}&"
|
|
f"password={constants.STK_PASSWORD}&"
|
|
"save-session=true"
|
|
)
|
|
|
|
self.stk_userid = loginPayload.attrib["userid"]
|
|
self.stk_token = loginPayload.attrib["token"]
|
|
|
|
log.info(f"STK user {loginPayload.attrib['username']} logged in successfully.")
|
|
if not self.stkPoll.is_running():
|
|
self.stkPoll.start()
|
|
|
|
@tasks.loop(minutes=1)
|
|
async def stkPoll(self):
|
|
try:
|
|
await self.stkPostReq(
|
|
"/api/v2/user/poll",
|
|
f"userid={self.stk_userid}&"
|
|
f"token={self.stk_token}"
|
|
)
|
|
except STKRequestError as e:
|
|
if str(e) in "Session not valid. Please sign in.":
|
|
log.warning("Session was invalidated. Reauthenticating...")
|
|
await self.authSTK()
|
|
else:
|
|
log.error("Poll request failed: %s", e)
|
|
except Exception:
|
|
log.exception("Poll request failed due to exception:")
|
|
|
|
@tasks.loop(minutes=10)
|
|
async def stkaddonsSessionPersister(self):
|
|
"""
|
|
A very convoluted way of retaining the SuperTuxKart Addons Website's
|
|
STK_SESSID. This is needed to make requests to the frontend API
|
|
"""
|
|
|
|
log.debug("Trying to obtain STK_SESSID")
|
|
|
|
req = await self.session.post("/login.php?action=submit",
|
|
headers={
|
|
**self.session.headers,
|
|
"Content-Type": "application/x-www-form-urlencoded"
|
|
},
|
|
data=(
|
|
f"username={constants.STK_USERNAME}&"
|
|
f"password={constants.STK_PASSWORD}"
|
|
))
|
|
|
|
# A successful login will result in a redirect so check if it happened.
|
|
if not len(req.history) == 1:
|
|
log.warning("Frontend API Login was not successful")
|
|
else:
|
|
self.stk_sessid = req.cookies["STK_SESSID"].value
|
|
log.debug("Successfully obtained STK_SESSID: %s",
|
|
self.stk_sessid)
|
|
|
|
|
|
async def on_command_error(self, ctx: commands.Context, error: commands.CommandError):
|
|
log.exception("%s: Command error occurred", ctx.command.name, exc_info=error)
|
|
if isinstance(error, commands.NoPrivateMessage):
|
|
return await ctx.author.send("You can't use this command in private messages.")
|
|
|
|
if isinstance(error, commands.NotOwner):
|
|
return await ctx.send("This command is restricted to the owner only.")
|
|
|
|
if isinstance(error, commands.CommandInvokeError):
|
|
error: commands.CommandInvokeError
|
|
original = error.original
|
|
|
|
if isinstance(original, STKRequestError):
|
|
return await ctx.send(embed=discord.Embed(
|
|
title="Sorry, an STK related error occurred.",
|
|
description=str(original),
|
|
color=self.accent_color
|
|
))
|
|
|
|
return await ctx.send(embed=discord.Embed(
|
|
title="Sorry, this shouldn't have happened. Guru Meditation.",
|
|
description=f"```\n{error.original.__class__.__name__}: {str(error.original)}\n```",
|
|
color=self.accent_color
|
|
))
|
|
|
|
async def on_app_command_error(self, interaction: discord.Interaction, error: app_commands.AppCommandError):
|
|
log.exception("%s: App command error occurred.", interaction.command.name, exc_info=error)
|
|
|
|
if isinstance(error, app_commands.CommandInvokeError):
|
|
error: app_commands.CommandInvokeError
|
|
original = error.original
|
|
|
|
if isinstance(original, STKRequestError):
|
|
return await interaction.response.send_message(embed=discord.Embed(
|
|
title="Sorry, an STK related error occurred.",
|
|
description=str(original),
|
|
color=self.accent_color
|
|
), ephemeral=True)
|
|
|
|
return await interaction.response.send_message(embed=discord.Embed(
|
|
title="Sorry, this shouldn't have happened. Guru Meditation.",
|
|
description=f"```\n{error.original.__class__.__name__}: {str(error.original)}\n```",
|
|
color=self.accent_color
|
|
), ephemeral=True)
|
|
|
|
async def afterStkAuth(self):
|
|
self.stkaddonsSessionPersister.start()
|
|
for extension in extensions:
|
|
log.debug("Loading extension %s", extension)
|
|
try:
|
|
await self.load_extension(extension)
|
|
except Exception:
|
|
log.exception(f"Unable to load extension {extension}.")
|
|
else:
|
|
log.debug("Successfully loaded extension %s.", extension)
|
|
|
|
async def setup_hook(self):
|
|
self.tree.error(self.on_app_command_error)
|
|
|
|
if not hasattr(self, "uptime"):
|
|
self.uptime = discord.utils.utcnow()
|
|
|
|
self.session = ClientSession(
|
|
"https://online.supertuxkart.net",
|
|
headers={
|
|
"User-Agent": f"linaSTK-Discord/{self.version} (+https://linastk.codeberg.page/linastk_ua.html)"
|
|
}
|
|
)
|
|
|
|
try:
|
|
await self.authSTK()
|
|
except Exception as e:
|
|
log.exception("STK account authentication failed. "
|
|
"See below for details.",
|
|
exc_info=e)
|
|
await self.close()
|
|
else:
|
|
await self.afterStkAuth()
|
|
|
|
async def close(self):
|
|
"""Shut down lina"""
|
|
|
|
log.info("lina is shutting down...")
|
|
if hasattr(self, 'session'):
|
|
# it's important to check if the session is closed in case
|
|
# something went wrong before the session is set up
|
|
if not self.session.closed:
|
|
try:
|
|
if self.stk_userid and self.stk_token:
|
|
# send a client quit request so that internal online
|
|
# counter is deducted and mark the STK account offline
|
|
await self.stkPostReq("/api/v2/user/client-quit",
|
|
f"userid={self.stk_userid}&"
|
|
f"token={self.stk_token}")
|
|
else:
|
|
log.warning("userid and token is absent. "
|
|
"will not send client quit request.")
|
|
finally:
|
|
await self.session.close()
|
|
|
|
await super().close()
|
|
|
|
async def start(self):
|
|
"""Bring lina to life"""
|
|
log.info("Starting bot...")
|
|
await super().start(constants.TOKEN, reconnect=True)
|
|
|
|
@property
|
|
def playertrack(self) -> Optional[PlayerTrack]:
|
|
"""Represents the PlayerTrack cog"""
|
|
return self.get_cog("PlayerTrack")
|
|
|
|
@property
|
|
def online(self) -> Optional[Online]:
|
|
"""Represents the Online cog"""
|
|
return self.get_cog("Online")
|
|
|
|
async def on_ready(self):
|
|
log.info(f"Bot {self.user} ({self.user.id}) is ready!")
|