fixes issue with finding ip prefixx in NetBox 4.2 and reassignes VRFs to IPs #397

Author: bb-Ricardo

module/netbox/object_classes.py

@@ -1442,6 +1442,7 @@ def update(self, data=None, read_from_netbox=False, source=None):
 
         super().update(data=data, read_from_netbox=read_from_netbox, source=source)
 
+
 class NBVLANGroup(NetBoxObject):
     name = "VLANGroup"
     api_path = "ipam/vlan-groups"
@@ -1477,8 +1478,9 @@ def resolve_relations(self):
 
         if isinstance(o_id, int) and o_type is not None and self.data_model_relation.get(o_type) is not None:
             self.data["scope_id"] = self.inventory.get_by_id(self.data_model_relation.get(o_type), nb_id=o_id)
-        elif not isinstance(o_id, NetBoxObject):
-            log.debug(f"{self.name} '{self.data.get('name')}' scope type '{o_type}' for '{grab(self, 'data.scope.name')}'  is currently not supported")
+        elif o_id is not None and not isinstance(o_id, NetBoxObject):
+            log.debug(f"{self.name} '{self.data.get('name')}' scope type '{o_type}' for "
+                      f"'{grab(self, 'data.scope.name')}' is currently not supported")
             self.data["scope_id"] = ""
 
         super().resolve_relations()
@@ -1500,16 +1502,16 @@ def matches_site_cluster(self, site=None, cluster=None) -> bool:
 
         """
         if isinstance(site, NBSite):
-            if isinstance(self.data["scope_id"], NBSite) and self.data["scope_id"] == site:
+            if isinstance(self.data.get("scope_id"), NBSite) and self.data.get("scope_id") == site:
                 return True
-            if (isinstance(self.data["scope_id"], NBSiteGroup) and
+            if (isinstance(self.data.get("scope_id"), NBSiteGroup) and
                     self.data["scope_id"] == grab(site, "data.group")):
                 return True
 
         if isinstance(cluster, NBCluster):
-            if isinstance(self.data["scope_id"], NBCluster) and self.data["scope_id"] == cluster:
+            if isinstance(self.data.get("scope_id"), NBCluster) and self.data.get("scope_id") == cluster:
                 return True
-            if (isinstance(self.data["scope_id"], NBClusterGroup) and
+            if (isinstance(self.data.get("scope_id"), NBClusterGroup) and
                     self.data["scope_id"] == grab(cluster, "data.group")):
                 return True
 
@@ -1545,19 +1547,29 @@ def __init__(self, *args, **kwargs):
         self.data_model = {
             "prefix": [IPv4Network, IPv6Network],
             "site": NBSite,
+            "scope_type": ["dcim.site", "dcim.sitegroup"],
+            "scope_id": [NBSite, NBSiteGroup],
             "tenant": NBTenant,
             "vlan": NBVLAN,
             "vrf": NBVRF,
             "description": 200,
             "tags": NBTagList
         }
+        # add relation between two attributes
+        self.data_model_relation = {
+            "dcim.site": NBSite,
+            "dcim.sitegroup": NBSiteGroup,
+            NBSite: "dcim.site",
+            NBSiteGroup: "dcim.sitegroup",
+        }
+
         super().__init__(*args, **kwargs)
 
     def update(self, data=None, read_from_netbox=False, source=None):
 
         # prefixes are parsed into ip_networks
         data_prefix = data.get(self.primary_key)
-        if data_prefix is not None and not isinstance(data_prefix, (IPv4Network, IPv6Network)):
+        if not isinstance(data_prefix, (IPv4Network, IPv6Network)) and data_prefix is not None:
             try:
                 data[self.primary_key] = ip_network(data_prefix)
             except ValueError as e:
@@ -1569,6 +1581,47 @@ def update(self, data=None, read_from_netbox=False, source=None):
         if read_from_netbox is False:
             raise ValueError(f"Adding {self.name} by this program is currently not implemented.")
 
+    def resolve_relations(self):
+
+        o_id = self.data.get("scope_id")
+        o_type = self.data.get("scope_type")
+
+        if isinstance(o_id, int) and o_type is not None and self.data_model_relation.get(o_type) is not None:
+            self.data["scope_id"] = self.inventory.get_by_id(self.data_model_relation.get(o_type), nb_id=o_id)
+        elif o_id is not None and not isinstance(o_id, NetBoxObject):
+            log.debug(f"{self.name} '{self.data.get('name')}' scope type '{o_type}' for "
+                      f"'{grab(self, 'data.scope.name')}' is currently not supported")
+            self.data["scope_id"] = ""
+
+        super().resolve_relations()
+
+    def matches_site(self, site=None) -> bool:
+        """
+        tries to figure out if this prefix matches a certain site or site group
+
+        Parameters
+        ----------
+        site: NBSite
+            the site object to match to
+
+        Returns
+        -------
+        bool: True if matches one of the params
+
+        """
+        if isinstance(site, NBSite):
+            if isinstance(self.data.get("scope_id"), NBSite) and self.data.get("scope_id") == site:
+                return True
+            if (isinstance(self.data.get("scope_id"), NBSiteGroup) and
+                    self.data["scope_id"] == grab(site, "data.group")):
+                return True
+
+            # compatible for NetBox versions < 4.2.0
+            if self.data.get("site") == site:
+                return True
+
+        return False
+
 
 class NBManufacturer(NetBoxObject):
     name = "manufacturer"
@@ -1686,7 +1739,7 @@ def __init__(self, *args, **kwargs):
             "type": NBClusterType,
             "tenant": NBTenant,
             "group": NBClusterGroup,
-            "scope_type": ["dcim.site", "dcim.sitegroup", "dcim.location", "dcim.region"],
+            "scope_type": ["dcim.site", "dcim.sitegroup"],
             "scope_id": NBSite,
             "tags": NBTagList
         }

module/sources/common/source_base.py

@@ -216,7 +216,7 @@ def return_longest_matching_prefix_for_ip(self, ip_to_match=None, site_name=None
 
         for prefix in self.inventory.get_all_items(NBPrefix):
 
-            if grab(prefix, "data.site") != site_object:
+            if not prefix.matches_site(site_object):
                 continue
 
             prefix_network = grab(prefix, f"data.{NBPrefix.primary_key}")
@@ -286,7 +286,7 @@ def add_update_interface(self, interface_object, device_object, interface_data,
             site_name = device_object_cluster.get_site_name()
         elif type(device_object) == NBDevice:
             interface_class = NBInterface
-            site_name = grab(device_object, "data.site.data.name")
+            site_name = device_object.get_site_name()
         elif device_object is None:
             log.error(f"No device/VM object submitted to attach interface '{grab(interface_data, 'name')}' to.")
             return None
@@ -415,7 +415,7 @@ def add_update_interface(self, interface_object, device_object, interface_data,
 
                 # check if IP address is of type IP interface (includes prefix length)
                 if type(ip_object) in [IPv6Address, IPv4Address]:
-                    log.warning(f"{log_text}. Unable to add IP address to NetBox.")
+                    log.warning(f"{log_text}. Unable to add IP address to NetBox")
                     continue
                 else:
                     log.debug2(log_text)
@@ -426,7 +426,7 @@ def add_update_interface(self, interface_object, device_object, interface_data,
                 if type(this_prefix) in [IPv4Network, IPv6Network]:
                     ip_object = ip_interface(f"{ip_object}/{this_prefix.prefixlen}")
                 else:
-                    log.warning(f"{matching_ip_prefix.name} got wrong format. Unable to add IP to NetBox")
+                    log.warning(f"{matching_ip_prefix.name} got wrong format. Unable to add IP address to NetBox")
                     continue
 
             # try to find matching IP address object
@@ -563,27 +563,21 @@ def add_update_interface(self, interface_object, device_object, interface_data,
                         ip_tenant = prefix_tenant
                         break
 
+            if possible_ip_vrf is not None:
+                nic_ip_data["vrf"] = possible_ip_vrf
+            if ip_tenant is not None:
+                nic_ip_data["tenant"] = ip_tenant
+
             if not isinstance(this_ip_object, NBIPAddress):
                 log.debug(f"No existing {NBIPAddress.name} object found. Creating a new one.")
 
-                if possible_ip_vrf is not None:
-                    nic_ip_data["vrf"] = possible_ip_vrf
-                if ip_tenant is not None:
-                    nic_ip_data["tenant"] = ip_tenant
-
                 this_ip_object = self.inventory.add_object(NBIPAddress, data=nic_ip_data, source=self)
 
             # update IP address with additional data if not already present
             else:
 
                 log.debug2(f"Found existing NetBox {NBIPAddress.name} object: {this_ip_object.get_display_name()}")
 
-                if grab(this_ip_object, "data.vrf") is None and possible_ip_vrf is not None:
-                    nic_ip_data["vrf"] = possible_ip_vrf
-
-                if grab(this_ip_object, "data.tenant") is None and ip_tenant is not None:
-                    nic_ip_data["tenant"] = ip_tenant
-
                 this_ip_object.update(data=nic_ip_data, source=self)
 
             ip_address_objects.append(this_ip_object)