From 17d7648b9c8e3ace4078b49214d405eb5bf3a007 Mon Sep 17 00:00:00 2001
From: checktheroads <mtt.lov@gmail.com>
Date: Fri, 13 Sep 2019 00:36:58 -0700
Subject: [PATCH] add VRF to AFI mapping model

---
 hyperglass/configuration/models/routers.py | 59 +++++++++++++++++++++-
 hyperglass/configuration/models/vrfs.py    | 12 ++---
 hyperglass/constants.py                    |  9 ++++
 3 files changed, 73 insertions(+), 7 deletions(-)

diff --git a/hyperglass/configuration/models/routers.py b/hyperglass/configuration/models/routers.py
index 762a5ce..cdaa079 100644
--- a/hyperglass/configuration/models/routers.py
+++ b/hyperglass/configuration/models/routers.py
@@ -20,6 +20,44 @@ from pydantic import validator
 from hyperglass.configuration.models._utils import clean_name
 from hyperglass.constants import Supported
 from hyperglass.exceptions import UnsupportedDevice
+from hyperglass.constants import afi_nos_map
+
+
+class AfiMap(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
+
+    @validator("ipv4", always=True)
+    def validate_ipv4(cls, v):  # noqa: N805
+        """If a map field is undefined, get a default value"""
+        if v is None:
+            v = afi_nos_map.get("default").get("ipv4")
+        return v
+
+    @validator("ipv6", always=True)
+    def validate_ipv6(cls, v):  # noqa: N805
+        """If a map field is undefined, get a default value"""
+        if v is None:
+            v = afi_nos_map.get("default").get("ipv6")
+        return v
+
+    @validator("ipv4_vpn", always=True)
+    def validate_ipv4_vpn(cls, v):  # noqa: N805
+        """If a map field is undefined, get a default value"""
+        if v is None:
+            v = afi_nos_map.get("default").get("ipv4_vpn")
+        return v
+
+    @validator("ipv6_vpn", always=True)
+    def validate_ipv6_vpn(cls, v):  # noqa: N805
+        """If a map field is undefined, get a default value"""
+        if v is None:
+            v = afi_nos_map.get("default").get("ipv6_vpn")
+        return v
 
 
 class Router(BaseSettings):
@@ -37,10 +75,13 @@ class Router(BaseSettings):
     nos: str
     commands: Union[str, None] = None
     vrfs: List[str] = ["default"]
+    afi_map: Union[AfiMap, None] = None
 
     @validator("nos")
     def supported_nos(cls, v):  # noqa: N805
-        """Validates that passed nos string is supported by hyperglass"""
+        """
+        Validates that passed nos string is supported by hyperglass.
+        """
         if not Supported.is_supported(v):
             raise UnsupportedDevice(f'"{v}" device type is not supported.')
         return v
@@ -52,10 +93,26 @@ class Router(BaseSettings):
 
     @validator("commands", always=True)
     def validate_commands(cls, v, values):  # noqa: N805
+        """
+        If a named command profile is not defined, use theNOS name.
+        """
         if v is None:
             v = values["nos"]
         return v
 
+    @validator("afi_map", always=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")))
+        return v
+
+
+Router.update_forward_refs()
+
 
 class Routers(BaseSettings):
     """Base model for devices class."""
diff --git a/hyperglass/configuration/models/vrfs.py b/hyperglass/configuration/models/vrfs.py
index e8aff34..28ead91 100644
--- a/hyperglass/configuration/models/vrfs.py
+++ b/hyperglass/configuration/models/vrfs.py
@@ -19,7 +19,7 @@ class Vrf(BaseSettings):
     """Model for per VRF/afi config in devices.yaml"""
 
     display_name: str
-    name: str
+    label: str
     afis: List[str]
 
 
@@ -36,26 +36,26 @@ class Vrfs(BaseSettings):
         vrfs: Vrf = {
             "default": {
                 "display_name": "Default",
-                "name": "default",
+                "label": "default",
                 "afis": ["ipv4, ipv6"],
             }
         }
-        names: List[str] = ["default"]
+        labels: List[str] = ["default"]
         _all: List[str] = ["default"]
 
         for (vrf_key, params) in input_params.items():
             vrf = clean_name(vrf_key)
             vrf_params = Vrf(**params)
             vrfs.update({vrf: vrf_params.dict()})
-            names.append(params.get("name"))
+            labels.append(params.get("label"))
             _all.append(vrf_key)
         for (vrf_key, params) in vrfs.items():
             setattr(Vrfs, vrf_key, params)
 
-        names: List[str] = list(set(names))
+        labels: List[str] = list(set(labels))
         _all: List[str] = list(set(_all))
         Vrfs.vrfs = vrfs
-        Vrfs.names = names
+        Vrfs.labels = labels
         Vrfs._all = _all
         return Vrfs()
 
diff --git a/hyperglass/constants.py b/hyperglass/constants.py
index 2d56ce2..28899f2 100644
--- a/hyperglass/constants.py
+++ b/hyperglass/constants.py
@@ -4,6 +4,15 @@ Global Constants for hyperglass
 
 protocol_map = {80: "http", 8080: "http", 443: "https", 8443: "https"}
 
+afi_nos_map = {
+    "default": {
+        "ipv4": "ipv4",
+        "ipv6": "ipv6",
+        "ipv4_vpn": "vpnv4",
+        "ipv6_vpn": "vpnv6",
+    }
+}
+
 
 class Supported:
     """