Skip to content

Commit

Permalink
Merge pull request #29 from srobo/kjk/logs
Browse files Browse the repository at this point in the history
Integrate logs bot functionality
raccube authored Aug 28, 2024
2 parents d8c8bdb + c194f5d commit 485afa7
Showing 4 changed files with 374 additions and 7 deletions.
2 changes: 2 additions & 0 deletions src/bot.py
Original file line number Diff line number Diff line change
@@ -20,6 +20,7 @@
PASSWORDS_CHANNEL_NAME,
)
from src.commands.join import join
from src.commands.logs import logs
from src.commands.team import (
Team,
new_team,
@@ -64,6 +65,7 @@ def __init__(
team.add_command(export_team)
self.tree.add_command(team, guild=self.guild)
self.tree.add_command(join, guild=self.guild)
self.tree.add_command(logs, guild=self.guild)

async def setup_hook(self) -> None:
# This copies the global commands over to your guild.
363 changes: 363 additions & 0 deletions src/commands/logs.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,363 @@
import os
import re
import sys
import shutil
import logging
import tempfile
from enum import Enum
from typing import IO, cast, List, Tuple, TYPE_CHECKING
from pathlib import Path
from zipfile import ZipFile, BadZipFile, is_zipfile, ZIP_DEFLATED
from datetime import date

import aiohttp
import discord
from discord import app_commands

from src.constants import TEAM_CHANNEL_PREFIX

if TYPE_CHECKING:
from src.bot import BotClient


class AnimationHandling(Enum):
none = 0
team = 1
separate = 2


logger = logging.getLogger("logs")
logger.setLevel(logging.INFO)
handler = logging.StreamHandler(sys.stdout)
handler.setLevel(logging.INFO)
logger.addHandler(handler)

# Don't post to team channels and force the guild used so testing can you DMs
DISCORD_TESTING = bool(os.getenv('DISCORD_TESTING'))
# Just post all messages to calling channel, allow DMs
DISCORD_DEBUG = bool(os.getenv('DISCORD_DEBUG'))
if DISCORD_TESTING or DISCORD_DEBUG:
# print all debug messages
logger.setLevel(logging.DEBUG)
handler.setLevel(logging.DEBUG)


async def log_and_reply(ctx: discord.interactions.Interaction["BotClient"], error_str: str) -> None:
logger.error(error_str)
await ctx.followup.send(content=error_str, ephemeral=True)


async def get_channel(
ctx: discord.interactions.Interaction["BotClient"],
channel_name: str,
) -> discord.TextChannel | None:
channel_name = channel_name.lower() # all text/voice channels are lowercase
guild = ctx.guild
if DISCORD_DEBUG:
# Always return calling channel
return cast(discord.TextChannel, ctx.channel)
if DISCORD_TESTING:
guild_id = os.getenv('DISCORD_GUILD')
if guild_id is None:
guild = None
else:
guild = ctx.client.get_guild(int(guild_id))

# get team's channel by name
if guild is None:
raise app_commands.NoPrivateMessage
channel = discord.utils.get(
guild.channels,
name=channel_name,
)

if not channel:
await log_and_reply(
ctx,
f"# Channel {channel_name} not found, unable to send message",
)
return None
elif not isinstance(channel, discord.TextChannel):
await log_and_reply(
ctx,
f"# {channel.name} is not a text channel, unable to send message",
)
return None

return channel


async def get_team_channel(
ctx: discord.interactions.Interaction["BotClient"],
archive_name: str,
zip_name: str,
) -> Tuple[str, discord.TextChannel | None]:
# extract team name from filename
tla_search = re.match(TEAM_CHANNEL_PREFIX + r'(.*?)[-.]', archive_name)
if not isinstance(tla_search, re.Match):
await log_and_reply(
ctx,
f"# Failed to extract a TLA from {archive_name} in {zip_name}",
)
return '', None

tla = tla_search.group(1)
channel = await get_channel(ctx, f"{TEAM_CHANNEL_PREFIX}{tla}")

return tla, channel


def pre_test_zipfile(archive_name: str, zip_name: str) -> bool:
if not archive_name.lower().endswith('.zip'): # skip non-zips
logger.debug(f"{archive_name} from {zip_name} is not a ZIP, skipping")
return False

# skip files not starting with TEAM_CHANNEL_PREFIX
if not archive_name.lower().startswith(TEAM_CHANNEL_PREFIX):
logger.debug(
f"{archive_name} from {zip_name} "
f"doesn't start with {TEAM_CHANNEL_PREFIX}, skipping",
)
return False
return True


def match_animation_files(log_name: str, animation_dir: Path) -> List[Path]:
match_num_search = re.search(r'match-([0-9]+)', log_name)
if not isinstance(match_num_search, re.Match):
logger.warning(f'Invalid match name: {log_name}')
return []
match_num = match_num_search[1]
logger.debug(f"Fetching animation files for match {match_num}")
match_files = animation_dir.glob(f'match-{match_num}.*')
return [data_file for data_file in match_files if data_file.suffix != '.mp4']


def insert_match_files(archive: Path, animation_dir: Path) -> None:
# append animations to archive
with ZipFile(archive, 'a', compression=ZIP_DEFLATED) as zipfile:
for log_name in zipfile.namelist():
if not log_name.endswith('.txt'):
continue

for animation_file in match_animation_files(log_name, animation_dir):
zipfile.write(animation_file.resolve(), animation_file.name)

# add textures subtree
for texture in (animation_dir / 'textures').glob('**/*'):
zipfile.write(
texture.resolve(),
texture.relative_to(animation_dir),
)


async def send_file(
ctx: discord.interactions.Interaction["BotClient"],
channel: discord.TextChannel,
archive: Path,
event_name: str,
msg_str: str = "Here are your logs",
logging_str: str = "Uploaded logs",
) -> bool:
try:
if DISCORD_TESTING: # don't actually send message in testing
if (archive.stat().st_size / 1000 ** 2) > 8:
# discord.HTTPException requires aiohttp.ClientResponse
await log_and_reply(
ctx,
f"# {archive.name} was too large to upload at "
f"{archive.stat().st_size / 1000 ** 2 :.3f} MiB",
)
return False
else:
await channel.send(
content=f"{msg_str} from {event_name if event_name else 'today'}",
file=discord.File(str(archive)),
)
logger.debug(
f"{logging_str} from {event_name if event_name else 'today'}",
)
except discord.HTTPException as e: # handle file size issues
if e.status == 413:
await log_and_reply(
ctx,
f"# {archive.name} was too large to upload at "
f"{archive.stat().st_size / 1000 ** 2 :.3f} MiB",
)
return False
else:
raise e
return True


def extract_animations(zipfile: ZipFile, tmpdir: Path, fully_extract: bool) -> bool:
animation_files = [
name for name in zipfile.namelist()
if name.split('/')[-1].startswith('animations')
and name.endswith('.zip')
]

if not animation_files:
return False

try:
zipfile.extract(animation_files[0], path=tmpdir)
except BadZipFile:
logger.warning("The animations zip was corrupt")
return False

# give the animations archive + folder if fixed name
shutil.move(str(tmpdir / animation_files[0]), str(tmpdir / 'animations.zip'))

if fully_extract:
with ZipFile(tmpdir / 'animations.zip') as animation_zip:
(tmpdir / 'animations').mkdir()
animation_zip.extractall(tmpdir / 'animations')
logger.debug("Extracting animations.zip")
return True


async def logs_upload(
ctx: discord.interactions.Interaction["BotClient"],
file: IO[bytes],
zip_name: str,
event_name: str,
team_animation: AnimationHandling, # None = don't upload animations
) -> None:
animations_found = False
try:
with tempfile.TemporaryDirectory() as tmpdir_name:
tmpdir = Path(tmpdir_name)
completed_tlas = []

with ZipFile(file) as zipfile:
if team_animation != AnimationHandling.none:
animations_found = extract_animations(zipfile, tmpdir, team_animation == AnimationHandling.team)

if not animations_found:
await log_and_reply(ctx, "animations Zip file is missing")

for archive_name in zipfile.namelist():
if not pre_test_zipfile(archive_name, zip_name):
continue

zipfile.extract(archive_name, path=tmpdir)

if not is_zipfile(tmpdir / archive_name): # test file is a valid zip
await log_and_reply(
ctx,
f"# {archive_name} from {zip_name} is not a valid ZIP file",
)
# The file will be removed with the temporary directory
continue

if team_animation == AnimationHandling.team and animations_found:
insert_match_files(tmpdir / archive_name, tmpdir / 'animations')

# get team's channel
tla, channel = await get_team_channel(ctx, archive_name, zip_name)
if not channel:
continue

# upload to team channel with message
if not await send_file(
ctx,
channel,
tmpdir / archive_name,
event_name,
logging_str=f"Uploaded logs for {tla}",
):
# try again without animations
# TODO test this clause in unit testing
if team_animation:
# extract original archive, modified version is overwritten
zipfile.extract(archive_name, path=tmpdir)

if await send_file( # retry with original archive
ctx,
channel,
tmpdir / archive_name,
event_name,
logging_str=f"Uploaded only logs for {tla}",
):
await log_and_reply(
ctx,
f"Only able to upload logs for {tla}, "
"no animations were served",
)

continue

completed_tlas.append(tla)

if team_animation == AnimationHandling.separate and animations_found:
common_channel = await get_channel(ctx, "general")
# upload animations.zip to common channel
if common_channel:
await send_file(
ctx,
common_channel,
tmpdir / 'animations.zip',
event_name,
msg_str="Here are the animation files",
logging_str="Uploaded animations",
)

await ctx.followup.send(content=
f"Successfully uploaded logs to {len(completed_tlas)} teams: "
f"{', '.join(completed_tlas)}",
)
except BadZipFile:
await log_and_reply(ctx, f"# {zip_name} is not a valid ZIP file")


@app_commands.command( # type:ignore[arg-type]
name="logs",
description="Get combined logs archive from URL for distribution to teams, avoids Discord's size limit",
)
@app_commands.describe(
url="URL to a zip of logs",
animations="How the animation files will be handled",
event_name="Optionally set the event name used in the bot's message to teams",
)
async def logs(
interaction: discord.interactions.Interaction['BotClient'],
url: str,
animations: AnimationHandling = AnimationHandling.none,
event_name: str | None = None,
) -> None:
logger.info(f"{interaction.user.name} started downloading logs from {url}")

with tempfile.TemporaryFile(suffix='.zip') as zipfile:
if url.endswith('.zip'):
filename = url.split("/")[-1]
else:
filename = f"logs_upload-{date.today()}.zip"

await interaction.response.defer(thinking=True) # provides feedback that the bot is processing
# download zip, using aiohttp
async with aiohttp.ClientSession() as session:
resp = await session.get(url)

if resp.status >= 400:
logger.error(
f"Download from {url} failed with error "
f"{resp.status}, {resp.reason}",
)
await interaction.followup.send(content="Zip file failed to download")
return

zipfile_data = await resp.read()

zipfile.write(zipfile_data)

# start processing from beginning of the file
zipfile.seek(0)

await logs_upload(
interaction,
zipfile,
filename,
event_name or "",
animations,
)
15 changes: 8 additions & 7 deletions src/commands/team.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
from typing import TYPE_CHECKING, Mapping
from typing import Mapping, TYPE_CHECKING

import discord
from discord import app_commands
@@ -12,6 +12,7 @@
ROLE_PREFIX,
TEAM_LEADER_ROLE,
TEAM_CATEGORY_NAME,
TEAM_CHANNEL_PREFIX,
PASSWORDS_CHANNEL_NAME,
TEAM_VOICE_CATEGORY_NAME,
)
@@ -77,7 +78,7 @@ async def new_team(interaction: discord.interactions.Interaction["BotClient"], t
)
channel = await guild.create_text_channel(
reason=TEAM_CREATED_REASON + interaction.user.name,
name=f"team-{tla.lower()}",
name=f"{TEAM_CHANNEL_PREFIX}{tla.lower()}",
topic=name,
category=category,
overwrites=permissions(interaction.client, role)
@@ -123,12 +124,12 @@ async def delete_team(interaction: discord.interactions.Interaction["BotClient"]
await member.kick(reason=reason)

for channel in guild.channels:
if channel.name.startswith(f"team-{tla.lower()}"):
if channel.name.startswith(f"{TEAM_CHANNEL_PREFIX}{tla.lower()}"):
await channel.delete(reason=reason)

await role.delete(reason=reason)

if isinstance(interaction.channel, discord.abc.GuildChannel) and not interaction.channel.name.startswith(f"team-{tla.lower()}"):
if isinstance(interaction.channel, discord.abc.GuildChannel) and not interaction.channel.name.startswith(f"{TEAM_CHANNEL_PREFIX}{tla.lower()}"):
await interaction.edit_original_response(content=f"Team {tla.upper()} has been deleted")
else:
await interaction.delete_original_response()
@@ -153,7 +154,7 @@ async def create_voice(interaction: discord.interactions.Interaction["BotClient"

category = discord.utils.get(guild.categories, name=TEAM_VOICE_CATEGORY_NAME)
channel = await guild.create_voice_channel(
f"team-{tla.lower()}",
f"{TEAM_CHANNEL_PREFIX}{tla.lower()}",
category=category,
overwrites=permissions(interaction.client, role)
)
@@ -182,15 +183,15 @@ async def create_team_channel(
await interaction.response.send_message("Team does not exist", ephemeral=True)
return

main_channel = discord.utils.get(guild.text_channels, name=f"team-{tla.lower()}")
main_channel = discord.utils.get(guild.text_channels, name=f"{TEAM_CHANNEL_PREFIX}{tla.lower()}")
category = discord.utils.get(guild.categories, name=TEAM_CATEGORY_NAME)

if category is None or main_channel is None:
await interaction.response.send_message(f"Team {tla.upper()} does not exist", ephemeral=True)
return

new_channel = await guild.create_text_channel(
name=f"team-{tla.lower()}-{suffix.lower()}",
name=f"{TEAM_CHANNEL_PREFIX}{tla.lower()}-{suffix.lower()}",
category=category,
overwrites=permissions(interaction.client, role),
position=main_channel.position + 1,
1 change: 1 addition & 0 deletions src/constants.py
Original file line number Diff line number Diff line change
@@ -21,6 +21,7 @@
PASSWORDS_CHANNEL_NAME = "role-passwords"

TEAM_CATEGORY_NAME = "Team Channels"
TEAM_CHANNEL_PREFIX = "team-"
TEAM_VOICE_CATEGORY_NAME = "Team Voice Channels"
TEAM_LEADER_ROLE = "Team Supervisor"

0 comments on commit 485afa7

Please sign in to comment.