Spiffo/cogs/console.py

"""
Reads and handles updates in the server console file.
"""

import re
import logging
from pathlib import Path
from dataclasses import dataclass

import httpx
import aiofiles
from discord import Embed, Colour
from discord.ext import commands, tasks

CONSOLE_FILE_PATH = Path(__file__).parent.parent / "data" / "server-console.txt"
log = logging.getLogger(__name__)


@dataclass(slots=True)
class ZomboidUser:
    guid: str
    ip: str
    steam_id: str
    access: str
    username: str
    connection_type: str

    @property
    def steam_url(self):
        return f"https://steamcommunity.com/profiles/{self.steam_id}"

    async def get_steam_profile_picture(self, steam_api_key: str):
        if not steam_api_key:
            log.warning("No steam API key, can't get profile picture.")
            return

        async with httpx.AsyncClient() as client:
            response = await client.get(url=f"https://api.steampowered.com/ISteamUser/GetPlayerSummaries/v2/?key={steam_api_key}&steamids={self.steam_id}")
            response.raise_for_status()

            data = response.json()
            return data["response"]["players"][0]["avatar"]


class ConsoleCog(commands.Cog):
    """
    Reads and handles the server-console.txt file.
    """
    _last_line_number = 0

    def __init__(self, bot: commands.Bot):
        self.bot = bot
        self.monitor_console.start()

    @tasks.loop(seconds=1)
    async def monitor_console(self):
        """
        Check the latest version of the console log file.
        """

        if not CONSOLE_FILE_PATH.exists():
            self.monitor_console.cancel()
            raise FileNotFoundError("Server console file doesn't exist, task cancelled.")

        async with aiofiles.open(CONSOLE_FILE_PATH, "r", encoding="utf-8") as file:

            # If we are at 0, restarting the bot would cause rapid fire of all log lines,
            # instead lets grab the latest line, to prevent spam.
            if self._last_line_number == 0:
                await file.seek(0, 2)
                self._last_line_number = await file.tell()
                return

            await file.seek(self._last_line_number)
            lines = await file.readlines()
            if not lines:
                log.debug("no new lines to read")
                return

            for line in lines:
                await self.process_console_line(line.strip())

            self._last_line_number = await file.tell()

    async def process_console_line(self, line: str):
        """
        Determine how to handle the given line from the server console.        
        """

        if "CheckModsNeedUpdate" in line:
            await self.handle_mod_needs_update(line)

        elif "ConnectionManager: [fully-connected]" in line:
            await self.handle_player_joined(line)

        elif "ConnectionManager: [disconnect]" in line:
            await self.handle_player_left(line)

    async def handle_mod_needs_update(self, line: str):
        log.debug("mod update instruction: %s", line)

    async def handle_player_joined(self, line: str):
        """
        Report when a user has joined the server into a specified Discord channel.
        Example of line:
            ConnectionManager: [fully-connected] "" connection: guid=*** ip=*** steam-id=*** access=admin username="corbz" connection-type="UDPRakNet"
        """
        re_pattern = r"guid=(\d+) ip=([\d\.]+) steam-id=(\d+) access=(\w+) username=\"([^\"]+)\" connection-type=\"([^\"]+)\""
        re_match = re.search(re_pattern, line)
        if not re_match:
            log.warning("failed to parse player data: %s", line)
            return

        user = ZomboidUser(
            guid=re_match.group(1),
            ip=re_match.group(2),
            steam_id=re_match.group(3),
            access=re_match.group(4),
            username=re_match.group(5),
            connection_type=re_match.group(6),
        )

        channel = self.bot.get_channel(self.bot.in_game_channel_id)
        channel = await self.bot.fetch_channel(self.bot.in_game_channel_id) if not channel else channel

        embed = Embed(
            title=user.username,
            url=user.steam_url,
            description="Player has joined the server",
            colour=Colour.brand_green()
        )
        embed.set_thumbnail(url=user.get_steam_profile_picture(self.bot.steam_api_key))

        await channel.send(embed=embed)

    async def handle_player_left(self, line: str):
        """
        Report when a user has left the server into a specified Discord channel.
        Example of line:
            ConnectionManager: [disconnect] "receive-disconnect" connection: guid=*** ip=*** steam-id=*** access=admin username="corbz" connection-type="Disconnected"
        """
        re_pattern = r"guid=(\d+) ip=([\d\.]+) steam-id=(\d+) access=(\w+) username=\"([^\"]+)\" connection-type=\"([^\"]+)\""
        re_match = re.search(re_pattern, line)
        if not re_match:
            log.warning("failed to parse player data: %s", line)
            return

        user = ZomboidUser(
            guid=re_match.group(1),
            ip=re_match.group(2),
            steam_id=re_match.group(3),
            access=re_match.group(4),
            username=re_match.group(5),
            connection_type=re_match.group(6),
        )

        channel = self.bot.get_channel(self.bot.in_game_channel_id)
        channel = await self.bot.fetch_channel(self.bot.in_game_channel_id) if not channel else channel

        embed = Embed(
            title=user.username,
            url=user.steam_url,
            description="Player has left the server",
            colour=Colour.brand_red()
        )
        embed.set_thumbnail(url=user.get_steam_profile_picture(self.bot.steam_api_key))

        await channel.send(embed=embed)


async def setup(bot: commands.Bot):
    cog = ConsoleCog(bot)
    await bot.add_cog(cog)
    log.info("Added %s cog", cog.__class__.__name__)