ipfabric_netbox 3.2.0__py3-none-any.whl → 3.2.2__py3-none-any.whl

ipfabric_netbox/utilities/ipfutils.py

@@ -1,13 +1,18 @@
 import json
 import logging
-import uuid
 from importlib import metadata
+from typing import Callable
+from typing import TYPE_CHECKING
 
+from core.choices import DataSourceStatusChoices
 from core.exceptions import SyncError
 from dcim.models import Device
+from dcim.models import Interface
+from dcim.models import MACAddress
 from django.conf import settings
+from django.core.exceptions import MultipleObjectsReturned
 from django.core.exceptions import ObjectDoesNotExist
-from django.shortcuts import get_object_or_404
+from django.db.models import Model
 from django.utils.text import slugify
 from django_tables2 import Column
 from ipfabric import IPFClient
@@ -16,11 +21,14 @@ from netbox.config import get_config
 from netutils.utils import jinja2_convenience_function
 
 from ..choices import IPFabricSourceTypeChoices
+from ..exceptions import SearchError
+from ..exceptions import SyncDataError
 from .nbutils import device_serial_max_length
 from .nbutils import order_devices
 from .nbutils import order_members
-from .nbutils import order_pn
-from .nbutils import order_vrf
+
+if TYPE_CHECKING:
+    from ..models import IPFabricBranch
 
 logger = logging.getLogger("ipfabric_netbox.utilities.ipf_utils")
 
@@ -126,19 +134,45 @@ class IPFabricSyncRunner(object):
         self.settings = settings
         self.transform_map = transform_map
         self.sync = sync
-        self.relationship_store = {}
-        self.siteUUID = {}
         if hasattr(self.sync, "logger"):
             self.logger = self.sync.logger
-        self.interface_count_total = 0
-        self.interface_count = 1
-        self.inventoryitem_count = 1
-        self.inventoryitem_count_total = 0
 
         if self.sync.snapshot_data.status != "loaded":
             raise SyncError("Snapshot not loaded in IP Fabric.")
 
-    def get_model_or_update(self, app, model, data, uuid=None):
+    @staticmethod
+    def handle_errors(func: Callable):
+        def wrapper(*args, **kwargs):
+            try:
+                return func(*args, **kwargs)
+            except Exception as err:
+                # Log the error to logger outside of job - console/file
+                logger.error(err, exc_info=True)
+                # Logging section for logs inside job - facing user
+                self = args[0]
+                if isinstance(err, SearchError):
+                    if self.settings.get(err.model):
+                        self.logger.log_failure(
+                            f"Aborting syncing {err.model} instance due to above error, please check your transform maps and/or existing data.",
+                            obj=self.sync,
+                        )
+                    else:
+                        self.logger.log_failure(
+                            f"Syncing {err.model} is disabled in settings, but hit above error trying to find the correct item. Please check your transform maps and/or existing data.",
+                            obj=self.sync,
+                        )
+                else:
+                    self.logger.log_failure(
+                        f"Syncing failed with: {err}. See above error for more details.",
+                        obj=self.sync,
+                    )
+                # Make sure the whole sync is failed when we encounter error
+                self.sync.status = DataSourceStatusChoices.FAILED
+                return None
+
+        return wrapper
+
+    def get_model_or_update(self, app, model, data):
         transform_map = self.transform_map.objects.filter(
             target_model__app_label=app, target_model__model=model
         ).first()
@@ -147,84 +181,57 @@ class IPFabricSyncRunner(object):
             raise SystemError(f"No transform map available for {app}: {model}")
 
         model_settings = self.settings.get(model, False)
-        object = None
-
-        if model_settings:
-            logger.info(f"Creating {model}")
-            object = transform_map.update_or_create_instance(
-                data=data,
-                uuid=uuid,
-                relationship_store=self.relationship_store,
-                tags=self.sync.tags.all(),
-                logger=self.logger,
-            )
-        else:
-            logger.info(f"Getting {model}")
-            coalesce_fields = transform_map.get_coalesce_fields(data)
-            object = get_object_or_404(
-                transform_map.target_model.model_class().objects.all(),
-                **coalesce_fields,
+        try:
+            context = transform_map.get_context(data)
+        except Exception as err:
+            message = f"Error getting context for `{model}`."
+            if isinstance(err, ObjectDoesNotExist):
+                message += " Could not find related object using on template in transform maps."
+            elif isinstance(err, MultipleObjectsReturned):
+                message += " Multiple objects returned using on template in transform maps, the template is not strict enough."
+            self.logger.log_failure(
+                f"<b>{message}</b><br/>data: `{data}`<br/>error: `{err}`", obj=self.sync
             )
+            raise SyncError(err) from err
 
-        store = self.relationship_store.get(uuid)
-
-        if store:
-            store[object._meta.model] = object
-        else:
-            self.relationship_store[uuid] = {object._meta.model: object}
+        queryset = transform_map.target_model.model_class().objects
 
-        return object
-
-    def create_interface(
-        self, device_interface, device_uuid, managed_ips, device_object, device
-    ):
-        device_interface["loginIp"] = device.get("loginIp")
-        interface_object = self.get_model_or_update(
-            "dcim", "interface", device_interface, uuid=device_uuid
-        )
-
-        self.logger.increment_statistics(
-            model="interface",
-            current=self.interface_count,
-            total=self.interface_count_total,
-        )
-        self.interface_count += 1
-
-        if self.settings.get("ipaddress"):
-            managed_ip = managed_ips.get(device_object.serial, {}).get(
-                interface_object.name
-            )
-            if managed_ip:
-                ip_address_obj = self.get_model_or_update(
-                    "ipam",
-                    "ipaddress",
-                    managed_ip,
+        object = None
+        try:
+            if model_settings:
+                logger.info(f"Creating {model}")
+                object = transform_map.update_or_create_instance(
+                    context=context,
+                    tags=self.sync.tags.all(),
                 )
-                try:
-                    other_device = Device.objects.get(primary_ip4=ip_address_obj)
-                    if other_device and device_object != other_device:
-                        other_device.primary_ip4 = None
-                        other_device.save()
-                except ObjectDoesNotExist:
-                    pass
-
-                if device.get("loginIp") == device_interface.get("primaryIp"):
-                    device_object.primary_ip4 = ip_address_obj
-                    device_object.save()
-
-        if self.settings.get("macaddress") and device_interface.get("mac"):
-            # Need to create MAC Address object before we can assign it to Interface
-            macaddress_data = {
-                "mac": device_interface.get("mac"),
-                "id": interface_object.id,
-            }
-            macaddress_object = self.get_model_or_update(
-                "dcim", "macaddress", macaddress_data, uuid=device_uuid
+            else:
+                logger.info(f"Getting {model}")
+                context.pop("defaults", None)
+                object = queryset.get(**context)
+        except queryset.model.DoesNotExist as err:
+            self.logger.log_failure(
+                f"<b>`{model}` with these keys not found: `{context}`.</b><br/>Original data: `{data}`.",
+                obj=self.sync,
             )
-            interface_object.primary_mac_address = macaddress_object
-            interface_object.save()
+            raise SearchError(model=model, context=context, data=data) from err
+        except queryset.model.MultipleObjectsReturned as err:
+            self.logger.log_failure(
+                f"<b>Multiple `{model}` with these keys found: `{context}`.</b><br/>Original data: `{data}`.",
+                obj=self.sync,
+            )
+            raise SearchError(model=model, context=context, data=data) from err
+        except Exception as err:
+            self.logger.log_failure(
+                f"""Failed to create instance of `{model}`:<br/>
+            message: `{err}`<br/>
+            raw data: `{data}`<br/>
+            context: `{context}`<br/>
+            """,
+                obj=self.sync,
+            )
+            raise SyncDataError(model=model, context=context, data=data) from err
 
-        return True
+        return object
 
     def collect_data(self):
         try:
@@ -295,7 +302,6 @@ class IPFabricSyncRunner(object):
                 query_filter = {
                     "and": [{"vendor": ["neq", vendor]} for vendor in excluded_vendors]
                 }
-                # filter = {"and": [{"vendor": ["neq", "aws"]}, {"vendor": ["neq", "azure"]}]}
 
                 if ingestion_sites := self.settings.get("sites"):
                     site_filter = {
@@ -304,13 +310,14 @@ class IPFabricSyncRunner(object):
                     query_filter["and"].append(site_filter)
 
                     self.logger.log_info(
-                        f"Creating site filter {json.dumps(site_filter)}", obj=self.sync
+                        f"Creating site filter `{json.dumps(site_filter)}`",
+                        obj=self.sync,
                     )
                 else:
                     site_filter = {}
 
                 data["site"] = self.client.inventory.sites.all(
-                    snapshot_id=self.settings["snapshot_id"]
+                    snapshot_id=self.settings["snapshot_id"], filters=site_filter
                 )
 
                 data["device"] = self.client.inventory.devices.all(
@@ -324,13 +331,17 @@ class IPFabricSyncRunner(object):
                 )
 
                 data["interface"] = self.client.inventory.interfaces.all(
-                    snapshot_id=self.settings["snapshot_id"]
+                    snapshot_id=self.settings["snapshot_id"], filters=site_filter
                 )
 
                 data["inventoryitem"] = self.client.inventory.pn.all(
                     snapshot_id=self.settings["snapshot_id"],
                     filters={
-                        "and": [{"sn": ["empty", False]}, {"name": ["empty", False]}]
+                        "and": [
+                            site_filter,
+                            {"sn": ["empty", False]},
+                            {"name": ["empty", False]},
+                        ]
                     },
                 )
 
@@ -356,7 +367,7 @@ class IPFabricSyncRunner(object):
                 data[
                     "ipaddress"
                 ] = self.client.technology.addressing.managed_ip_ipv4.all(
-                    snapshot_id=self.settings["snapshot_id"]
+                    snapshot_id=self.settings["snapshot_id"], filters=site_filter
                 )
         except Exception as e:
             self.logger.log_failure(
@@ -413,18 +424,7 @@ class IPFabricSyncRunner(object):
 
         self.logger.log_info("Ordering Part Numbers", obj=self.sync)
 
-        part_numbers = order_pn(data.get("inventoryitem", []))
-
-        self.logger.log_info("Ordering VRF's", obj=self.sync)
-
-        vrfs = order_vrf(data["vrf"])
-
-        managed_ips = {}
-        site_dict = {}
         interface_dict = {}
-        for site in data["site"]:
-            site_dict[site["siteName"]] = site
-
         for interface in data["interface"]:
             if int_sn := interface.get("sn"):
                 if interface_dict.get(int_sn):
@@ -448,6 +448,7 @@ class IPFabricSyncRunner(object):
             )
             raise SyncError(f"Error collecting source column name for interface: {e}")
 
+        managed_ips = {}
         for ip in data["ipaddress"]:
             # Find corresponding interface list by serial number (sn)
             device_interfaces = interface_dict.get(ip["sn"], [])
@@ -463,141 +464,259 @@ class IPFabricSyncRunner(object):
                     int_name = ip[interface_key]
                 else:
                     int_name = ip["intName"]
-
-                if managed_ips.get(ip["sn"]):
-                    managed_ips[ip["sn"]][int_name] = ip
+                if ip["sn"] not in managed_ips:
+                    managed_ips[ip["sn"]] = {int_name: [ip]}
+                elif int_name not in managed_ips.get(ip["sn"]):
+                    managed_ips[ip["sn"]][int_name] = [ip]
                 else:
-                    managed_ips[ip["sn"]] = {int_name: ip}
+                    managed_ips[ip["sn"]][int_name].append(ip)
+
+        for model, item_count in [
+            ("site", len(data.get("site", []))),
+            ("device", len(devices)),
+            ("interface", len(data.get("interface", []))),
+            ("inventoryitem", len(data.get("inventoryitem", []))),
+            ("vlan", len(data.get("vlan", []))),
+            ("vrf", len(data.get("vrf", []))),
+            ("prefix", len(data.get("prefix", []))),
+            # TODO: Since we sync only those assigned to interfaces, we are skipping some IPs
+            # TODO: This is fixable by syncing IPs separately from interface and only assign them on interfaces
+            ("ipaddress", len(data.get("ipaddress", []))),
+        ]:
+            if self.settings.get(model):
+                self.logger.init_statistics(model, item_count)
 
         return (
-            site_dict,
+            data["site"],
             devices,
             interface_dict,
-            part_numbers,
-            vrfs,
+            data["inventoryitem"],
+            data["vrf"],
             data["vlan"],
             data["prefix"],
             managed_ips,
         )
 
-    def sync_devices(self, branch=None):
-        self.logger.log_info("Starting device sync", obj=self.sync)
-
-        (
-            site_dict,
-            devices,
-            interface_dict,
-            part_numbers,
-            vrfs,
-            vlans,
-            networks,
-            managed_ips,
-        ) = self.collect_data()
-        vlan_count = 1
-        vrf_count = 1
-        network_count = 1
-        device_vrfs_total = 0
-
-        for device_count, device in enumerate(devices, start=1):
-            logger.info(f"Device {device_count} out of {len(devices)}")
-            self.logger.increment_statistics(
-                model="device", current=device_count, total=len(devices)
+    @handle_errors
+    def sync_model(self, app_label: str, model: str, data: dict | None) -> Model | None:
+        """Sync a single item to NetBox."""
+        if not data:
+            return None
+        return self.get_model_or_update(app_label, model, data)
+
+    def sync_items(
+        self,
+        items,
+        app_label: str,
+        model: str,
+        cf: bool = False,
+        branch: "IPFabricBranch" = None,
+    ) -> None:
+        """Sync list of items to NetBox."""
+        if not self.settings.get(model):
+            self.logger.log_info(
+                f"Did not ask to sync {model}s, skipping.", obj=self.sync
             )
+            return
+
+        for item in items:
+            synced_object = self.sync_model(app_label=app_label, model=model, data=item)
+            if synced_object is None:
+                continue
+            # Only log when we successfully synced the item
+            self.logger.increment_statistics(model=model)
+
+            if cf:
+                synced_object.custom_field_data[
+                    "ipfabric_source"
+                ] = self.sync.snapshot_data.source.pk
+                if branch:
+                    synced_object.custom_field_data["ipfabric_branch"] = branch.pk
+                synced_object.save()
+
+    @handle_errors
+    def sync_devices(
+        self,
+        branch,
+        devices,
+        interface_dict,
+        managed_ips,
+    ):
+        for model, name in [
+            ("manufacturer", "manufacturers"),
+            ("devicetype", "device types"),
+            ("platform", "platforms"),
+            ("devicerole", "device roles"),
+            ("virtualchassis", "virtual chassis"),
+            ("device", "devices"),
+            ("inventoryitem", "device inventory items"),
+        ]:
+            if not self.settings.get(model):
+                self.logger.log_info(
+                    f"Did not ask to sync {name}, skipping", obj=self.sync
+                )
 
-            device_uuid = str(uuid.uuid4())
-
-            site_object = self.get_model_or_update(
-                "dcim", "site", site_dict[device["siteName"]], uuid=device_uuid
+        devices_total = len(devices)
+
+        for device in devices:
+            if self.sync_model("dcim", "manufacturer", device) is None:
+                continue
+            if self.sync_model("dcim", "devicetype", device) is None:
+                continue
+            if self.sync_model("dcim", "platform", device) is None:
+                continue
+            if self.sync_model("dcim", "devicerole", device) is None:
+                continue
+
+            virtual_chassis = device.get("virtual_chassis", {})
+            self.sync_model("dcim", "virtualchassis", virtual_chassis)
+
+            if (device_object := self.sync_model("dcim", "device", device)) is None:
+                continue
+
+            if self.settings.get("device"):
+                device_object.custom_field_data[
+                    "ipfabric_source"
+                ] = self.sync.snapshot_data.source.pk
+                if branch:
+                    device_object.custom_field_data["ipfabric_branch"] = branch.pk
+                device_object.save()
+
+            self.logger.increment_statistics(model="device")
+            logger.info(
+                f"Device {self.logger.log_data.get('statistics', {}).get('device', {}).get('current')} out of {devices_total}"
             )
 
-            self.get_model_or_update("dcim", "manufacturer", device, uuid=device_uuid)
-            self.get_model_or_update("dcim", "devicetype", device, uuid=device_uuid)
+            # The Device exists now, so we can update master of the VC.
+            # The logic is handled in transform maps.
+            self.sync_model("dcim", "virtualchassis", virtual_chassis)
 
-            self.get_model_or_update("dcim", "platform", device, uuid=device_uuid)
-
-            self.get_model_or_update("dcim", "devicerole", device, uuid=device_uuid)
+            device_interfaces = interface_dict.get(device.get("sn"), [])
+            for device_interface in device_interfaces:
+                interface_object = self.sync_interface(
+                    device_interface, managed_ips, device_object, device
+                )
+                if interface_object is None:
+                    continue
+                self.logger.increment_statistics(model="interface")
+
+    @handle_errors
+    def sync_ipaddress(
+        self,
+        managed_ip: dict | None,
+        device_object: Device | None,
+        primary_ip: str | None,
+        login_ip: str | None,
+    ):
+        if not self.settings.get("ipaddress") or not managed_ip:
+            return
+        ip_address_obj = self.get_model_or_update(
+            "ipam",
+            "ipaddress",
+            managed_ip,
+        )
+        if ip_address_obj is None:
+            return
+        self.logger.increment_statistics(model="ipaddress")
 
-            device_object = self.get_model_or_update(
-                "dcim", "device", device, uuid=device_uuid
+        try:
+            # Removing other IP is done in .signals.clear_other_primary_ip
+            # But do it here too so the change is shown in StagedChange diff
+            other_device = Device.objects.get(primary_ip4=ip_address_obj)
+            if other_device and device_object != other_device:
+                other_device.primary_ip4 = None
+                other_device.save()
+        except ObjectDoesNotExist:
+            pass
+
+        if login_ip == primary_ip:
+            try:
+                device_object.primary_ip4 = ip_address_obj
+                device_object.save()
+            except ValueError as err:
+                self.logger.log_failure(
+                    f"Error assigning primary IP to device: {err}", obj=self.sync
+                )
+                return None
+        return ip_address_obj
+
+    @handle_errors
+    def sync_macaddress(
+        self, data: dict | None, interface_object: Interface
+    ) -> MACAddress | None:
+        if not self.settings.get("macaddress") or not data:
+            return None
+        # Need to create MAC Address object before we can assign it to Interface
+        # TODO: Figure out how to do this using transform maps
+        macaddress_data = {
+            "mac": data,
+            "id": interface_object.pk,
+        }
+        macaddress_object = self.get_model_or_update(
+            "dcim", "macaddress", macaddress_data
+        )
+        try:
+            interface_object.primary_mac_address = macaddress_object
+            interface_object.save()
+        except ValueError as err:
+            self.logger.log_failure(
+                f"Error assigning MAC Address to interface: {err}", obj=self.sync
+            )
+            return None
+        return macaddress_object
+
+    @handle_errors
+    def sync_interface(
+        self,
+        device_interface: dict,
+        managed_ips: dict,
+        device_object: Device | None,
+        device: dict,
+    ):
+        device_interface["loginIp"] = device.get("loginIp")
+        interface_object = self.get_model_or_update(
+            "dcim", "interface", device_interface
+        )
+        if interface_object is None:
+            return None
+
+        for ipaddress in managed_ips.get(device_object.serial, {}).get(
+            interface_object.name, []
+        ):
+            self.sync_ipaddress(
+                ipaddress,
+                device_object,
+                device_interface.get("primaryIp"),
+                device.get("loginIp"),
             )
 
-            site_object.custom_field_data[
-                "ipfabric_source"
-            ] = self.sync.snapshot_data.source.pk
-
-            device_object.custom_field_data[
-                "ipfabric_source"
-            ] = self.sync.snapshot_data.source.pk
-            if branch:
-                site_object.custom_field_data["ipfabric_branch"] = branch.pk
-                device_object.custom_field_data["ipfabric_branch"] = branch.pk
-
-            site_object.save()
-            device_object.save()
-
-            if self.settings.get("virtualchassis"):
-                if member := device.get("virtual_chassis"):
-                    self.get_model_or_update("dcim", "virtualchassis", member)
-                    device_object = self.get_model_or_update(
-                        "dcim", "device", device, uuid=device_uuid
-                    )
+        self.sync_macaddress(device_interface.get("mac"), interface_object)
 
-            if device_object and self.settings.get("interface"):
-                device_interfaces = interface_dict.get(device.get("sn"), [])
-                self.interface_count_total += len(device_interfaces)
-                for device_interface in device_interfaces:
-                    self.create_interface(
-                        device_interface,
-                        device_uuid,
-                        managed_ips,
-                        device_object,
-                        device,
-                    )
-                    # x = threading.Thread(target=self.create_interface, args=((device_interface, device_uuid, managed_ips, device_object, device)))
-                    # threads.append(x)
-                    # x.start()
-
-            if self.settings.get("vrf"):
-                device_vrfs = vrfs.get(device_object.serial, [])
-                device_vrfs_total += len(device_vrfs)
-                for vrf in device_vrfs:
-                    self.get_model_or_update("ipam", "vrf", vrf, uuid=device_uuid)
-                    self.logger.increment_statistics(
-                        model="vrf", current=vrf_count, total=device_vrfs_total
-                    )
-                    vrf_count += 1
-
-            device_count += 1
-
-        if self.settings.get("inventoryitem"):
-            devices = Device.objects.all()
-            for device in devices:
-                device_parts = part_numbers.get(device.serial, [])
-                self.inventoryitem_count_total += len(device_parts)
-                for part in device_parts:
-                    self.get_model_or_update("dcim", "inventoryitem", part)
-                    self.logger.increment_statistics(
-                        model="inventory_item",
-                        current=self.inventoryitem_count,
-                        total=self.inventoryitem_count_total,
-                    )
-                    self.inventoryitem_count += 1
-
-        if self.settings.get("vlan"):
-            for vlan in vlans:
-                self.get_model_or_update("ipam", "vlan", vlan)
-                self.logger.increment_statistics(
-                    model="vlan", current=vlan_count, total=len(vlans)
-                )
-                vlan_count += 1
+        return interface_object
 
-        if self.settings.get("prefix"):
-            for network in networks:
-                self.get_model_or_update("ipam", "prefix", network)
-                self.logger.increment_statistics(
-                    model="prefix", current=network_count, total=len(networks)
-                )
-                network_count += 1
+    def collect_and_sync(self, branch=None) -> None:
+        self.logger.log_info("Starting data sync.", obj=self.sync)
+        (
+            sites,
+            devices,
+            interface_dict,
+            inventory_items,
+            vrfs,
+            vlans,
+            networks,
+            managed_ips,
+        ) = self.collect_data()
 
-    def sync(self):
-        self.sync_devices()
+        self.sync_items(
+            app_label="dcim", model="site", items=sites, cf=True, branch=branch
+        )
+        self.sync_devices(
+            branch,
+            devices,
+            interface_dict,
+            managed_ips,
+        )
+        self.sync_items(app_label="dcim", model="inventoryitem", items=inventory_items)
+        self.sync_items(app_label="ipam", model="vlan", items=vlans)
+        self.sync_items(app_label="ipam", model="vrf", items=vrfs)
+        self.sync_items(app_label="ipam", model="prefix", items=networks)