k.eaven
bec06bd909
- Added multi-guild support with comma-separated AllowedGuildIDs config - Per-guild embed tracking with isolated state - Per-server action locking to prevent race conditions across guilds - Cross-guild embed updates after power actions - Debug logging for embed update triggers - Fixed resource usage field formatting (brackets now properly closed) - Updated config.ini.sample, generate_config.py for backward compatibility - Updated README.md with multi-guild documentation - Updated tests for multi-guild support
523 lines
21 KiB
Python
523 lines
21 KiB
Python
"""
|
||
Embed Management Module
|
||
|
||
Handles creation, tracking, and persistence of Discord server status embeds.
|
||
Manages embed locations storage and provides utilities for embed lifecycle.
|
||
Supports multi-guild deployments where each guild tracks servers independently.
|
||
"""
|
||
|
||
import json
|
||
import logging
|
||
from collections import defaultdict
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
from typing import Any, Dict, List, Optional, Tuple
|
||
|
||
import discord
|
||
|
||
from api import PanelAPI
|
||
from config import Config
|
||
from exceptions import EmbedNotFoundError
|
||
from metrics import ServerMetricsManager
|
||
from models import ServerData, ServerResources
|
||
|
||
logger = logging.getLogger("pterodisbot")
|
||
|
||
|
||
class EmbedManager:
|
||
"""
|
||
Manages server status embeds lifecycle across multiple guilds.
|
||
|
||
Handles:
|
||
- Embed creation with real-time server data
|
||
- Persistent storage of embed locations per guild (guild_id -> server_id -> location)
|
||
- Embed refresh and recreation per guild
|
||
- Cleanup of deleted/missing embeds
|
||
- Rate limit protection during bulk operations
|
||
- Migration from legacy single-guild storage format
|
||
|
||
Attributes:
|
||
config: Application configuration
|
||
api: Panel API client
|
||
metrics_manager: Metrics manager for graphs
|
||
storage_path: Path to embed locations JSON file
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
config: Config,
|
||
api: PanelAPI,
|
||
metrics_manager: ServerMetricsManager,
|
||
) -> None:
|
||
"""
|
||
Initialize the embed manager.
|
||
|
||
Args:
|
||
config: Validated application configuration
|
||
api: Initialized PanelAPI instance
|
||
metrics_manager: ServerMetricsManager instance
|
||
"""
|
||
self.config = config
|
||
self.api = api
|
||
self.metrics_manager = metrics_manager
|
||
self.storage_path = config.embed_storage_path
|
||
# Structure: {guild_id: {server_id: {"channel_id": str, "message_id": str}}}
|
||
self.embed_locations: Dict[int, Dict[str, Dict[str, str]]] = defaultdict(dict)
|
||
# Previous states tracking moved to bot.py for multi-guild support
|
||
|
||
# Ensure storage directory exists
|
||
self.storage_path.parent.mkdir(parents=True, exist_ok=True)
|
||
logger.info("Initialized EmbedManager with multi-guild support")
|
||
|
||
async def load_locations(self) -> None:
|
||
"""Load saved embed locations from JSON storage file with migration support."""
|
||
logger.debug("Loading embed locations from storage")
|
||
if not self.storage_path.exists():
|
||
logger.info("No existing embed locations file found")
|
||
return
|
||
|
||
try:
|
||
with open(self.storage_path, "r") as f:
|
||
data = json.load(f)
|
||
|
||
# Detect format:
|
||
# New format: {guild_id: {server_id: location}} where values are dicts with channel_id/message_id
|
||
# Old format: {server_id: location} where location has channel_id/message_id directly
|
||
is_new_format = False
|
||
if data:
|
||
first_key = next(iter(data.keys()))
|
||
first_value = data[first_key]
|
||
# New format: value is a dict of server locations
|
||
# Old format: value is a location dict with channel_id/message_id
|
||
if isinstance(first_value, dict) and "channel_id" not in first_value:
|
||
# Value is {server_id: {channel_id, message_id}} - new format
|
||
is_new_format = True
|
||
elif isinstance(first_value, dict) and "channel_id" in first_value:
|
||
# Value is {channel_id, message_id} - old format
|
||
is_new_format = False
|
||
# If first_key is numeric and >= 17 digits, also treat as new format
|
||
elif first_key.isdigit() and len(first_key) >= 17:
|
||
is_new_format = True
|
||
|
||
if is_new_format:
|
||
# New format: {guild_id: {server_id: location}}
|
||
self.embed_locations = defaultdict(dict)
|
||
for guild_id_str, servers in data.items():
|
||
guild_id = int(guild_id_str)
|
||
self.embed_locations[guild_id] = servers
|
||
total_embeds = sum(len(s) for s in self.embed_locations.values())
|
||
logger.info(f"Loaded embed locations for {len(self.embed_locations)} guilds ({total_embeds} total embeds)")
|
||
else:
|
||
# Legacy format: {server_id: location} - migrate to first allowed guild
|
||
legacy_guild_id = self.config.allowed_guild_ids[0] if self.config.allowed_guild_ids else 0
|
||
if legacy_guild_id:
|
||
self.embed_locations = defaultdict(dict)
|
||
self.embed_locations[legacy_guild_id] = data
|
||
logger.info(f"Migrated legacy embed locations to guild {legacy_guild_id}")
|
||
# Save in new format immediately
|
||
await self.save_locations()
|
||
else:
|
||
logger.warning("No allowed guilds configured, cannot migrate legacy embed locations")
|
||
|
||
except Exception as e:
|
||
logger.error(f"Failed to load embed locations: {str(e)}")
|
||
|
||
async def save_locations(self) -> None:
|
||
"""Save current embed locations to JSON storage file."""
|
||
logger.debug("Saving embed locations to storage")
|
||
try:
|
||
# Convert defaultdict to regular dict for JSON serialization
|
||
# Guild IDs as strings for JSON compatibility
|
||
data = {str(gid): servers for gid, servers in self.embed_locations.items()}
|
||
with open(self.storage_path, "w") as f:
|
||
json.dump(data, f, indent=2)
|
||
logger.debug("Successfully saved embed locations")
|
||
except Exception as e:
|
||
logger.error(f"Failed to save embed locations: {str(e)}")
|
||
|
||
def get_guild_locations(self, guild_id: int) -> Dict[str, Dict[str, str]]:
|
||
"""
|
||
Get all tracked embeds for a specific guild.
|
||
|
||
Args:
|
||
guild_id: Discord guild ID
|
||
|
||
Returns:
|
||
Dictionary mapping server_id to location info
|
||
"""
|
||
return self.embed_locations.get(guild_id, {})
|
||
|
||
def get_all_tracked_servers(self) -> List[Tuple[int, str]]:
|
||
"""
|
||
Get all (guild_id, server_id) pairs currently being tracked.
|
||
|
||
Returns:
|
||
List of (guild_id, server_id) tuples
|
||
"""
|
||
result = []
|
||
for guild_id, servers in self.embed_locations.items():
|
||
for server_id in servers:
|
||
result.append((guild_id, server_id))
|
||
return result
|
||
|
||
async def track_embed(self, guild_id: int, server_id: str, message: discord.Message) -> None:
|
||
"""
|
||
Track a newly created embed in storage for a specific guild.
|
||
|
||
Args:
|
||
guild_id: Discord guild ID
|
||
server_id: The server's panel identifier
|
||
message: Discord message containing the embed
|
||
"""
|
||
logger.debug(f"Tracking new embed for server {server_id} in guild {guild_id}, channel {message.channel.id}")
|
||
self.embed_locations[guild_id][server_id] = {
|
||
"channel_id": str(message.channel.id),
|
||
"message_id": str(message.id),
|
||
}
|
||
await self.save_locations()
|
||
|
||
async def untrack_embed(self, guild_id: int, server_id: str) -> None:
|
||
"""
|
||
Remove a server's embed from tracking for a specific guild.
|
||
|
||
Args:
|
||
guild_id: Discord guild ID
|
||
server_id: The server's panel identifier
|
||
"""
|
||
self.embed_locations[guild_id].pop(server_id, None)
|
||
# Clean up empty guild entries
|
||
if not self.embed_locations[guild_id]:
|
||
del self.embed_locations[guild_id]
|
||
await self.save_locations()
|
||
|
||
async def create_server_embed(
|
||
self, server_data: ServerData, resources: ServerResources, guild_id: int
|
||
) -> Tuple[discord.Embed, discord.ui.View]:
|
||
"""
|
||
Create a status embed and interactive view for a server.
|
||
|
||
Args:
|
||
server_data: Server information from panel
|
||
resources: Current resource usage data
|
||
guild_id: Discord guild ID this embed belongs to
|
||
|
||
Returns:
|
||
Tuple of (embed, view) ready for display
|
||
"""
|
||
attributes = server_data.attributes
|
||
identifier = attributes.identifier
|
||
name = attributes.name
|
||
description = attributes.description or "No description available"
|
||
|
||
logger.debug(f"Building status embed for server {name} ({identifier}) in guild {guild_id}")
|
||
|
||
# Parse resource data
|
||
current_state = resources.attributes.current_state.title()
|
||
is_suspended = attributes.suspended
|
||
|
||
# Create embed with appropriate color based on status
|
||
embed = discord.Embed(
|
||
title=f"{name} - {current_state}",
|
||
description=description,
|
||
color=(
|
||
discord.Color.blue()
|
||
if current_state.lower() == "running"
|
||
else discord.Color.red()
|
||
),
|
||
timestamp=datetime.now(),
|
||
)
|
||
|
||
embed.add_field(name="🆔 Server ID", value=f"`{identifier}`", inline=True)
|
||
|
||
if is_suspended:
|
||
embed.add_field(name="ℹ️ Status", value="⛔ `Suspended`", inline=True)
|
||
else:
|
||
embed.add_field(name="ℹ️ Status", value="✅ `Active`", inline=True)
|
||
|
||
# Add resource usage if server is running
|
||
if current_state.lower() != "offline":
|
||
embed = self._add_resource_fields(embed, resources, attributes)
|
||
embed = self._add_graph_fields(embed, identifier)
|
||
|
||
embed.set_footer(text="Last updated")
|
||
|
||
# Create interactive view with control buttons (pass guild_id for validation)
|
||
from views import ServerStatusView
|
||
|
||
view = ServerStatusView(
|
||
server_id=identifier,
|
||
server_name=name,
|
||
api=self.api,
|
||
server_data=server_data.model_dump(),
|
||
guild_id=guild_id,
|
||
)
|
||
|
||
logger.debug(f"Successfully built status components for {name} in guild {guild_id}")
|
||
return embed, view
|
||
|
||
def _add_resource_fields(
|
||
self, embed: discord.Embed, resources: ServerResources, attributes: Any
|
||
) -> discord.Embed:
|
||
"""Add resource usage fields to embed."""
|
||
resource_attrs = resources.attributes.resources
|
||
limits = attributes.limits
|
||
|
||
# Current usage
|
||
cpu_usage = round(resource_attrs.cpu_absolute, 2)
|
||
memory_usage = round(resource_attrs.memory_bytes / (1024**2), 2)
|
||
disk_usage = round(resource_attrs.disk_bytes / (1024**2), 2)
|
||
network_rx = round(resource_attrs.network_rx_bytes / (1024**2), 2)
|
||
network_tx = round(resource_attrs.network_tx_bytes / (1024**2), 2)
|
||
|
||
# Maximum allocated resources
|
||
cpu_limit = limits.cpu
|
||
memory_limit = limits.memory
|
||
disk_limit = limits.disk
|
||
|
||
# Format limit values - display ∞ for unlimited (0 limit)
|
||
def format_limit(value: int) -> str:
|
||
if value == 0:
|
||
return f"{'∞':>8}"
|
||
return f"{value:>8}"
|
||
|
||
# Format uptime
|
||
uptime_ms = resource_attrs.uptime
|
||
if uptime_ms > 0:
|
||
uptime_seconds = uptime_ms // 1000
|
||
if uptime_seconds < 60:
|
||
uptime_text = f"`{uptime_seconds}s`"
|
||
elif uptime_seconds < 3600:
|
||
uptime_text = f"`{uptime_seconds // 60}m {uptime_seconds % 60}s`"
|
||
elif uptime_seconds < 86400:
|
||
hours = uptime_seconds // 3600
|
||
minutes = (uptime_seconds % 3600) // 60
|
||
uptime_text = f"`{hours}h {minutes}m`"
|
||
else:
|
||
days = uptime_seconds // 86400
|
||
hours = (uptime_seconds % 86400) // 3600
|
||
uptime_text = f"`{days}d {hours}h`"
|
||
else:
|
||
uptime_text = "`Just started`"
|
||
|
||
embed.add_field(name="⏱️ Uptime", value=uptime_text, inline=True)
|
||
|
||
# Resource usage text box
|
||
usage_text = (
|
||
f"```properties\n"
|
||
f"CPU : [{cpu_usage:>8} / {format_limit(cpu_limit)}] %\n"
|
||
f"Memory : [{memory_usage:>8} / {format_limit(memory_limit)}] MiB\n"
|
||
f"Disk : [{disk_usage:>8} / {format_limit(disk_limit)}] MiB\n"
|
||
f"```"
|
||
)
|
||
|
||
embed.add_field(name="📊 Resource Usage", value=usage_text, inline=False)
|
||
embed.add_field(name="Network In", value=f"📥 `{network_rx} MiB`", inline=True)
|
||
embed.add_field(name="Network Out", value=f"📤 `{network_tx} MiB`", inline=True)
|
||
|
||
return embed
|
||
|
||
def _add_graph_fields(
|
||
self, embed: discord.Embed, server_id: str
|
||
) -> discord.Embed:
|
||
"""Add metrics graph fields to embed if data available."""
|
||
server_graphs = self.metrics_manager.get_server_graphs(server_id)
|
||
if server_graphs and server_graphs.has_sufficient_data:
|
||
summary = server_graphs.get_data_summary()
|
||
|
||
graph_description = (
|
||
f">>> `Data points: {summary['point_count']}/6`\n"
|
||
f"`CPU trend: {summary['cpu_trend']} • Memory trend: {summary['memory_trend']}`"
|
||
)
|
||
|
||
embed.add_field(
|
||
name="📈 Usage Trends (Last Minute)",
|
||
value=graph_description,
|
||
inline=False,
|
||
)
|
||
embed.set_image(url=f"attachment://metrics_graph_{server_id}.png")
|
||
|
||
return embed
|
||
|
||
async def refresh_all_embeds(self, guild_id: int) -> Tuple[int, int]:
|
||
"""
|
||
Perform a complete refresh of all server status embeds for a specific guild.
|
||
|
||
Creates new embeds and deletes old ones to prevent duplication.
|
||
|
||
Args:
|
||
guild_id: Discord guild ID to refresh embeds for
|
||
|
||
Returns:
|
||
Tuple of (deleted_count, created_count)
|
||
"""
|
||
logger.info(f"Starting full refresh of all server embeds for guild {guild_id}")
|
||
from bot import PterodactylBot # Avoid circular import
|
||
bot = PterodactylBot.get_instance()
|
||
|
||
deleted_count = 0
|
||
created_count = 0
|
||
skipped_count = 0
|
||
|
||
guild_locations = self.get_guild_locations(guild_id)
|
||
if not guild_locations:
|
||
logger.info(f"No embeds tracked for guild {guild_id}, nothing to refresh")
|
||
return 0, 0
|
||
|
||
try:
|
||
# Get current server list
|
||
servers = await self.api.get_servers()
|
||
if not servers:
|
||
logger.warning("No servers found in panel during refresh")
|
||
return 0, 0
|
||
|
||
server_cache = {s["attributes"]["identifier"]: s for s in servers}
|
||
|
||
# Create new embeds
|
||
new_embeds = {}
|
||
|
||
for server_id, server_data in server_cache.items():
|
||
if server_id not in guild_locations:
|
||
skipped_count += 1
|
||
continue
|
||
|
||
channel_id = guild_locations[server_id]["channel_id"]
|
||
channel = bot.get_channel(int(channel_id))
|
||
if not channel:
|
||
logger.warning(f"Channel {channel_id} not found for server {server_id} in guild {guild_id}")
|
||
continue
|
||
|
||
# Verify channel belongs to the correct guild
|
||
if channel.guild.id != guild_id:
|
||
logger.warning(f"Channel {channel_id} not in guild {guild_id}, skipping")
|
||
continue
|
||
|
||
try:
|
||
resources = await self.api.get_server_resources(server_id)
|
||
server_obj = ServerData.model_validate(server_data)
|
||
resources_obj = ServerResources.model_validate(resources)
|
||
|
||
embed, view = await self.create_server_embed(server_obj, resources_obj, guild_id)
|
||
message = await channel.send(embed=embed, view=view)
|
||
|
||
new_embeds[server_id] = {
|
||
"channel_id": str(channel.id),
|
||
"message_id": str(message.id),
|
||
}
|
||
created_count += 1
|
||
logger.info(f"Created new embed for server {server_data['attributes']['name']} in guild {guild_id}")
|
||
|
||
await asyncio.sleep(1) # Rate limit protection
|
||
except Exception as e:
|
||
logger.error(f"Failed to create new embed for server {server_id} in guild {guild_id}: {str(e)}")
|
||
|
||
if not new_embeds:
|
||
logger.warning(f"No new embeds created during refresh for guild {guild_id}")
|
||
return 0, 0
|
||
|
||
# Delete old embeds
|
||
for server_id, location in list(guild_locations.items()):
|
||
try:
|
||
channel = bot.get_channel(int(location["channel_id"]))
|
||
if channel:
|
||
# Verify channel belongs to the correct guild
|
||
if channel.guild.id != guild_id:
|
||
logger.warning(f"Channel {location['channel_id']} not in guild {guild_id}, skipping delete")
|
||
continue
|
||
try:
|
||
message = await channel.fetch_message(int(location["message_id"]))
|
||
await message.delete()
|
||
deleted_count += 1
|
||
logger.debug(f"Deleted old embed for server {server_id} in guild {guild_id}")
|
||
await asyncio.sleep(0.5)
|
||
except discord.NotFound:
|
||
logger.debug(f"Old embed for server {server_id} in guild {guild_id} already deleted")
|
||
except Exception as e:
|
||
logger.error(f"Failed to delete old embed for server {server_id} in guild {guild_id}: {str(e)}")
|
||
except Exception as e:
|
||
logger.error(f"Error processing old embed for server {server_id} in guild {guild_id}: {str(e)}")
|
||
|
||
logger.info(f"Guild {guild_id}: Deleted {deleted_count} old embeds, created {created_count} new embeds")
|
||
|
||
# Update storage for this guild only
|
||
self.embed_locations[guild_id] = new_embeds
|
||
await self.save_locations()
|
||
|
||
return deleted_count, created_count
|
||
|
||
except Exception as e:
|
||
logger.error(f"Critical error during embed refresh for guild {guild_id}: {str(e)}")
|
||
raise
|
||
|
||
async def update_single_embed(
|
||
self,
|
||
guild_id: int,
|
||
server_id: str,
|
||
location: Dict[str, str],
|
||
server_data: ServerData,
|
||
resources: ServerResources,
|
||
current_state: str,
|
||
cpu_usage: float,
|
||
) -> bool:
|
||
"""
|
||
Update a single server's status embed for a specific guild.
|
||
|
||
Args:
|
||
guild_id: Discord guild ID
|
||
server_id: Server identifier
|
||
location: Embed location dict with channel_id and message_id
|
||
server_data: Current server data
|
||
resources: Current resource data
|
||
current_state: Current power state
|
||
cpu_usage: Current CPU usage
|
||
|
||
Returns:
|
||
True if embed was updated, False otherwise
|
||
"""
|
||
from bot import PterodactylBot
|
||
bot = PterodactylBot.get_instance()
|
||
|
||
try:
|
||
embed, view = await self.create_server_embed(server_data, resources, guild_id)
|
||
|
||
channel = bot.get_channel(int(location["channel_id"]))
|
||
if not channel:
|
||
logger.warning(f"Channel {location['channel_id']} not found for server {server_id} in guild {guild_id}")
|
||
return False
|
||
|
||
# Verify channel belongs to the correct guild
|
||
if channel.guild.id != guild_id:
|
||
logger.warning(f"Channel {location['channel_id']} not in guild {guild_id}")
|
||
return False
|
||
|
||
message = await channel.fetch_message(int(location["message_id"]))
|
||
|
||
# Handle graph attachments
|
||
files = []
|
||
server_graphs = self.metrics_manager.get_server_graphs(server_id)
|
||
|
||
if current_state == "running" and server_graphs and server_graphs.has_sufficient_data:
|
||
combined_graph = server_graphs.generate_combined_graph()
|
||
if combined_graph:
|
||
files.append(
|
||
discord.File(
|
||
combined_graph,
|
||
filename=f"metrics_graph_{server_id}.png",
|
||
)
|
||
)
|
||
# Empty files list removes existing attachments
|
||
|
||
await message.edit(embed=embed, view=view, attachments=files)
|
||
logger.debug(f"Updated status for {server_data.attributes.name} in guild {guild_id}")
|
||
return True
|
||
|
||
except discord.NotFound:
|
||
logger.warning(f"Embed for server {server_id} in guild {guild_id} not found, removing from tracking")
|
||
await self.untrack_embed(guild_id, server_id)
|
||
return False
|
||
except Exception as e:
|
||
logger.error(f"Failed to update status for server {server_id} in guild {guild_id}: {str(e)}")
|
||
return False
|
||
|
||
|
||
# Import at bottom to avoid circular imports
|
||
import asyncio |