NexaFabric/backend/app/services/providers/proxmox.py

from typing import Any

import httpx

from app.services.providers.base import Provider, ProviderConnection


class ProxmoxProvider(Provider):
    name = "proxmox"
    ignored_guest_interface_prefixes = (
        "br-",
        "cali",
        "cni",
        "docker",
        "flannel",
        "kube",
        "lo",
        "veth",
        "virbr",
    )

    def auth_header(self, token: str) -> str:
        token = token.strip()
        if token.startswith("PVEAPIToken="):
            return token
        return f"PVEAPIToken={token}"

    async def test_connection(self, connection: ProviderConnection) -> dict[str, Any]:
        headers = {"Authorization": self.auth_header(connection.token)}
        async with httpx.AsyncClient(verify=connection.verify_tls, timeout=10) as client:
            response = await client.get(
                f"{connection.api_url.rstrip('/')}/api2/json/version",
                headers=headers,
            )
            response.raise_for_status()
            return response.json().get("data", {})

    async def sync_inventory(self, connection: ProviderConnection) -> dict[str, list[dict[str, Any]]]:
        base_url = connection.api_url.rstrip("/")
        headers = {"Authorization": self.auth_header(connection.token)}
        async with httpx.AsyncClient(verify=connection.verify_tls, timeout=20) as client:
            resources = await client.get(
                f"{base_url}/api2/json/cluster/resources",
                headers=headers,
            )
            resources.raise_for_status()
            data = resources.json().get("data", [])

        nodes = [item for item in data if item.get("type") == "node"]
        workloads = [item for item in data if item.get("type") in {"qemu", "lxc"}]
        async with httpx.AsyncClient(verify=connection.verify_tls, timeout=8) as client:
            for workload in workloads:
                await self.enrich_workload_ips(client, base_url, headers, workload)
        networks = await self.list_networks(connection)
        return {"nodes": nodes, "workloads": workloads, "networks": networks}

    async def enrich_workload_ips(
        self,
        client: httpx.AsyncClient,
        base_url: str,
        headers: dict[str, str],
        workload: dict[str, Any],
    ) -> None:
        node = workload.get("node")
        vmid = workload.get("vmid")
        kind = workload.get("type")
        workload["ip_addresses"] = []
        if not node or not vmid:
            return

        if kind == "qemu":
            try:
                response = await client.get(
                    f"{base_url}/api2/json/nodes/{node}/qemu/{vmid}/agent/network-get-interfaces",
                    headers=headers,
                )
                if response.status_code >= 400:
                    return
                interfaces = response.json().get("data", {}).get("result", [])
                for interface in interfaces:
                    interface_name = str(interface.get("name") or "").lower()
                    if interface_name.startswith(self.ignored_guest_interface_prefixes):
                        continue
                    for address in interface.get("ip-addresses", []):
                        ip_address = address.get("ip-address")
                        prefix = address.get("prefix")
                        if ip_address and ":" not in ip_address and prefix is not None:
                            workload["ip_addresses"].append(f"{ip_address}/{prefix}")
            except httpx.HTTPError:
                return

        if kind == "lxc":
            try:
                response = await client.get(
                    f"{base_url}/api2/json/nodes/{node}/lxc/{vmid}/config",
                    headers=headers,
                )
                if response.status_code >= 400:
                    return
                config = response.json().get("data", {})
                for key, value in config.items():
                    if key.startswith("net") and isinstance(value, str):
                        for part in value.split(","):
                            if part.startswith("ip="):
                                ip_address = part.removeprefix("ip=")
                                if ip_address != "dhcp" and ":" not in ip_address:
                                    workload["ip_addresses"].append(ip_address)
            except httpx.HTTPError:
                return

    async def list_networks(self, connection: ProviderConnection) -> list[dict[str, Any]]:
        headers = {"Authorization": self.auth_header(connection.token)}
        async with httpx.AsyncClient(verify=connection.verify_tls, timeout=20) as client:
            resources = await client.get(
                f"{connection.api_url.rstrip('/')}/api2/json/cluster/resources",
                headers=headers,
            )
            resources.raise_for_status()
            data = resources.json().get("data", [])
        return [item for item in data if item.get("type") in {"network", "sdn"}]

    async def preview_rules(self, connection: ProviderConnection, rules: list[dict[str, Any]]) -> dict[str, Any]:
        return {
            "provider": self.name,
            "read_only": connection.read_only,
            "generated": rules,
            "warnings": ["Preview only. No Proxmox firewall changes were sent."],
        }

    def firewall_rules_url(self, connection: ProviderConnection, target: dict[str, Any]) -> str:
        kind = "lxc" if target.get("kind") == "lxc" else "qemu"
        node = target["node"]
        vmid = target["vmid"]
        return f"{connection.api_url.rstrip('/')}/api2/json/nodes/{node}/{kind}/{vmid}/firewall/rules"

    def policy_marker(self, rule: dict[str, Any]) -> str:
        return f"NexaFabric policy={rule.get('policy_id')}"

    async def delete_existing_policy_rules(
        self,
        client: httpx.AsyncClient,
        headers: dict[str, str],
        rules_url: str,
        marker: str,
    ) -> list[dict[str, Any]]:
        existing_response = await client.get(rules_url, headers=headers)
        existing_response.raise_for_status()
        existing_rules = existing_response.json().get("data", [])
        deletions = []
        for existing_rule in sorted(existing_rules, key=lambda item: int(item.get("pos", 0)), reverse=True):
            comment = str(existing_rule.get("comment") or "")
            pos = existing_rule.get("pos")
            if marker in comment and pos is not None:
                delete_response = await client.delete(f"{rules_url}/{pos}", headers=headers)
                delete_response.raise_for_status()
                deletions.append({"pos": pos, "comment": comment})
        return deletions

    async def apply_rules(self, connection: ProviderConnection, rules: list[dict[str, Any]]) -> dict[str, Any]:
        if connection.read_only:
            return {"applied": False, "reason": "Cluster is read-only", "rules": rules}
        missing_mapping = [rule for rule in rules if "provider_target" not in rule]
        if missing_mapping:
            return {
                "applied": False,
                "reason": "Live apply requires every rule to resolve to a concrete Proxmox VM/LXC target.",
                "rules": missing_mapping,
            }

        headers = {"Authorization": self.auth_header(connection.token)}
        grouped: dict[tuple[str, str, str, str], list[dict[str, Any]]] = {}
        for rule in rules:
            target = rule["provider_target"]
            marker = self.policy_marker(rule)
            key = (str(target["node"]), str(target["kind"]), str(target["vmid"]), marker)
            grouped.setdefault(key, []).append(rule)

        applied_rules = []
        deleted_rules = []
        audit_only_rules = []
        async with httpx.AsyncClient(verify=connection.verify_tls, timeout=20) as client:
            for target_rules in grouped.values():
                target = target_rules[0]["provider_target"]
                rules_url = self.firewall_rules_url(connection, target)
                marker = self.policy_marker(target_rules[0])
                deleted_rules.extend(await self.delete_existing_policy_rules(client, headers, rules_url, marker))

                for rule in target_rules:
                    if rule.get("audit_only"):
                        audit_only_rules.append(
                            {
                                "policy_id": rule.get("policy_id"),
                                "target": target,
                                "reason": "Audit mode does not enforce or write blocking Proxmox rules.",
                            }
                        )
                        continue
                    provider_rule = rule.get("provider_rule")
                    if not provider_rule:
                        return {
                            "applied": False,
                            "reason": "Resolved rule is missing provider_rule payload.",
                            "rule": rule,
                        }
                    create_response = await client.post(rules_url, headers=headers, data=provider_rule)
                    create_response.raise_for_status()
                    applied_rules.append({"target": target, "rule": provider_rule, "result": create_response.json().get("data")})

        return {
            "applied": True,
            "rules_written": len(applied_rules),
            "rules_deleted": len(deleted_rules),
            "audit_only": audit_only_rules,
            "rules": applied_rules,
        }