igmp-proxy: T4912: Rewrite show IGMP proxy commands in the new op-mode format

2 files changed, 86 insertions, 109 deletions

diff --git a/op-mode-definitions/show-ip-multicast.xml.in b/op-mode-definitions/show-ip-multicast.xml.in
index 80d83b424..d8f16d997 100644
--- a/op-mode-definitions/show-ip-multicast.xml.in
+++ b/op-mode-definitions/show-ip-multicast.xml.in
@@ -13,13 +13,13 @@
                 <properties>
                   <help>Show multicast interfaces</help>
                 </properties>
-                <command>if ps -C igmpproxy &amp;&gt;/dev/null; then ${vyos_op_scripts_dir}/show_igmpproxy.py --interface; else echo IGMP proxy not configured; fi</command>
+                <command>${vyos_op_scripts_dir}/igmp-proxy.py show_interface</command>
               </leafNode>
               <leafNode name="mfc">
                 <properties>
                   <help>Show multicast fowarding cache</help>
                 </properties>
-                <command>if ps -C igmpproxy &amp;&gt;/dev/null; then ${vyos_op_scripts_dir}/show_igmpproxy.py --mfc; else echo IGMP proxy not configured; fi</command>
+                <command>${vyos_op_scripts_dir}/igmp-proxy.py show_mfc</command>
               </leafNode>
               <leafNode name="summary">
                 <properties>
diff --git a/src/op_mode/show_igmpproxy.py b/src/op_mode/igmp-proxy.py
index 4714e494b..2a8f284bc 100755
--- a/src/op_mode/show_igmpproxy.py
+++ b/src/op_mode/igmp-proxy.py
@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 #
-# Copyright (C) 2018 VyOS maintainers and contributors
+# Copyright (C) 2023 VyOS maintainers and contributors
 #
 # This program is free software; you can redistribute it and/or modify
 # it under the terms of the GNU General Public License version 2 or later as
@@ -19,90 +19,43 @@
 #    Display istatistics from IPv4 IGMP proxy.
 #    Used by the "run show ip multicast" command tree.
 
-import sys
-import jinja2
-import argparse
 import ipaddress
+import json
+import jinja2
 import socket
+import sys
 
 import vyos.config
+import vyos.opmode
 
-# Output Template for "show ip multicast interface" command
-#
-# Example:
-# Interface  BytesIn      PktsIn       BytesOut     PktsOut      Local
-# eth0       0.0b         0            0.0b         0            xxx.xxx.xxx.65
-# eth1       0.0b         0            0.0b         0            xxx.xxx.xx.201
-# eth0.3     0.0b         0            0.0b         0            xxx.xxx.x.7
-# tun1       0.0b         0            0.0b         0            xxx.xxx.xxx.2
-vif_out_tmpl = """
-{% for r in data %}
-{{ "%-10s"|format(r.interface) }} {{ "%-12s"|format(r.bytes_in) }} {{ "%-12s"|format(r.pkts_in) }} {{ "%-12s"|format(r.bytes_out) }} {{ "%-12s"|format(r.pkts_out) }} {{ "%-15s"|format(r.loc) }}
-{% endfor %}
-"""
+from vyos.util import bytes_to_human, print_error
 
-# Output Template for "show ip multicast mfc" command
-#
-# Example:
-# Group             Origin            In    Out           Pkts        Bytes        Wrong
-# xxx.xxx.xxx.250   xxx.xx.xxx.75     --
-# xxx.xxx.xx.124    xx.xxx.xxx.26     --
-mfc_out_tmpl = """
-{% for r in data %}
-{{ "%-15s"|format(r.group) }} {{ "%-15s"|format(r.origin) }} {{ "%-12s"|format(r.pkts) }} {{ "%-12s"|format(r.bytes) }} {{ "%-12s"|format(r.wrong) }} {{ "%-10s"|format(r.iif) }} {{ "%-20s"|format(r.oifs|join(', ')) }}
-{% endfor %}
-"""
-
-parser = argparse.ArgumentParser()
-parser.add_argument("--interface", action="store_true", help="Interface Statistics")
-parser.add_argument("--mfc", action="store_true", help="Multicast Forwarding Cache")
-
-def byte_string(size):
-    # convert size to integer
-    size = int(size)
-
-    # One Terrabyte
-    s_TB = 1024 * 1024 * 1024 * 1024
-    # One Gigabyte
-    s_GB = 1024 * 1024 * 1024
-    # One Megabyte
-    s_MB = 1024 * 1024
-    # One Kilobyte
-    s_KB = 1024
-    # One Byte
-    s_B  = 1
-
-    if size > s_TB:
-        return str(round((size/s_TB), 2)) + 'TB'
-    elif size > s_GB:
-        return str(round((size/s_GB), 2)) + 'GB'
-    elif size > s_MB:
-        return str(round((size/s_MB), 2)) + 'MB'
-    elif size > s_KB:
-        return str(round((size/s_KB), 2)) + 'KB'
-    else:
-        return str(round((size/s_B), 2)) + 'b'
-
-    return None
-
-def kernel2ip(addr):
+def _is_configured():
+    """Check if IGMP proxy is configured"""
+    return vyos.config.Config().exists_effective('protocols igmp-proxy')
+
+def _is_running():
+    """Check if IGMP proxy is currently running"""
+    return not vyos.util.run('ps -C igmpproxy')
+
+def _kernel_to_ip(addr):
     """
-    Convert any given addr from Linux Kernel to a proper, IPv4 address
+    Convert any given address from Linux kernel to a proper, IPv4 address
     using the correct host byte order.
     """
 
     # Convert from hex 'FE000A0A' to decimal '4261415434'
     addr = int(addr, 16)
-    # Kernel ABI _always_ uses network byteorder
+    # Kernel ABI _always_ uses network byte order
     addr = socket.ntohl(addr)
 
-    return ipaddress.IPv4Address( addr )
+    return str(ipaddress.IPv4Address(addr))
 
-def do_mr_vif():
+def _process_mr_vif():
     """
     Read contents of file /proc/net/ip_mr_vif and print a more human
-    friendly version to the command line. IPv4 addresses present as
-    32bit integers in hex format are converted to IPv4 notation, too.
+    friendly version to the command line. IPv4 addresses presented as
+    32-bit integers in hex format are converted to IPv4 notation too.
     """
 
     with open('/proc/net/ip_mr_vif', 'r') as f:
@@ -135,22 +88,21 @@ def do_mr_vif():
                 'pkts_out' : line.split()[5],
 
                 # convert raw byte number to something more human readable
-                # Note: could be replaced by Python3 hurry.filesize module
-                'bytes_in' : byte_string( line.split()[2] ),
-                'bytes_out': byte_string( line.split()[4] ),
+                'bytes_in' : bytes_to_human(int(line.split()[2])),
+                'bytes_out': bytes_to_human(int(line.split()[4])),
 
                 # convert IP address from hex 'FE000A0A' to decimal '4261415434'
-                'loc'      : kernel2ip( line.split()[7] ),
+                'loc'      : _kernel_to_ip(line.split()[7]),
             }
             result['data'].append(data)
 
     return result
 
-def do_mr_mfc():
+def _process_mr_mfc():
     """
     Read contents of file /proc/net/ip_mr_cache and print a more human
-    friendly version to the command line. IPv4 addresses present as
-    32bit integers in hex format are converted to IPv4 notation, too.
+    friendly version to the command line. IPv4 addresses presented as
+    32-bit integers in hex format are converted to IPv4 notation too.
     """
 
     with open('/proc/net/ip_mr_cache', 'r') as f:
@@ -184,8 +136,8 @@ def do_mr_mfc():
         for line in f:
             data = {
                 # convert IP address from hex 'FE000A0A' to decimal '4261415434'
-                'group' : kernel2ip( line.split()[0] ),
-                'origin': kernel2ip( line.split()[1] ),
+                'group' : _kernel_to_ip(line.split()[0]),
+                'origin': _kernel_to_ip(line.split()[1]),
 
                 'iif'   : '--',
                 'pkts'  : '',
@@ -194,10 +146,10 @@ def do_mr_mfc():
                 'oifs'  : []
             }
 
-            iif = int( line.split()[2] )
+            iif = int(line.split()[2])
             if not ((iif == -1) or (iif == 65535)):
                 data['pkts']  = line.split()[3]
-                data['bytes'] = byte_string( line.split()[4] )
+                data['bytes'] = bytes_to_human(int(line.split()[4]))
                 data['wrong'] = line.split()[5]
 
                 # convert index to real interface name
@@ -205,37 +157,62 @@ def do_mr_mfc():
 
                 # convert each output interface index to a real interface name
                 for oif in line.split()[6:]:
-                    idx = int( oif.split(':')[0] )
-                    data['oifs'].append( vif[idx]['interface'] )
+                    idx = int(oif.split(':')[0])
+                    data['oifs'].append(vif[idx]['interface'])
 
             result['data'].append(data)
 
     return result
 
-if __name__ == '__main__':
-    args = parser.parse_args()
-
-    # Do nothing if service is not configured
-    c = vyos.config.Config()
-    if not c.exists_effective('protocols igmp-proxy'):
-        print("IGMP proxy is not configured")
-        sys.exit(0)
-
-    if args.interface:
-        data = do_mr_vif()
-        if data:
-            tmpl = jinja2.Template(vif_out_tmpl)
-            print(tmpl.render(data))
-
-        sys.exit(0)
-    elif args.mfc:
-        data = do_mr_mfc()
-        if data:
-            tmpl = jinja2.Template(mfc_out_tmpl)
-            print(tmpl.render(data))
-
-        sys.exit(0)
-    else:
-        parser.print_help()
-        sys.exit(1)
 
+# Output template for "show ip multicast interface" command
+#
+# Example:
+# Interface  BytesIn      PktsIn       BytesOut     PktsOut      Local
+# eth0       0.0 B        0            0.0 B        0            xxx.xxx.xxx.65
+# eth1       0.0 B        0            0.0 B        0            xxx.xxx.xx.201
+# eth0.3     0.0 B        0            0.0 B        0            xxx.xxx.x.7
+# tun1       0.0 B        0            0.0 B        0            xxx.xxx.xxx.2
+def show_interface(raw: bool):
+    vif_out_template = """{%- for r in data -%}
+{{ "%-10s"|format(r.interface) }} {{ "%-12s"|format(r.bytes_in) }} {{ "%-12s"|format(r.pkts_in) }} {{ "%-12s"|format(r.bytes_out) }} {{ "%-12s"|format(r.pkts_out) }} {{ "%-15s"|format(r.loc) }}
+{% endfor %}"""
+    if data := _process_mr_vif():
+        if raw:
+            return json.loads(json.dumps(data))
+        else:
+            return jinja2.Template(vif_out_template).render(data)
+
+# Output template for "show ip multicast mfc" command
+#
+# Example:
+# Group             Origin            In    Out           Pkts        Bytes        Wrong
+# xxx.xxx.xxx.250   xxx.xx.xxx.75     --
+# xxx.xxx.xx.124    xx.xxx.xxx.26     --
+def show_mfc(raw: bool):
+    mfc_out_template = """{%- for r in data -%}
+{{ "%-15s"|format(r.group) }} {{ "%-15s"|format(r.origin) }} {{ "%-12s"|format(r.pkts) }} {{ "%-12s"|format(r.bytes) }} {{ "%-12s"|format(r.wrong) }} {{ "%-10s"|format(r.iif) }} {{ "%-20s"|format(r.oifs|join(', ')) }}
+{% endfor %}"""
+    if data := _process_mr_mfc():
+        if raw:
+            return json.loads(json.dumps(data))
+        else:
+            return jinja2.Template(mfc_out_template).render(data)
+
+
+if not _is_configured():
+    print_error('IGMP proxy is not configured.')
+    sys.exit(0)
+if not _is_running():
+    print_error('IGMP proxy is not running.')
+    sys.exit(0)
+
+
+if __name__ == "__main__":
+    try:
+        res = vyos.opmode.run(sys.modules[__name__])
+        if res:
+            print(res)
+    except (ValueError, vyos.opmode.Error) as e:
+        print_error(e)
+        sys.exit(1)