From 047f0a5a804966a431649704f2e72af48e54c931 Mon Sep 17 00:00:00 2001
From: wdesmedt <walter.de.smedt@gmail.com>
Date: Wed, 3 Jan 2024 16:21:51 +0000
Subject: [PATCH] Add support for LAGs and Ethernet Segments (ES)

---
 nornir_srl/connections/srlinux.py | 40 ++++++++++++++++++++++
 nornir_srl/fsc.py                 | 57 +++++++++++++++++++++++++++++++
 2 files changed, 97 insertions(+)

diff --git a/nornir_srl/connections/srlinux.py b/nornir_srl/connections/srlinux.py
index 95c5215..0c2cc38 100644
--- a/nornir_srl/connections/srlinux.py
+++ b/nornir_srl/connections/srlinux.py
@@ -596,6 +596,46 @@ def get_nwi_itf(self, nw_instance: str = "*") -> Dict[str, Any]:
 
         res = jmespath.search(path_spec["jmespath"], resp[0])
         return {"nwi_itfs": res}
+    
+    def get_lag(self, lag_id: str = "*") -> Dict[str, Any]:
+        path_spec = {
+            "path": f"/interface[name=lag{lag_id}]",
+            "jmespath": '"interface"[].{lag:name, oper:"oper-state",mtu:mtu,num:lag.member|length(@),"min":lag."min-links",desc:description, type:lag."lag-type", speed:lag."lag-speed","stby-sig":ethernet."standby-signaling",\
+                  "lacp-key":lag.lacp."admin-key","lacp-itvl":lag.lacp.interval,"lacp-mode":lag.lacp."lacp-mode","lacp-sysid":lag.lacp."system-id-mac","lacp-prio":lag.lacp."system-priority",\
+                    members:lag.member[].{"member-itf":name, "member-oper":"oper-state","act":lacp."activity"}}',
+            "datatype": "state",
+        }
+        resp = self.get(
+            paths=[path_spec.get("path", "")], datatype=path_spec["datatype"]
+        )
+        res = jmespath.search(path_spec["jmespath"], resp[0])
+        return {"lag": res}
+
+    def get_es(self) -> Dict[str, Any]:
+        path_spec = {
+            "path": f"/system/network-instance/protocols/evpn/ethernet-segments",
+            "jmespath": '"system/network-instance/protocols/evpn/ethernet-segments"."bgp-instance"[]."ethernet-segment"[].{name:name, esi:esi, "mh-mode":"multi-homing-mode",\
+                oper:"oper-state",itfs:interface[]."ethernet-interface"|join(\' \',@), vrfs:association."network-instance"[].{ni:name, "peers":"_peers"}}',
+            "datatype": "state",
+        }
+        def set_es_peers(resp):
+            for bgp_inst in resp[0].get("system/network-instance/protocols/evpn/ethernet-segments", {}).get("bgp-instance", []):
+                for es in bgp_inst.get("ethernet-segment", []):
+                    for vrf in es.get("association", {}).get("network-instance", []):
+                        es_peers = vrf["bgp-instance"][0].get("computed-designated-forwarder-candidates", {}).get("designated-forwarder-candidate", [])
+                        vrf["_peers"] = ", ".join(
+                                f"{peer['address']} (DF)" if peer["designated-forwarder"] else peer["address"]
+                                for peer in es_peers
+                        )
+        if not "evpn" in self.get(paths=["/system/features"], datatype="state")[0]["system/features"]:
+            return {"es": []}
+        resp = self.get(
+            paths=[path_spec.get("path", "")], datatype=path_spec["datatype"]
+        )
+        set_es_peers(resp)
+        res = jmespath.search(path_spec["jmespath"], resp[0])
+        return {"es": res}
+
 
     def get(
         self,
diff --git a/nornir_srl/fsc.py b/nornir_srl/fsc.py
index 2667000..3226435 100644
--- a/nornir_srl/fsc.py
+++ b/nornir_srl/fsc.py
@@ -482,6 +482,35 @@ def _subinterfaces(task: Task) -> Result:
         i_filter=ctx.obj["i_filter"],
     )
 
+@cli.command()
+@click.pass_context
+@click.option(
+    "--field-filter",
+    "-f",
+    multiple=True,
+    help='filter fields with <field-name>=<glob-pattern>, e.g. -f state=up -f admin_state="ena*". Fieldnames correspond to column names of a report',
+)
+def lag(ctx: Context, field_filter: Optional[List] = None):
+    """Displays LAGs of nodes"""
+
+    def _lag(task: Task) -> Result:
+        device = task.host.get_connection(CONNECTION_NAME, task.nornir.config)
+        return Result(host=task.host, result=device.get_lag())
+
+    f_filter = (
+        {k: v for k, v in [f.split("=") for f in field_filter]} if field_filter else {}
+    )
+    result = ctx.obj["target"].run(
+        task=_lag, name="lag", raise_on_error=False
+    )
+    print_report(
+        result=result,
+        name="LAGs",
+        failed_hosts=result.failed_hosts,
+        box_type=ctx.obj["box_type"],
+        f_filter=f_filter,
+        i_filter=ctx.obj["i_filter"],
+    )
 
 @cli.command()
 @click.pass_context
@@ -667,6 +696,34 @@ def _lldp_neighbors(task: Task) -> Result:
         i_filter=ctx.obj["i_filter"],
     )
 
+@cli.command()
+@click.pass_context
+@click.option(
+    "--field-filter",
+    "-f",
+    multiple=True,
+    help='filter fields with <field-name>=<glob-pattern>, e.g. -f name=ge-0/0/0 -f admin_state="ena*". Fieldnames correspond to column names of a report',
+)
+def es(ctx: Context, field_filter: Optional[List] = None):
+    """Displays Ethernet Segments"""
+
+    def _es(task: Task) -> Result:
+        device = task.host.get_connection(CONNECTION_NAME, task.nornir.config)
+        return Result(host=task.host, result=device.get_es())
 
+    f_filter = (
+        {k: v for k, v in [f.split("=") for f in field_filter]} if field_filter else {}
+    )
+
+    result = ctx.obj["target"].run(task=_es, name="es", raise_on_error=False)
+    print_report(
+        result=result,
+        name="Ethernet Segments",
+        failed_hosts=result.failed_hosts,
+        box_type=ctx.obj["box_type"],
+        f_filter=f_filter,
+        i_filter=ctx.obj["i_filter"]
+    )
+    
 if __name__ == "__main__":
     cli(obj={})