ai_collections/self_created/Tools/proxmox_management.py

"""
title: Proxmox API Access
author: Pakobbix
author_url: https://gitea.zephyre.one/Pakobbix/
version: 0.1.0
"""

#!/usr/bin/env python3

import asyncio
import requests
from typing import Callable, Any
from pydantic import BaseModel, Field


class EventEmitter:
    def __init__(self, event_emitter: Callable[[dict], Any] = None):
        self.event_emitter = event_emitter

    async def progress_update(self, description):
        await self.emit(description)

    async def error_update(self, description):
        await self.emit(description, "error", True)

    async def success_update(self, description):
        await self.emit(description, "success", True)

    async def emit(self, description="Unknown State", status="in_progress", done=False):
        if self.event_emitter:
            await self.event_emitter(
                {
                    "type": "status",
                    "data": {
                        "status": status,
                        "description": description,
                        "done": done,
                    },
                }
            )


class function_collection:
    def request_api(self, url, token, method="GET"):
        headers = {
            "Authorization": f"PVEAPIToken={token}",
            "Content-Type": "application/json",
        }
        if method.upper() == "GET":
            response = requests.get(url, headers=headers, verify=False)
        elif method.upper() == "POST":
            response = requests.post(url, headers=headers, verify=False)
        else:
            raise ValueError(f"Unsupported HTTP method: {method}")

        if response.text:
            return response.json()
        return None

    def _check_lxc_or_vm(self, url, token, id_or_name) -> str | None:
        """
        Check if the given ID corresponds to a logical container (LXC) or a virtual machine (VM).

        :param id_or_name: The ID or name of the resource to check.
        :return: "lxc" if it's an LXC, "vm" if it's a VM, or None if not found.
        """
        url = f"{url}/api2/json/cluster/resources"
        response = self.request_api(url, token, method="GET")

        if "data" in response:
            for item in response["data"]:
                # Only check items that are vms or lxcs
                if item.get("type") in ["qemu", "lxc"]:
                    if (
                        str(item.get("vmid")) == str(id_or_name)
                        or item.get("name", "").lower() == str(id_or_name).lower()
                    ):
                        information_list = {
                            "vmid": item["vmid"],
                            "name": item["name"],
                            "type": item["type"],
                        }
                        return information_list
        return None


class Tools:
    class Valves(BaseModel):
        proxmox_url: str = Field(
            "https://localhost:8006", description="The Proxmox API URL"
        )
        api_token: str = Field(
            "api_token_here", description="The API token for authentication"
        )

    def __init__(self):
        self.valves = self.Valves()
        self.function_call = function_collection()

    ###########################################################################
    ######################### Proxmox Node Management #########################
    ###########################################################################

    async def running_version_check(
        self, __event_emitter__: Callable[[dict], Any] = None
    ) -> str | None:
        """
        Check the Proxmox API version.

        :param event_emitter: Optional callable for emitting events.
        :return: The Proxmox API version or None if an error occurs.
        """
        emitter = EventEmitter(__event_emitter__)
        await emitter.progress_update("Checking Proxmox API version")

        url = f"{self.valves.proxmox_url}/api2/json/version"

        try:
            version_info = self.function_call.request_api(
                url, self.valves.api_token, method="GET"
            )
            await emitter.success_update("Proxmox API version retrieved successfully")
            final_output = str(version_info)
            return (
                final_output
                + "\nDas ist nur die momentan laufenende Version! Eventuell ist aber bereits eine neuere Version Verfügbar."
            )
        except Exception as e:
            await emitter.error_update(
                f"Error retrieving Proxmox API version: {str(e)}"
            )
            return None

    async def list_nodes(
        self, __event_emitter__: Callable[[dict], Any] = None
    ) -> str | None:
        """
        List all nodes in the Proxmox cluster.

        :param event_emitter: Optional callable for emitting events.
        :return: List of nodes or None if an error occurs.
        """
        emitter = EventEmitter(__event_emitter__)
        await emitter.progress_update("Starting node listing process")

        url = f"{self.valves.proxmox_url}/api2/json/nodes"

        try:
            nodes = self.function_call.request_api(
                url, self.valves.api_token, method="GET"
            )
            await emitter.success_update("Node listing completed successfully")
            return str(nodes)
        except Exception as e:
            await emitter.error_update(f"Error during node listing: {str(e)}")
            return None

    async def get_node_status(
        self, __event_emitter__: Callable[[dict], Any] = None
    ) -> str | None:
        """
        Get the status of a specific node in the Proxmox cluster.

        :param node_name: The name of the node to retrieve status for.
        :param event_emitter: Optional callable for emitting events.
        :return: The status of the node or None if an error occurs.
        """
        emitter = EventEmitter(__event_emitter__)
        await emitter.progress_update(f"Getting status for node: zephyre")

        url = f"{self.valves.proxmox_url}/api2/json/nodes/zephyre/status"

        try:
            status = self.function_call.request_api(
                url, self.valves.api_token, method="GET"
            )
            await emitter.success_update(f"Node status retrieved successfully")
            return str(status)
        except Exception as e:
            await emitter.error_update(f"Error retrieving node status: {str(e)}")
            return None

    async def node_apt_update(
        self, __event_emitter__: Callable[[dict], Any] = None
    ) -> str | None:
        """
        Update the APT packages on a specific node in the Proxmox cluster.

        :param node_name: The name of the node to update.
        :param event_emitter: Optional callable for emitting events.
        :return: The response from the API or None if an error occurs.
        """
        emitter = EventEmitter(__event_emitter__)
        await emitter.progress_update(f"Updating APT packages for node: zephyre")

        url = f"{self.valves.proxmox_url}/api2/json/nodes/zephyre/apt/update"

        try:
            response = self.function_call.request_api(
                url, self.valves.api_token, method="GET"
            )
            await emitter.success_update(
                f"Node zephyre APT packages updated successfully"
            )
            return str(response)
        except Exception as e:
            await emitter.error_update(
                f"Error updating node zephyre APT packages: {str(e)}"
            )
            return None

    async def proxmox_node_log(
        self, __event_emitter__: Callable[[dict], Any] = None
    ) -> str | None:
        """
        Get the log of a specific node in the Proxmox cluster.

        :param event_emitter: Optional callable for emitting events.
        :return: The log of the node or None if an error occurs.
        """
        emitter = EventEmitter(__event_emitter__)
        await emitter.progress_update(f"Getting log for node: zephyre")

        url = f"{self.valves.proxmox_url}/api2/json/nodes/zephyre/log"

        try:
            log = self.function_call.request_api(url, self.valves.api_token)
            await emitter.success_update(f"Node log retrieved successfully")
            return str(log)
        except Exception as e:
            await emitter.error_update(f"Error retrieving node log: {str(e)}")
            return None

    ###########################################################################
    ######################### Proxmox VM Management ##########################
    ###########################################################################

    async def list_all_vms(
        self, __event_emitter__: Callable[[dict], Any] = None
    ) -> str | None:
        """
        List all VMs in the Proxmox cluster.

        :param event_emitter: Optional callable for emitting events.
        :return: List of VMs or None if an error occurs.
        """
        emitter = EventEmitter(__event_emitter__)
        await emitter.progress_update("Starting VM listing process")

        url = f"{self.valves.proxmox_url}/api2/json/cluster/resources"

        try:
            vms = self.function_call.request_api(
                url, self.valves.api_token, method="GET"
            )
            await emitter.success_update("VM listing completed successfully")
            return str(vms)
        except Exception as e:
            await emitter.error_update(f"Error during VM listing: {str(e)}")
            return None

    async def restart_vm(
        self, vm_id_or_name: str, __event_emitter__: Callable[[dict], Any] = None
    ) -> str | None:
        """
        Restart a virtual machine in the Proxmox cluster.

        :param vm_id: The ID of the VM to restart.
        :param event_emitter: Optional callable for emitting events.
        :return: The response from the API or None if an error occurs.
        """
        emitter = EventEmitter(__event_emitter__)
        await emitter.progress_update(f"Restarting VM: {vm_id_or_name}")

        type = self.function_call._check_lxc_or_vm(
            self.valves.proxmox_url, self.valves.api_token, vm_id_or_name
        )
        url = f"{self.valves.proxmox_url}/api2/json/nodes/zephyre/{type['type']}/{type['vmid']}/status/restart"

        try:
            response = self.function_call.request_api(
                url, self.valves.api_token, method="POST"
            )
            await emitter.success_update(f"VM {vm_id_or_name} restarted successfully")
            return str(response)
        except Exception as e:
            await emitter.error_update(f"Error restarting VM {vm_id_or_name}: {str(e)}")
            return None

    async def shutdown_vm(
        self, vm_id: str, __event_emitter__: Callable[[dict], Any] = None
    ) -> str | None:
        """
        Shutdown a virtual machine in the Proxmox cluster.

        :param vm_id: The ID of the VM to shutdown.
        :param event_emitter: Optional callable for emitting events.
        :return: The response from the API or None if an error occurs.
        """
        emitter = EventEmitter(__event_emitter__)
        await emitter.progress_update(f"Shutting down VM: {vm_id}")

        type = self.function_call._check_lxc_or_vm(
            self.valves.proxmox_url, self.valves.api_token, vm_id
        )
        url = f"{self.valves.proxmox_url}/api2/json/nodes/zephyre/{type['type']}/{type['vmid']}/status/shutdown"

        try:
            response = self.function_call.request_api(
                url, self.valves.api_token, method="POST"
            )
            await emitter.success_update(f"VM {vm_id} shut down successfully")
            return str(response)
        except Exception as e:
            await emitter.error_update(f"Error shutting down VM {vm_id}: {str(e)}")
            return None

    async def start_vm(
        self, vm_id: str, __event_emitter__: Callable[[dict], Any] = None
    ) -> str | None:
        """
        Start a virtual machine in the Proxmox cluster.

        :param vm_id: The ID of the VM to start.
        :param event_emitter: Optional callable for emitting events.
        :return: The response from the API or None if an error occurs.
        """
        emitter = EventEmitter(__event_emitter__)
        await emitter.progress_update(f"Starting VM: {vm_id}")

        type = self.function_call._check_lxc_or_vm(
            self.valves.proxmox_url, self.valves.api_token, vm_id
        )
        url = f"{self.valves.proxmox_url}/api2/json/nodes/zephyre/{type['type']}/{type['vmid']}/status/start"

        try:
            response = self.function_call.request_api(
                url, self.valves.api_token, method="POST"
            )
            await emitter.success_update(f"VM {vm_id} started successfully")
            return str(response)
        except Exception as e:
            await emitter.error_update(f"Error starting VM {vm_id}: {str(e)}")
            return f"Error starting VM {vm_id}: {str(e)}"

    async def get_specific_vm_status(
        self, vm_id_or_name: str, __event_emitter__: Callable[[dict], Any] = None
    ) -> str | None:
        """
        Get the status of a specific virtual machine in the Proxmox cluster.

        :param vm_id: The ID of the VM to retrieve status for.
        :param event_emitter: Optional callable for emitting events.
        :return: The status of the VM or None if an error occurs.
        """
        emitter = EventEmitter(__event_emitter__)
        await emitter.progress_update(f"Getting status for VM: {vm_id_or_name}")

        type = self.function_call._check_lxc_or_vm(
            self.valves.proxmox_url, self.valves.api_token, vm_id_or_name
        )
        await emitter.progress_update(
            f"Getting status for {type['type']}: {type['name']}"
        )
        url = f"{self.valves.proxmox_url}/api2/json/nodes/zephyre/{type['type']}/{type['vmid']}/status/current"

        try:
            status = self.function_call.request_api(
                url, self.valves.api_token, method="GET"
            )
            await emitter.success_update(
                f"VM {vm_id_or_name} status retrieved successfully"
            )
            return str(status)
        except Exception as e:
            await emitter.error_update(
                f"Error retrieving VM {vm_id_or_name} status: {str(e)}"
            )
            return None

    async def get_specific_vm_configuration_and_hardware(
        self, vm_id_or_vm_name: str, __event_emitter__: Callable[[dict], Any] = None
    ) -> str | None:
        """
        Get detailed information about a specific virtual machine or container in the Proxmox cluster.

        :param vm_id: The ID of the VM or container to retrieve information for.
        :param event_emitter: Optional callable for emitting events.
        :return: The information of the VM/container or None if an error occurs.
        """
        emitter = EventEmitter(__event_emitter__)
        await emitter.progress_update(
            f"Getting information for resource: {vm_id_or_vm_name}"
        )

        try:
            type = self.function_call._check_lxc_or_vm(
                self.valves.proxmox_url, self.valves.api_token, vm_id_or_vm_name
            )
            url = f"{self.valves.proxmox_url}/api2/json/nodes/zephyre/{type['type']}/{type['vmid']}/config"

            info = self.function_call.request_api(
                url, self.valves.api_token, method="GET"
            )
            await emitter.success_update(
                f"Resource {vm_id_or_vm_name} information retrieved successfully"
            )
            return str(info)
        except Exception as e:
            await emitter.error_update(
                f"Error retrieving resource {vm_id_or_vm_name} information: {str(e)}"
            )
            return None