synchronization.py

from dcim.models import Device, DeviceRole, DeviceType, Site
from ipam.models import Prefix
from dcim.choices import DeviceStatusChoices
from virtualization.choices import VirtualMachineStatusChoices
from virtualization.models import VirtualMachine
from extras.models import Tag
from rest_framework import serializers
import json
import logging
from core.models import Job
from .models import AWXInventory

logger = logging.getLogger(__name__)

group_prefixes = {Site: "site_", DeviceRole: "devicerole_", DeviceType: "devicetype_", Prefix: "prefix_", Tag:"tag_"}


class TagSerializer(serializers.BaseSerializer):
    def to_representation(self, instance):
        return {
            "name": "{}{}".format(group_prefixes[Tag], instance.slug.replace("-", "_")),
            "variables": json.dumps({"netbox_tag_name": instance.name}),
        }

class SiteSerializer(serializers.BaseSerializer):
    def to_representation(self, instance):
        return {
            "name": "{}{}".format(group_prefixes[Site], instance.slug.replace("-", "_")),
            "variables": json.dumps(
                {
                    "netbox_site_status": instance.status,
                }
            ),
        }


class DeviceRoleSerializer(serializers.BaseSerializer):
    def to_representation(self, instance):
        return {
            "name": "{}{}".format(group_prefixes[DeviceRole], instance.slug.replace("-", "_")),
            "variables": json.dumps({})
        }


class DeviceTypeSerializer(serializers.BaseSerializer):
    def to_representation(self, instance):
        return {
            "name": "{}{}".format(group_prefixes[DeviceType], instance.slug.replace("-", "_")),
            "description": instance.description,
            "variables": json.dumps(
                {
                    "netbox_devicetype_model": instance.model,
                }
            ),
        }


class PrefixSerializer(serializers.BaseSerializer):
    def to_representation(self, instance):
        return {
            "name": "{}{}".format(group_prefixes[Prefix], str(instance.prefix).replace(".", "_").replace("/", "_")),
            "variables": json.dumps({"netbox_prefix": "{}".format(str(instance.prefix))}),
        }


class InterfaceSerializer(serializers.BaseSerializer):
    def to_representation(self, instance):
        ip_addresses = []
        for ip in instance.ip_addresses.all():
            serializer = IPAddressSerializer(ip)
            ip_addresses.append(serializer.data)
        return {
            "name": instance.name,
            "mac": instance.mac_address,
            "ip_addresses": ip_addresses
        }


class IPAddressSerializer(serializers.BaseSerializer):
    def to_representation(self, instance):
        return {
            "address": str(instance.address),
            "dns_name": instance.dns_name
        }


class DeviceSerializer(serializers.BaseSerializer):
    def to_representation(self, instance):
        variables = {}
        variables["netbox_interfaces"] = []
        for interface in instance.interfaces.all():
            serializer = InterfaceSerializer(interface)
            variables["netbox_interfaces"].append(serializer.data)
        return {
            "name": getattr(instance.primary_ip4, 'dns_name', instance.name),
            "description": instance.description,
            "enabled": instance.status == DeviceStatusChoices.STATUS_ACTIVE,
            "variables": json.dumps(variables),
        }


class VMInterfaceSerializer(serializers.BaseSerializer):
    def to_representation(self, instance):
        ip_addresses = []
        for ip in instance.ip_addresses.all():
            serializer = IPAddressSerializer(ip)
            ip_addresses.append(serializer.data)
        return {
            "name": instance.name,
            "mac": instance.mac_address,
            "ip_addresses": ip_addresses
        }


class VMSerializer(serializers.BaseSerializer):
    def to_representation(self, instance):
        variables = {
            "netbox_virtualmachine_name": instance.name,
            "netbox_virtualmachine_vcpus": float(instance.vcpus) if instance.vcpus is not None else 0.0,
            "netbox_virtualmachine_memory": instance.memory or 0,
            "netbox_virtualmachine_disk": instance.disk or 0,
        }
        variables["netbox_interfaces"] = []
        for interface in instance.interfaces.all():
            serializer = VMInterfaceSerializer(interface)
            variables["netbox_interfaces"].append(serializer.data)
        return {
            "name": getattr(instance.primary_ip4, 'dns_name', instance.name),
            "description": instance.description,
            "enabled": instance.status == VirtualMachineStatusChoices.STATUS_ACTIVE,
            "variables": json.dumps(variables),
        }


class AWXHostSerializer(serializers.BaseSerializer):
    def to_internal_value(self, data):
        return {
            "name": data.get('name'),
            "description": data.get('description'),
            "enabled": data.get('enabled'),
            "variables": data.get('variables')
        }


class AWXGroupSerializer(serializers.BaseSerializer):
    def to_internal_value(self, data):
        return {
            "name": data.get('name'),
            "variables": data.get('variables')
        }


serializers = {
    Site: SiteSerializer,
    DeviceRole: DeviceRoleSerializer,
    DeviceType: DeviceTypeSerializer,
    Prefix: PrefixSerializer,
    Device: DeviceSerializer,
    VirtualMachine: VMSerializer,
    Tag: TagSerializer
}


def sync_host(inventory, sender, instance):
    serializer = serializers[sender](instance)
    host = inventory.get_host(serializer.data["name"])
    if host is None:
        # If the host doesn't exist, create it.
        inventory.create_host(serializer.data)
        host = inventory.get_host(serializer.data["name"])
    else:
        # If the host exists, update it if necessary.
        host_serializer = AWXHostSerializer(data=host)
        if not host_serializer.is_valid() or host_serializer.validated_data != serializer.data:
            inventory.update_host(host["id"], serializer.data)
    
    # Sync group associations
    current_groups = host["summary_fields"]["groups"]["results"]
    
    # Handle associations for each type (Site, DeviceRole, DeviceType, and Tags)
    if instance.site:
        sync_host_group_association(inventory, host, Site, instance.site, current_groups)
    if instance.role:
        sync_host_group_association(inventory, host, DeviceRole, instance.role, current_groups)
    if isinstance(instance, Device) and instance.device_type:
        sync_host_group_association(inventory, host, DeviceType, instance.device_type, current_groups)
    if tags := instance.tags.all():
        for tag in tags:
            sync_host_group_association(inventory, host, Tag, tag, current_groups)

    # Disassociate groups that are no longer relevant
    disassociate_removed_groups(inventory, host, instance, current_groups)


def sync_host_group_association(inventory, host, object_type, instance, current_groups):
    """
    Associates the host with the appropriate AWX group.
    """
    serializer = serializers[object_type](instance)
    group_name = serializer.data["name"]
    
    # Filter out the groups that match the current object_type (e.g., Site, Tag)
    relevant_groups = [
        group for group in current_groups if group["name"].startswith(group_prefixes[object_type])
    ]
    
    # Check if the host is already in the desired group
    if group_name not in [group["name"] for group in relevant_groups]:
        group = inventory.get_group(group_name)
        if group:
            inventory.associate_host_group(host["id"], group["id"])
            logger.info(f"Host {host['name']} associated with group {group_name}.")


def disassociate_removed_groups(inventory, host, instance, current_groups):
    """
    Disassociates the host from AWX groups that are no longer relevant.
    """
    # Collect the current groups that should remain (based on the instance's attributes)
    valid_group_names = set()
    
    if instance.site:
        valid_group_names.add(f"{group_prefixes[Site]}{instance.site.slug.replace('-', '_')}")
    if instance.role:
        valid_group_names.add(f"{group_prefixes[DeviceRole]}{instance.role.slug.replace('-', '_')}")
    if isinstance(instance, Device) and instance.device_type:
        valid_group_names.add(f"{group_prefixes[DeviceType]}{instance.device_type.slug.replace('-', '_')}")
    if tags := instance.tags.all():
        valid_group_names.update(
            f"{group_prefixes[Tag]}{tag.slug.replace('-', '_')}" for tag in tags
        )
    # Disassociate from groups that are no longer valid
    for group in current_groups:
        if group["name"] not in valid_group_names:
            inventory.disassociate_host_group(host["id"], group["id"])
            logger.info(f"Host {host['name']} disassociated from group {group['name']}.")

def delete_host(inventory, sender, instance):
    serializer = serializers[sender](instance)
    inventory.delete_host(serializer.data["name"])


def sync_group(inventory, sender, instance):
    serializer = serializers[sender](instance)
    group = inventory.get_group(serializer.data["name"])
    if group is None:
        inventory.create_group(serializer.data)
        group = inventory.get_group(serializer.data["name"])
    else:
        group_serializer = AWXGroupSerializer(data=group)
        if not group_serializer.is_valid() or group_serializer.validated_data != serializer.data:
            print("Group will be updated due to:")
            print(group_serializer.validated_data)
            print(serializer.data)
            inventory.update_group(serializer.data["name"], serializer.data)


def delete_group(inventory, sender, instance):
    serializer = serializers[sender](instance)
    inventory.delete_group(serializer.data["name"])


def sync_all(job):
    inventory = job.object
    logger.info("Performing full inventory sync for inventory {}".format(inventory.inventory_id))
    job.start()
    for site in Site.objects.all():
        sync_group(inventory, Site, site)
    for device_role in DeviceRole.objects.all():
        sync_group(inventory, DeviceRole, device_role)
    for device_type in DeviceType.objects.all():
        sync_group(inventory, DeviceType, device_type)
    for tag in Tag.objects.all():
        sync_group(inventory, Tag, tag)

    for device in Device.objects.all():
        if not device.primary_ip4 is None and device.primary_ip4.dns_name:
            sync_host(inventory, Device, device)
    for vm in VirtualMachine.objects.all():
        if not vm.primary_ip4 is None and vm.primary_ip4.dns_name:
            sync_host(inventory, VirtualMachine, vm)
    job.terminate()