add VRF support

hyperglass/command/construct.py

@@ -56,11 +56,11 @@ class Construct:
         Constructs AFI string. If query_vrf is specified, AFI prefix is
         "vpnv", if not, AFI prefix is "ipv"
         """
-        ip_version = ipaddress.ip_network(query_target).version
+        protocol = ipaddress.ip_network(query_target).version
         if query_vrf:
-            afi = f"vpnv{ip_version}"
+            afi = f"ipv{protocol}_vpn"
         else:
-            afi = f"ipv{ip_version}"
+            afi = f"ipv{protocol}"
         return afi
 
     def ping(self):
@@ -71,28 +71,36 @@ class Construct:
         )
 
         query = []
-        afi = self.query_afi(self.query_target, self.query_vrf)
-        source = self.get_src(self.device, afi)
+        query_vrfs = self.query_vrf
 
-        if self.transport == "rest":
-            query = json.dumps(
-                {
-                    "query_type": "ping",
-                    "afi": afi,
-                    "vrf": self.query_vrf,
-                    "source": source,
-                    "target": self.query_target,
-                }
-            )
-        elif self.transport == "scrape":
-            cmd = self.device_commands(self.device.commands, afi, "ping")
-            query = cmd.format(
-                target=self.query_target, source=source, vrf=self.query_vrf
-            )
+        for vrf in query_vrfs:
+            query_afi = self.query_afi(self.query_target, vrf)
+            afi_path = f"self.device.afis.{query_afi}"
+            afi = getattr(afi_path, "label")
+            vrf_label = vrfs.get(vrf).get("label")
+            vrf_source = getattr(afi_path, "source")
+
+            if self.transport == "rest":
+                vrf_query = json.dumps(
+                    {
+                        "query_type": "ping",
+                        "afi": afi,
+                        "vrf": vrf_label,
+                        "source": vrf_source,
+                        "target": self.query_target,
+                    }
+                )
+            elif self.transport == "scrape":
+                cmd = self.device_commands(self.device.commands, afi, "ping")
+                query.append(
+                    cmd.format(
+                        target=self.query_target, source=vrf_source, vrf=vrf_label
+                    )
+                )
+            query.append(vrf_query)
 
         logger.debug(f"Constructed query: {query}")
-
-        return [query]
+        return query
 
     def traceroute(self):
         """

hyperglass/configuration/__init__.py

@@ -88,7 +88,7 @@ except ValidationError as validation_errors:
         raise ConfigInvalid(
             field=": ".join([str(item) for item in error["loc"]]),
             error_msg=error["msg"],
-        ) from None
+        )
 
 # Validate that VRFs configured on a device are actually defined
 for dev in devices.hostnames:

hyperglass/configuration/models/routers.py

@@ -6,6 +6,7 @@ Imports config variables and overrides default class attributes.
 Validates input for overridden parameters.
 """
 # Standard Library Imports
+import operator
 from ipaddress import IPv4Address
 from ipaddress import IPv6Address
 from typing import List
@@ -15,6 +16,7 @@ from typing import Union
 from pydantic import BaseSettings
 from pydantic import IPvAnyAddress
 from pydantic import validator
+from logzero import logger
 
 # Project Imports
 from hyperglass.configuration.models._utils import clean_name
@@ -23,13 +25,20 @@ from hyperglass.exceptions import UnsupportedDevice
 from hyperglass.constants import afi_nos_map
 
 
-class AfiMap(BaseSettings):
+class Afi(BaseSettings):
+    """Model for AFI definitions"""
+
+    label: str
+    source: IPvAnyAddress
+
+
+class Afis(BaseSettings):
     """Model for AFI map"""
 
-    ipv4: Union[str, None] = None
-    ipv6: Union[str, None] = None
-    ipv4_vpn: Union[str, None] = None
-    ipv6_vpn: Union[str, None] = None
+    ipv4: Union[Afi, None] = None
+    ipv6: Union[Afi, None] = None
+    ipv4_vpn: Union[Afi, None] = None
+    ipv6_vpn: Union[Afi, None] = None
 
 
 class Router(BaseSettings):
@@ -37,8 +46,6 @@ class Router(BaseSettings):
 
     address: Union[IPvAnyAddress, str]
     network: str
-    src_addr_ipv4: IPv4Address
-    src_addr_ipv6: IPv6Address
     credential: str
     proxy: Union[str, None] = None
     location: str
@@ -47,7 +54,7 @@ class Router(BaseSettings):
     nos: str
     commands: Union[str, None] = None
     vrfs: List[str] = ["default"]
-    afi_map: Union[AfiMap, None] = None
+    afis: Afis
 
     @validator("nos")
     def supported_nos(cls, v):  # noqa: N805
@@ -72,20 +79,24 @@ class Router(BaseSettings):
             v = values["nos"]
         return v
 
-    @validator("afi_map", always=True)
+    @validator("afis", pre=True)
     def validate_afis(cls, v, values):  # noqa: N805
         """
         If an AFI map is not defined, try to get one based on the
         NOS name. If that doesn't exist, use a default.
         """
-        if v is None:
-            v = AfiMap(**afi_nos_map.get(values["nos"], afi_nos_map.get("default")))
+        logger.debug(f"V In: {v}")
+        for (afi_name, afi_params) in {
+            afi: params for afi, params in v.items() if params is not None
+        }.items():
+            if afi_params.get("label") is None:
+                label = afi_nos_map.get(values["nos"], None)
+                if label is None:
+                    label = afi_nos_map["default"][afi_name]
+                v[afi_name].update({"label": label})
         return v
 
 
-Router.update_forward_refs()
-
-
 class Routers(BaseSettings):
     """Base model for devices class."""