Use CiliumCIDRGroup to simplify Cilium network policies on AKS

#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.11"
# dependencies = [
#   "azure-identity>=1.17.0",
#   "azure-mgmt-network>=28.0.0",
#   "pyyaml>=6.0.2",
# ]
# ///

"""Generate CiliumCIDRGroup resources from Azure VNets and subnets.

This script always emits both:
- one CiliumCIDRGroup per VNet (all VNet address spaces)
- one CiliumCIDRGroup per subnet (subnet address prefix/prefixes)
"""

from __future__ import annotations

import argparse
import re
from typing import Any

import yaml
from azure.identity import DefaultAzureCredential
from azure.mgmt.network import NetworkManagementClient


def sanitize_name(raw: str) -> str:
    """Return a DNS-1123 compatible Kubernetes resource name."""
    sanitized = re.sub(r"[^a-z0-9-]+", "-", raw.lower()).strip("-")
    sanitized = re.sub(r"-{2,}", "-", sanitized)
    sanitized = sanitized[:63].strip("-")
    return sanitized or "cidrgroup"


def parse_resource_group_from_id(resource_id: str) -> str:
    """Extract the Azure resource group from a resource ID."""
    parts = resource_id.split("/")
    try:
        index = parts.index("resourceGroups")
        return parts[index + 1]
    except (ValueError, IndexError) as exc:
        raise ValueError(
            f"Could not parse resource group from id: {resource_id}"
        ) from exc


def make_ccg(
    *,
    name: str,
    subscription: str,
    resource_group: str,
    vnet: str,
    subnet: str | None,
    cidrs: list[str],
) -> dict[str, Any]:
    """Build a CiliumCIDRGroup manifest as a Python dictionary."""
    labels: dict[str, str] = {
        "source": "azure",
        "azure/subscription": subscription,
        "azure/resource-group": resource_group,
        "azure/vnet": vnet,
    }
    if subnet:
        labels["azure/subnet"] = subnet

    return {
        "apiVersion": "cilium.io/v2",
        "kind": "CiliumCIDRGroup",
        "metadata": {
            "name": name,
            "labels": labels,
        },
        "spec": {
            "externalCIDRs": cidrs,
        },
    }


def main() -> None:
    """Generate CiliumCIDRGroup manifests and print them as YAML."""
    parser = argparse.ArgumentParser(
        description="Generate CiliumCIDRGroup resources from Azure VNets and subnets."
    )
    parser.add_argument("--subscription", required=True, help="Azure subscription ID")
    parser.add_argument("--resource-group", help="Optional resource group filter")
    args = parser.parse_args()

    credential = DefaultAzureCredential()
    client = NetworkManagementClient(credential, args.subscription)

    vnets = (
        list(client.virtual_networks.list(args.resource_group))
        if args.resource_group
        else list(client.virtual_networks.list_all())
    )

    ccg_manifests: list[dict[str, Any]] = []

    for vnet in vnets:
        if not vnet.id or not vnet.name:
            continue

        resource_group = parse_resource_group_from_id(vnet.id)
        vnet_name = vnet.name

        vnet_cidrs = []
        if vnet.address_space and vnet.address_space.address_prefixes:
            vnet_cidrs = list(vnet.address_space.address_prefixes)

        if vnet_cidrs:
            ccg_manifests.append(
                make_ccg(
                    name=sanitize_name(f"azure-{vnet_name}"),
                    subscription=args.subscription,
                    resource_group=resource_group,
                    vnet=vnet_name,
                    subnet=None,
                    cidrs=vnet_cidrs,
                )
            )

        for subnet in client.subnets.list(resource_group, vnet_name):
            if not subnet.name:
                continue

            subnet_cidrs: list[str] = []
            if getattr(subnet, "address_prefixes", None):
                subnet_cidrs.extend(subnet.address_prefixes)
            elif getattr(subnet, "address_prefix", None):
                subnet_cidrs.append(subnet.address_prefix)

            if not subnet_cidrs:
                continue

            ccg_manifests.append(
                make_ccg(
                    name=sanitize_name(f"azure-{subnet.name}"),
                    subscription=args.subscription,
                    resource_group=resource_group,
                    vnet=vnet_name,
                    subnet=subnet.name,
                    cidrs=subnet_cidrs,
                )
            )

    print(yaml.safe_dump_all(ccg_manifests, sort_keys=False))


if __name__ == "__main__":
    main()