Started repo and finished the bot in one shot
This commit is contained in:
@@ -0,0 +1,169 @@
|
||||
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)
|
||||
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 '🔴'
|
||||
code = r['status_code'] if r['status_code'] is not None else '—'
|
||||
latency = f"{r['latency_ms']} ms" if r['latency_ms'] is not None else '—'
|
||||
value = f'{emoji} `HTTP {code}` · {latency}'
|
||||
if r.get('db') is not None:
|
||||
value += f' · db {r["db"]}'
|
||||
if r.get('version'):
|
||||
value += f' · v{r["version"]}'
|
||||
if r.get('detail'):
|
||||
value += f'\n*{r["detail"]}*'
|
||||
embed.add_field(name=r['name'], value=value, 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()
|
||||
Reference in New Issue
Block a user