checktheroads-hyperglass/hyperglass/configuration/models/vrfs.py

"""Validate VRF configuration variables."""

# Standard Library Imports
from ipaddress import IPv4Address
from ipaddress import IPv4Network
from ipaddress import IPv6Address
from ipaddress import IPv6Network
from typing import Dict
from typing import List
from typing import Optional

# Third Party Imports
from pydantic import FilePath
from pydantic import IPvAnyNetwork
from pydantic import StrictBool
from pydantic import StrictStr
from pydantic import constr
from pydantic import validator

# Project Imports
from hyperglass.configuration.models._utils import HyperglassModel
from hyperglass.configuration.models._utils import HyperglassModelExtra


class InfoConfigParams(HyperglassModelExtra):
    """Validation model for per-help params."""

    title: Optional[StrictStr]


class InfoConfig(HyperglassModel):
    """Validation model for help configuration."""

    enable: StrictBool = True
    file: Optional[FilePath]
    params: InfoConfigParams = InfoConfigParams()


class Info(HyperglassModel):
    """Validation model for per-VRF, per-Command help."""

    bgp_aspath: InfoConfig = InfoConfig()
    bgp_community: InfoConfig = InfoConfig()
    bgp_route: InfoConfig = InfoConfig()
    ping: InfoConfig = InfoConfig()
    traceroute: InfoConfig = InfoConfig()


class DeviceVrf4(HyperglassModel):
    """Validation model for IPv4 AFI definitions."""

    vrf_name: StrictStr
    source_address: IPv4Address


class DeviceVrf6(HyperglassModel):
    """Validation model for IPv6 AFI definitions."""

    vrf_name: StrictStr
    source_address: IPv6Address


class Vrf(HyperglassModel):
    """Validation model for per VRF/afi config in devices.yaml."""

    name: StrictStr
    display_name: StrictStr
    info: Info = Info()
    ipv4: Optional[DeviceVrf4]
    ipv6: Optional[DeviceVrf6]
    access_list: List[Dict[constr(regex=("allow|deny")), IPvAnyNetwork]] = [
        {"allow": IPv4Network("0.0.0.0/0")},
        {"allow": IPv6Network("::/0")},
    ]

    @validator("ipv4", "ipv6", pre=True, always=True)
    def set_default_vrf_name(cls, value, values):
        """If per-AFI name is undefined, set it to the global VRF name.

        Returns:
            {str} -- VRF Name
        """
        if isinstance(value, DefaultVrf) and value.vrf_name is None:
            value["vrf_name"] = values["name"]
        elif isinstance(value, Dict) and value.get("vrf_name") is None:
            value["vrf_name"] = values["name"]
        return value

    @validator("access_list", pre=True)
    def validate_action(cls, value):
        """Transform ACL networks to IPv4Network/IPv6Network objects.

        Returns:
            {object} -- IPv4Network/IPv6Network object
        """
        for li in value:
            for action, network in li.items():
                if isinstance(network, (IPv4Network, IPv6Network)):
                    li[action] = str(network)
        return value

    def __hash__(self):
        """Make VRF object hashable so the object can be deduplicated with set().

        Returns:
            {int} -- Hash of VRF name
        """
        return hash((self.name,))

    def __eq__(self, other):
        """Make VRF object comparable so the object can be deduplicated with set().

        Arguments:
            other {object} -- Object to compare

        Returns:
            {bool} -- True if comparison attributes are the same value
        """
        return self.name == other.name


class DefaultVrf(HyperglassModel):
    """Validation model for default routing table VRF."""

    name: StrictStr = "default"
    display_name: StrictStr = "Global"
    info: Info = Info()
    access_list: List[Dict[constr(regex=("allow|deny")), IPvAnyNetwork]] = [
        {"allow": IPv4Network("0.0.0.0/0")},
        {"allow": IPv6Network("::/0")},
    ]

    class DefaultVrf4(HyperglassModel):
        """Validation model for IPv4 default routing table VRF definition."""

        vrf_name: StrictStr = "default"
        source_address: IPv4Address

    class DefaultVrf6(HyperglassModel):
        """Validation model for IPv6 default routing table VRF definition."""

        vrf_name: StrictStr = "default"
        source_address: IPv6Address

    ipv4: Optional[DefaultVrf4]
    ipv6: Optional[DefaultVrf6]