1 files changed, 22 insertions, 20 deletions
diff --git a/src/op_mode/show_neigh.py b/src/op_mode/neighbor.py
index d874bd544..b50e75d91 100755
--- a/src/op_mode/show_neigh.py
+++ b/src/op_mode/neighbor.py
@@ -28,29 +28,32 @@
 # ]
 
 import sys
+import typing
 
+import opmode
 
-def get_raw_data(family, device=None, state=None):
+
+def get_raw_data(family, interface=None, state=None):
     from json import loads
     from vyos.util import cmd
 
-    if device:
-        device = f"dev {device}"
+    if interface:
+        interface = f"dev {interface}"
     else:
-        device = ""
+        interface = ""
 
     if state:
         state = f"nud {state}"
     else:
         state = ""
 
-    neigh_cmd = f"ip --family {family} --json neighbor list {device} {state}"
+    neigh_cmd = f"ip --family {family} --json neighbor list {interface} {state}"
 
     data = loads(cmd(neigh_cmd))
 
     return data
 
-def get_formatted_output(family, device=None, state=None):
+def format_neighbors(neighs, interface=None):
     from tabulate import tabulate
 
     def entry_to_list(e, intf=None):
@@ -68,35 +71,34 @@ def get_formatted_output(family, device=None, state=None):
         # Device field is absent from outputs of `ip neigh list dev ...`
         if "dev" in e:
             dev = e["dev"]
-        elif device:
-            dev = device
+        elif interface:
+            dev = interface
         else:
             raise ValueError("interface is not defined")
 
         return [dst, dev, lladdr, state]
 
-    neighs = get_raw_data(family, device=device, state=state)
     neighs = map(entry_to_list, neighs)
 
     headers = ["Address", "Interface", "Link layer address",  "State"]
     return tabulate(neighs, headers)
 
-if __name__ == '__main__':
-    from argparse import ArgumentParser
+def show(raw: bool, family: str, interface: typing.Optional[str], state: typing.Optional[str]):
+    """ Display neighbor table contents """
+    data = get_raw_data(family, interface, state=state)
 
-    parser = ArgumentParser()
-    parser.add_argument("-f", "--family", type=str, default="inet", help="Address family")
-    parser.add_argument("-i", "--interface", type=str, help="Network interface")
-    parser.add_argument("-s", "--state", type=str, help="Neighbor table entry state")
+    if raw:
+        return data
+    else:
+        return format_neighbors(data, interface)
 
-    args = parser.parse_args()
 
-    if args.state:
-        if args.state not in ["reachable", "failed", "stale", "permanent"]:
-            raise ValueError(f"""Incorrect state "{args.state}"! Must be one of: reachable, stale, failed, permanent""")
+if __name__ == '__main__':
+    from argparse import ArgumentParser
 
     try:
-        print(get_formatted_output(args.family, device=args.interface, state=args.state))
+        print(opmode.run(sys.modules[__name__]))
     except ValueError as e:
         print(e)
         sys.exit(1)
+