feat: Interfaces now support setting the logical router #610

Author: adambaumeister

panos/base.py

@@ -1506,42 +1506,19 @@ def _convert_var(cls, value, vartype):
         elif vartype == "bool":
             return yesno(value)
 
-    def _set_reference(
+    def _get_all_objects_by_type(
         self,
-        reference_name,
         reference_type,
-        reference_var,
-        var_type,
-        exclusive,
         refresh,
-        update,
         running_config,
-        return_type,
         name_only,
-        **kwargs
+        reference_var
     ):
-        """Used by helper methods to set references between objects
-
-        For example, set_zone() would set the zone for an interface by creating a reference from
-        the zone to the interface. If the desired reference already exists then nothing happens.
-
-        This function has two modes:  refresh=True and refresh=False.  You
-        should only ever use refresh=False if:
-
-            1) all reference objects are in the current pan-os-python object tree
-            2) all reference objects are children attached to nearest_pandevice()
-            3) this is for firewall only, not a template / template stack
-            4) you're using firewall.vsys, not the device.Vsys object
-
-        If any of the above do not apply, you should be using refresh=True.
-
+        """
+        Returns all of the objects of the given type, ensuring that `reference_var` is refreshed and available to be
+        set, and the parent pandevice
         """
         parent = None
-        update_needed = False
-
-        if return_type not in ("bool", "object"):
-            raise ValueError("Unknown return_type specified: {0}".format(return_type))
-
         if refresh:
             """
             pan-os-python is too flexible:  users can use simple vsys mode or a
@@ -1616,6 +1593,22 @@ def _set_reference(
             parent = self.nearest_pandevice()
             allobjects = parent.findall(reference_type)
 
+        return parent, allobjects
+
+    def _update_reference_in_objects(
+        self,
+        parent,
+        allobjects,
+        reference_name,
+        reference_type,
+        reference_var,
+        var_type,
+        exclusive,
+        update,
+        return_type,
+        **kwargs
+    ):
+        update_needed = False
         # Find any current references to self and remove them, unless it is the desired reference
         if exclusive:
             for obj in allobjects:
@@ -1678,6 +1671,61 @@ def _set_reference(
         if return_type == "bool":
             return update_needed
 
+    def _set_reference(
+        self,
+        reference_name,
+        reference_type,
+        reference_var,
+        var_type,
+        exclusive,
+        refresh,
+        update,
+        running_config,
+        return_type,
+        name_only,
+        **kwargs
+    ):
+        """Used by helper methods to set references between objects
+
+        For example, set_zone() would set the zone for an interface by creating a reference from
+        the zone to the interface. If the desired reference already exists then nothing happens.
+
+        This function has two modes:  refresh=True and refresh=False.  You
+        should only ever use refresh=False if:
+
+            1) all reference objects are in the current pan-os-python object tree
+            2) all reference objects are children attached to nearest_pandevice()
+            3) this is for firewall only, not a template / template stack
+            4) you're using firewall.vsys, not the device.Vsys object
+
+        If any of the above do not apply, you should be using refresh=True.
+
+        """
+        if return_type not in ("bool", "object"):
+            raise ValueError("Unknown return_type specified: {0}".format(return_type))
+
+        # Get all the objects and the parent
+        parent, allobjects = self._get_all_objects_by_type(
+            reference_type,
+            refresh,
+            running_config,
+            name_only,
+            reference_var
+        )
+
+        return self._update_reference_in_objects(
+            parent,
+            allobjects,
+            reference_name,
+            reference_type,
+            reference_var,
+            var_type,
+            exclusive,
+            update,
+            return_type,
+            **kwargs
+        )
+
     def xml_merge(self, root, elements):
         """Merges other elements into the root element.
 

panos/network.py

@@ -20,6 +20,7 @@
 import logging
 import re
 import xml.etree.ElementTree as ET
+from keyword import kwlist
 
 import panos
 import panos.errors as err
@@ -522,6 +523,74 @@ def set_vlan(
             False,
         )
 
+    def set_logical_router(
+        self,
+        lr_name,
+        refresh=False,
+        update=False,
+        running_config=False,
+        return_type="object",
+        vrf_name="default",
+        **kwargs
+    ):
+        """adds the given interface to the VRF by name.
+
+        This is more complicated than `set_virtual_router` as the logical routers have child VRF child elements, which
+        is where the interfaces are configured.
+
+        This will use the VRF name 'default' by default.
+
+        Args:
+            logical_router_name (str): The name of the LogicalRouter or
+                a :class:`panos.network.LogicalRouter` instance
+            refresh (bool): Refresh the relevant current state of the device
+                before taking action (Default: False)
+            update (bool): Apply the changes to the device (Default: False)
+            running_config: If refresh is True, refresh from the running
+                configuration (Default: False)
+            return_type (str): Specify what this function returns, can be
+                either 'object' (the default) or 'bool'.  If this is 'object',
+                then the return value is the LogicalRouter in question.  If
+                this is 'bool', then the return value is a boolean that tells
+                you about if the live device needs updates (update=False) or
+                was updated (update=True).
+        """
+
+        # First we get all the logical routers
+        parent, all_logical_routers = self._get_all_objects_by_type(
+            LogicalRouter,
+            refresh,
+            running_config,
+            name_only=False,
+            reference_var="interface"
+        )
+        lr: LogicalRouter | None
+
+        lr = next((lr for lr in all_logical_routers if lr.uid == lr_name), None)
+        if not lr:
+            # If the LR isn't found, create it instead
+            lr = LogicalRouter(name=lr_name)
+            parent.add(lr)
+            vrf = Vrf(name=vrf_name)
+            lr.vrf = [vrf]
+            lr.create()
+
+        # Pass all the vrfs to the update method
+        return self._update_reference_in_objects(
+            lr,
+            lr.vrf,
+            reference_name=vrf_name,
+            reference_var="interface",
+            reference_type=Vrf,
+            var_type="list",
+            return_type=return_type,
+            update=update,
+            exclusive=True,
+            **kwargs
+        )
+
+
+
     def get_counters(self):
         """Pull the counters for an interface
 

tests/live/test_network.py

@@ -1,6 +1,9 @@
 import random
+import traceback
+from inspect import trace
 
 from panos import device, network
+from panos.network import Interface
 from tests.live import testlib
 
 
@@ -586,6 +589,44 @@ def cleanup_dependencies(self, fw, state):
         except Exception:
             pass
 
+class TestVirtualRouterInterfaces(testlib.FwFlow):
+    def create_dependencies(self, fw, state):
+        # Disable ARE on the device so we're just using VR setup
+        advanced_routing_engine = device.AdvancedRoutingEngine(enable=False)
+        fw.add(advanced_routing_engine)
+        advanced_routing_engine.create()
+
+        state.eth_obj = None
+        state.eth = testlib.get_available_interfaces(fw, max_interfaces=7)[0]
+
+        state.eth_obj = network.EthernetInterface(
+            state.eth, "layer3", testlib.random_ip("/24")
+        )
+        fw.add(state.eth_obj)
+        state.eth_obj.create()
+
+        state.obj = network.VirtualRouter(testlib.random_name())
+        fw.add(state.obj)
+
+    def setup_state_obj(self, fw, state):
+        pass
+
+    def test_10_set_vr(self, fw, state_map):
+        state = self.sanity(fw, state_map)
+        state.eth_obj.set_virtual_router(state.obj.name)
+
+    def cleanup_dependencies(self, fw, state):
+        return
+        try:
+            state.vr.delete()
+        except Exception:
+            pass
+
+        try:
+            state.eth_obj.delete()
+        except Exception:
+            pass
+
 
 class TestStaticRouteV6(testlib.FwFlow):
     def create_dependencies(self, fw, state):
@@ -1659,13 +1700,21 @@ def create_dependencies(self, fw, state):
         state.advanced_routing_engine_obj.create()
 
         state.eth_obj = None
-        state.eth = testlib.get_available_interfaces(fw)[0]
+        available_interfaces = testlib.get_available_interfaces(fw, 2, 7)
+        state.eth = available_interfaces[0]
+        state.eth_2 = available_interfaces[1]
 
         state.eth_obj = network.EthernetInterface(
             state.eth, "layer3", testlib.random_ip("/24")
         )
+        state.eth_obj_2 = network.EthernetInterface(
+            state.eth_2, "layer3", testlib.random_ip("/24")
+        )
         fw.add(state.eth_obj)
+        fw.add(state.eth_obj_2)
         state.eth_obj.create()
+        state.eth_obj_2.create()
+
 
     def setup_state_obj(self, fw, state):
         vrf = network.Vrf(
@@ -1690,18 +1739,41 @@ def setup_state_obj(self, fw, state):
         )
         lr = network.LogicalRouter(testlib.random_name())
         lr.add(vrf)
+
+        lr_2 = network.LogicalRouter(testlib.random_name())
+        vrf2 = network.Vrf(
+            "default"
+        )
+        lr_2.add(vrf2)
+
         state.obj = lr
+        state.obj_2 = lr_2
         fw.add(state.obj)
+        fw.add(state.obj_2)
+        state.obj.create()
+        state.obj_2.create()
+
+    def test_10_set_lr_for_interface(self, fw, state_map):
+        """Test setting the LR for an interface instead of the other way around"""
+        state = self.sanity(fw, state_map)
+        eth: Interface = state.eth_obj_2
+        eth.set_logical_router(state.obj_2, update=True)
 
     def update_state_obj(self, fw, state):
         state.obj.ad_static = random.randint(10, 240)
         state.obj.ad_rip = random.randint(10, 240)
 
     def cleanup_dependencies(self, fw, state):
         try:
+            fw.add(state.obj)
+            fw.add(state.obj_2)
+            state.obj.delete()
+            state.obj_2.delete()
             state.eth_obj.delete()
-            state.advanced_routing_engine_obj.delete()
-        except Exception:
+            state.eth_obj_2.delete()
+        except Exception as e:
+            print(traceback.format_exc())
+            print(e)
             pass
 
 

tests/live/testlib.py

@@ -1,4 +1,5 @@
 import random
+import traceback
 
 import pytest
 
@@ -47,11 +48,11 @@ def random_mac():
     return ":".join("{0:02x}".format(random.randint(0, 255)) for x in range(6))
 
 
-def get_available_interfaces(con, num=1):
+def get_available_interfaces(con, num=1, max_interfaces=10):
     ifaces = network.EthernetInterface.refreshall(con, add=False)
     ifaces = set(x.name for x in ifaces)
 
-    all_interfaces = set("ethernet1/{0}".format(x) for x in range(1, 10))
+    all_interfaces = set("ethernet1/{0}".format(x) for x in range(1, max_interfaces))
     available = all_interfaces.difference(ifaces)
 
     ans = []
@@ -73,6 +74,8 @@ def test_01_setup_dependencies(self, fw, state_map):
         except Exception as e:
             print("SETUP ERROR: {0}".format(e))
             state.err = True
+            print(traceback.format_exc())
+
             pytest.skip("Setup failed")
 
     def create_dependencies(self, fw, state):