Merge pull request #1425 from sever-sever/T4145

conntrack: T4145: Modify conntrack to format command runner

2 files changed, 38 insertions, 13 deletions

diff --git a/op-mode-definitions/show-conntrack.xml.in b/op-mode-definitions/show-conntrack.xml.in
index 792623d7d..8d921e6a5 100644
--- a/op-mode-definitions/show-conntrack.xml.in
+++ b/op-mode-definitions/show-conntrack.xml.in
@@ -16,7 +16,13 @@
                 <properties>
                   <help>Show conntrack entries for IPv4 protocol</help>
                 </properties>
-                <command>sudo ${vyos_op_scripts_dir}/show_conntrack.py</command>
+                <command>sudo ${vyos_op_scripts_dir}/conntrack.py show --family inet</command>
+              </node>
+              <node name="ipv6">
+                <properties>
+                  <help>Show conntrack entries for IPv6 protocol</help>
+                </properties>
+                <command>sudo ${vyos_op_scripts_dir}/conntrack.py show --family inet6</command>
               </node>
             </children>
           </node>
diff --git a/src/op_mode/show_conntrack.py b/src/op_mode/conntrack.py
index 089a3e454..1441d110f 100755
--- a/src/op_mode/show_conntrack.py
+++ b/src/op_mode/conntrack.py
@@ -14,17 +14,21 @@
 # You should have received a copy of the GNU General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
+import sys
 import xmltodict
 
 from tabulate import tabulate
 from vyos.util import cmd
+from vyos.util import run
 
+import vyos.opmode
 
-def _get_raw_data():
+
+def _get_xml_data(family):
     """
     Get conntrack XML output
     """
-    return cmd(f'sudo  conntrack --dump --output xml')
+    return cmd(f'sudo conntrack --dump --family {family} --output xml')
 
 
 def _xml_to_dict(xml):
@@ -32,26 +36,34 @@ def _xml_to_dict(xml):
     Convert XML to dictionary
     Return: dictionary
     """
-    parse = xmltodict.parse(xml)
+    parse = xmltodict.parse(xml, attr_prefix='')
     # If only one conntrack entry we must change dict
     if 'meta' in parse['conntrack']['flow']:
         return dict(conntrack={'flow': [parse['conntrack']['flow']]})
     return parse
 
 
-def _get_formatted_output(xml):
+def _get_raw_data(family):
+    """
+    Return: dictionary
+    """
+    xml = _get_xml_data(family)
+    return _xml_to_dict(xml)
+
+
+def get_formatted_output(dict_data):
     """
     :param xml:
     :return: formatted output
     """
     data_entries = []
-    dict_data = _xml_to_dict(xml)
+    #dict_data = _get_raw_data(family)
     for entry in dict_data['conntrack']['flow']:
         orig_src, orig_dst, orig_sport, orig_dport = {}, {}, {}, {}
         reply_src, reply_dst, reply_sport, reply_dport = {}, {}, {}, {}
         proto = {}
         for meta in entry['meta']:
-            direction = meta['@direction']
+            direction = meta['direction']
             if direction in ['original']:
                 if 'layer3' in meta:
                     orig_src = meta['layer3']['src']
@@ -61,7 +73,7 @@ def _get_formatted_output(xml):
                         orig_sport = meta['layer4']['sport']
                     if meta.get('layer4').get('dport'):
                         orig_dport = meta['layer4']['dport']
-                    proto = meta['layer4']['@protoname']
+                    proto = meta['layer4']['protoname']
             if direction in ['reply']:
                 if 'layer3' in meta:
                     reply_src = meta['layer3']['src']
@@ -71,7 +83,7 @@ def _get_formatted_output(xml):
                         reply_sport = meta['layer4']['sport']
                     if meta.get('layer4').get('dport'):
                         reply_dport = meta['layer4']['dport']
-                    proto = meta['layer4']['@protoname']
+                    proto = meta['layer4']['protoname']
             if direction == 'independent':
                 conn_id = meta['id']
                 timeout = meta['timeout']
@@ -90,13 +102,20 @@ def _get_formatted_output(xml):
     return output
 
 
-def show(raw: bool):
-    conntrack_data = _get_raw_data()
+def show(raw: bool, family: str):
+    family = 'ipv6' if family == 'inet6' else 'ipv4'
+    conntrack_data = _get_raw_data(family)
     if raw:
         return conntrack_data
     else:
-        return _get_formatted_output(conntrack_data)
+        return get_formatted_output(conntrack_data)
 
 
 if __name__ == '__main__':
-    print(show(raw=False))
+    try:
+        res = vyos.opmode.run(sys.modules[__name__])
+        if res:
+            print(res)
+    except ValueError as e:
+        print(e)
+        sys.exit(1)