diff options
Diffstat (limited to 'cogs/vm_management.py')
| -rw-r--r-- | cogs/vm_management.py | 185 |
1 files changed, 185 insertions, 0 deletions
diff --git a/cogs/vm_management.py b/cogs/vm_management.py new file mode 100644 index 0000000..fd76eaa --- /dev/null +++ b/cogs/vm_management.py @@ -0,0 +1,185 @@ +import discord +from discord import app_commands +from discord.ext import commands +from typing import Literal +import asyncio +from proxmoxer import ProxmoxAPI +import config + + +class VMManagement(commands.Cog): + """VM and LXC container management commands.""" + + def __init__(self, bot: commands.Bot): + self.bot = bot + + def get_proxmox(self) -> ProxmoxAPI: + """Get a Proxmox API connection.""" + return ProxmoxAPI( + config.PROXMOX_HOST, + user=config.PROXMOX_USER, + token_name=config.PROXMOX_TOKEN_NAME, + token_value=config.PROXMOX_TOKEN_VALUE, + verify_ssl=config.PROXMOX_VERIFY_SSL, + ) + + async def _check_owner(self, interaction: discord.Interaction) -> bool: + """Check if user is owner, send error if not.""" + if interaction.user.id != config.OWNER_ID: + await interaction.response.send_message( + "You don't have permission to use this command.", ephemeral=True + ) + return False + return True + + def _get_vm_type(self, proxmox: ProxmoxAPI, vmid: int) -> str | None: + """Determine if VMID is a QEMU VM or LXC container.""" + # Check QEMU VMs + try: + vms = proxmox.nodes(config.PROXMOX_NODE).qemu.get() + for vm in vms: + if vm.get("vmid") == vmid: + return "qemu" + except Exception: + pass + + # Check LXC containers + try: + lxcs = proxmox.nodes(config.PROXMOX_NODE).lxc.get() + for lxc in lxcs: + if lxc.get("vmid") == vmid: + return "lxc" + except Exception: + pass + + return None + + def _get_vm_name(self, proxmox: ProxmoxAPI, vmid: int, vm_type: str) -> str: + """Get the name of a VM or container.""" + try: + if vm_type == "qemu": + config_data = proxmox.nodes(config.PROXMOX_NODE).qemu(vmid).config.get() + else: + config_data = proxmox.nodes(config.PROXMOX_NODE).lxc(vmid).config.get() + return config_data.get("hostname", config_data.get("name", f"VM {vmid}")) + except Exception: + return f"VM {vmid}" + + vm_group = app_commands.Group( + name="vm", + description="Manage VMs and containers (owner only)", + allowed_contexts=app_commands.AppCommandContext(guild=True, dm_channel=True, private_channel=True), + allowed_installs=app_commands.AppInstallationType(guild=True, user=True), + ) + + @vm_group.command(name="start", description="Start a VM or container") + @app_commands.describe(vmid="The VMID to start") + async def start_vm(self, interaction: discord.Interaction, vmid: int): + if not await self._check_owner(interaction): + return + + await interaction.response.defer(ephemeral=True) + + try: + proxmox = self.get_proxmox() + vm_type = self._get_vm_type(proxmox, vmid) + + if not vm_type: + await interaction.followup.send(f"VM/Container {vmid} not found.", ephemeral=True) + return + + vm_name = self._get_vm_name(proxmox, vmid, vm_type) + + # Start the VM/container + if vm_type == "qemu": + await asyncio.to_thread( + lambda: proxmox.nodes(config.PROXMOX_NODE).qemu(vmid).status.start.post() + ) + else: + await asyncio.to_thread( + lambda: proxmox.nodes(config.PROXMOX_NODE).lxc(vmid).status.start.post() + ) + + type_name = "VM" if vm_type == "qemu" else "Container" + await interaction.followup.send( + f":green_circle: {type_name} **{vmid}** ({vm_name}) is starting...", ephemeral=True + ) + + except Exception as e: + await interaction.followup.send(f"Failed to start VM {vmid}: {e}", ephemeral=True) + + @vm_group.command(name="stop", description="Stop a VM or container") + @app_commands.describe(vmid="The VMID to stop") + async def stop_vm(self, interaction: discord.Interaction, vmid: int): + if not await self._check_owner(interaction): + return + + await interaction.response.defer(ephemeral=True) + + try: + proxmox = self.get_proxmox() + vm_type = self._get_vm_type(proxmox, vmid) + + if not vm_type: + await interaction.followup.send(f"VM/Container {vmid} not found.", ephemeral=True) + return + + vm_name = self._get_vm_name(proxmox, vmid, vm_type) + + # Stop the VM/container + if vm_type == "qemu": + await asyncio.to_thread( + lambda: proxmox.nodes(config.PROXMOX_NODE).qemu(vmid).status.stop.post() + ) + else: + await asyncio.to_thread( + lambda: proxmox.nodes(config.PROXMOX_NODE).lxc(vmid).status.stop.post() + ) + + type_name = "VM" if vm_type == "qemu" else "Container" + await interaction.followup.send( + f":red_circle: {type_name} **{vmid}** ({vm_name}) is stopping...", ephemeral=True + ) + + except Exception as e: + await interaction.followup.send(f"Failed to stop VM {vmid}: {e}", ephemeral=True) + + @vm_group.command(name="restart", description="Restart a VM or container") + @app_commands.describe(vmid="The VMID to restart") + async def restart_vm(self, interaction: discord.Interaction, vmid: int): + if not await self._check_owner(interaction): + return + + await interaction.response.defer(ephemeral=True) + + try: + proxmox = self.get_proxmox() + vm_type = self._get_vm_type(proxmox, vmid) + + if not vm_type: + await interaction.followup.send(f"VM/Container {vmid} not found.", ephemeral=True) + return + + vm_name = self._get_vm_name(proxmox, vmid, vm_type) + + # Reboot the VM/container + if vm_type == "qemu": + await asyncio.to_thread( + lambda: proxmox.nodes(config.PROXMOX_NODE).qemu(vmid).status.reboot.post() + ) + else: + await asyncio.to_thread( + lambda: proxmox.nodes(config.PROXMOX_NODE).lxc(vmid).status.reboot.post() + ) + + type_name = "VM" if vm_type == "qemu" else "Container" + await interaction.followup.send( + f":arrows_counterclockwise: {type_name} **{vmid}** ({vm_name}) is restarting...", ephemeral=True + ) + + except Exception as e: + await interaction.followup.send(f"Failed to restart VM {vmid}: {e}", ephemeral=True) + + +async def setup(bot: commands.Bot): + await bot.add_cog(VMManagement(bot)) |