HealthCheckBot/main.py

import asyncio
import json
import logging

import aiohttp
import discord
from discord import app_commands

import checks
import config

log = logging.getLogger('healthbot')


class HealthBot(discord.Client):
    def __init__(self):
        intents = discord.Intents.default()
        super().__init__(intents=intents)
        self.tree = app_commands.CommandTree(self)
        self.session: aiohttp.ClientSession | None = None
        self.status_messages: list[dict] = []
        self._monitor_task: asyncio.Task | None = None

    def _load_state(self):
        try:
            with open(config.STATE_FILE, 'r') as f:
                data = json.load(f)
            self.status_messages = [
                {'channel_id': int(m['channel_id']), 'message_id': int(m['message_id'])}
                for m in data.get('messages', [])
            ]
        except (OSError, ValueError, TypeError, KeyError):
            self.status_messages = []

    def _save_state(self):
        try:
            with open(config.STATE_FILE, 'w') as f:
                json.dump({'messages': self.status_messages}, f)
        except OSError:
            log.warning('Could not save state file %s', config.STATE_FILE)

    def register_status_message(self, channel_id, message_id):
        if any(m['message_id'] == message_id for m in self.status_messages):
            return
        self.status_messages.append({'channel_id': channel_id, 'message_id': message_id})
        self._save_state()

    def forget_status_message(self, message_id):
        self.status_messages = [
            m for m in self.status_messages if m['message_id'] != message_id
        ]
        self._save_state()

    async def _ensure_session(self):
        if self.session is None or self.session.closed:
            self.session = aiohttp.ClientSession()
        return self.session

    async def run_checks(self):
        session = await self._ensure_session()
        return await checks.run_checks(session, config.CHECK_TIMEOUT)

    def build_embed(self, results):
        all_up = all(r['up'] for r in results if r.get('critical', True))
        color = discord.Color.green() if all_up else discord.Color.red()
        embed = discord.Embed(
            title='Service Health',
            color=color,
            timestamp=discord.utils.utcnow(),
        )
        for r in results:
            emoji = '🟢' if r['up'] else '🔴'
            latency = f"{r['latency_ms']} ms" if r['latency_ms'] is not None else '—'
            if r.get('kind') == 'minecraft':
                players = r.get('players') or '—'
                hc = f'{emoji} `{players} players` · {latency}'
                if r.get('version'):
                    hc += f' · v{r["version"]}'
                if r.get('motd'):
                    hc += f' · *{r["motd"]}*'
            else:
                code = r['status_code'] if r['status_code'] is not None else '—'
                hc = f'{emoji} `HTTP {code}` · {latency}'
                if r.get('db') is not None:
                    hc += f' · db {r["db"]}'
                if r.get('version'):
                    hc += f' · v{r["version"]}'
            if r.get('detail'):
                hc += f' · *{r["detail"]}*'

            lines = [f'**healthcheck** {hc}']
            for label in ('Cloudflare', 'Google'):
                d = r.get('dns', {}).get(label)
                if d and d.get('ok'):
                    val = '✓ ' + ', '.join(d['ips'])
                else:
                    err = d.get('error') if d else 'unknown'
                    val = f'✗ {err}'
                lines.append(f'**DNS resolution ({label})** {val}')

            embed.add_field(name=r['name'], value='\n'.join(lines), inline=False)
        embed.set_footer(text=f'Updates every {config.CHECK_INTERVAL}s')
        return embed

    async def update_embeds(self):
        if not self.status_messages:
            return
        results = await self.run_checks()
        embed = self.build_embed(results)

        kept = []
        for entry in self.status_messages:
            channel = self.get_channel(entry['channel_id'])
            if channel is None:
                try:
                    channel = await self.fetch_channel(entry['channel_id'])
                except (discord.NotFound, discord.Forbidden, discord.HTTPException):
                    continue
            try:
                message = await channel.fetch_message(entry['message_id'])
                await message.edit(embed=embed)
                kept.append(entry)
            except (discord.NotFound, discord.Forbidden, discord.HTTPException):
                continue

        self.status_messages = kept
        self._save_state()

    async def _monitor_loop(self):
        await self.wait_until_ready()
        log.info('Monitor loop started (interval %ss)', config.CHECK_INTERVAL)
        while not self.is_closed():
            try:
                await self.update_embeds()
            except Exception as exc:
                log.exception('Monitor tick failed: %s', exc)
            await asyncio.sleep(config.CHECK_INTERVAL)

    async def on_ready(self):
        log.info('Logged in as %s (id %s)', self.user, self.user.id)
        self._load_state()
        await self.tree.sync()
        log.info('Synced slash commands; tracking %d status message(s)', len(self.status_messages))
        if self._monitor_task is None or self._monitor_task.done():
            self._monitor_task = self.loop.create_task(self._monitor_loop())

    async def close(self):
        if self.session is not None and not self.session.closed:
            await self.session.close()
        await super().close()


bot = HealthBot()


@bot.tree.command(name='status', description='Post a live-updating health embed for the services')
async def status(interaction: discord.Interaction):
    await interaction.response.defer()
    results = await bot.run_checks()
    embed = bot.build_embed(results)

    existing = next(
        (m for m in bot.status_messages if m['channel_id'] == interaction.channel_id),
        None,
    )
    if existing:
        channel = interaction.channel
        try:
            message = await channel.fetch_message(existing['message_id'])
            await message.edit(embed=embed)
            await interaction.followup.send(
                'Refreshed the existing status embed.', ephemeral=True,
            )
            return
        except (discord.NotFound, discord.Forbidden, discord.HTTPException):
            bot.forget_status_message(existing['message_id'])

    message = await interaction.followup.send(embed=embed)
    bot.register_status_message(interaction.channel_id, message.id)


def main():
    logging.basicConfig(level=logging.INFO, format='%(asctime)s %(name)s %(levelname)s %(message)s')
    bot.run(config.DISCORD_TOKEN)


if __name__ == '__main__':
    main()