From 62a1f5eb19dd8994d2dfaf66b31829fa622d0b2a Mon Sep 17 00:00:00 2001
From: kellyyeh <42761586+kellyyeh@users.noreply.github.com>
Date: Thu, 23 Sep 2021 22:01:26 -0700
Subject: [PATCH] Add CLI Support for IPv6 Helpers and DHCPv6 Relay Counters
 (#8593)

---
 .../cli-plugin-tests/mock_config.py           |  18 +++
 .../test_show_dhcp6relay_counters.py          |  43 +++++++
 .../test_show_dhcpv6_helper.py                |  40 +++++++
 .../clear/plugins/clear_dhcp6relay_counter.py |  32 +++++
 .../cli/show/plugins/show_dhcp_relay.py       | 113 +++++++++++++++++-
 files/build_templates/manifest.json.j2        |   3 +-
 rules/docker-dhcp-relay.mk                    |   1 +
 rules/functions                               |   1 +
 sonic-slave-buster/Dockerfile.j2              |   4 +
 sonic-slave-stretch/Dockerfile.j2             |   4 +
 10 files changed, 257 insertions(+), 2 deletions(-)
 create mode 100644 dockers/docker-dhcp-relay/cli-plugin-tests/mock_config.py
 create mode 100644 dockers/docker-dhcp-relay/cli-plugin-tests/test_show_dhcp6relay_counters.py
 create mode 100644 dockers/docker-dhcp-relay/cli-plugin-tests/test_show_dhcpv6_helper.py
 create mode 100644 dockers/docker-dhcp-relay/cli/clear/plugins/clear_dhcp6relay_counter.py

diff --git a/dockers/docker-dhcp-relay/cli-plugin-tests/mock_config.py b/dockers/docker-dhcp-relay/cli-plugin-tests/mock_config.py
new file mode 100644
index 000000000000..ed04367fbba7
--- /dev/null
+++ b/dockers/docker-dhcp-relay/cli-plugin-tests/mock_config.py
@@ -0,0 +1,18 @@
+TEST_DATA = [ 
+    [
+        "DHCPv6_Helpers",
+        {
+            "config_db": {
+                "DHCP_RELAY": {
+                    "Vlan1000": {
+                        "dhcpv6_servers": [
+                            "fc02:2000::1",
+                            "fc02:2000::2"
+                        ],
+                        "dhcpv6_option|rfc6939_support": "true"
+                    }
+                }
+            },
+        },
+    ],
+]
diff --git a/dockers/docker-dhcp-relay/cli-plugin-tests/test_show_dhcp6relay_counters.py b/dockers/docker-dhcp-relay/cli-plugin-tests/test_show_dhcp6relay_counters.py
new file mode 100644
index 000000000000..f640ef1de6ce
--- /dev/null
+++ b/dockers/docker-dhcp-relay/cli-plugin-tests/test_show_dhcp6relay_counters.py
@@ -0,0 +1,43 @@
+import sys
+import os
+from unittest import mock
+sys.path.append('../cli/show/plugins/')
+import show_dhcp_relay as show
+
+from click.testing import CliRunner
+
+try:
+    modules_path = os.path.join(os.path.dirname(__file__), "../../../src/sonic-utilities")
+    test_path = os.path.join(modules_path, "tests")
+    mock_table_path = os.path.join(test_path, "mock_tables")
+    sys.path.insert(0, modules_path)
+    sys.path.insert(0, test_path)
+    sys.path.insert(0, mock_table_path)
+    import dbconnector
+except KeyError:
+    pass
+
+expected_counts = """\
+  Message Type     Vlan1000
+--------------  -----------
+       Solicit
+     Advertise
+       Request
+       Confirm
+         Renew
+        Rebind
+         Reply
+       Release
+       Decline
+ Relay-Forward
+   Relay-Reply
+
+"""
+
+class TestDhcp6RelayCounters(object):
+
+    def test_show_counts(self):           
+        runner = CliRunner()
+        result = runner.invoke(show.dhcp6relay_counters.commands["counts"], ["-i Vlan1000"])
+        assert result.output == expected_counts
+
diff --git a/dockers/docker-dhcp-relay/cli-plugin-tests/test_show_dhcpv6_helper.py b/dockers/docker-dhcp-relay/cli-plugin-tests/test_show_dhcpv6_helper.py
new file mode 100644
index 000000000000..1f079c5ac965
--- /dev/null
+++ b/dockers/docker-dhcp-relay/cli-plugin-tests/test_show_dhcpv6_helper.py
@@ -0,0 +1,40 @@
+import pytest
+import sys
+import os
+sys.path.append('../cli/show/plugins/')
+import show_dhcp_relay as show
+from click.testing import CliRunner
+from swsscommon import swsscommon
+from mock_config import TEST_DATA
+from parameterized import parameterized
+from pyfakefs.fake_filesystem_unittest import patchfs
+
+try:
+    sys.path.insert(0, '../../../src/sonic-host-services/tests/common')
+    from mock_configdb import MockConfigDb
+    swsscommon.ConfigDBConnector = MockConfigDb 
+except KeyError:
+    pass
+    
+expected_table = """\
+--------  ------------
+Vlan1000  fc02:2000::1
+          fc02:2000::2
+--------  ------------
+"""
+
+DBCONFIG_PATH = '/var/run/redis/sonic-db/database_config.json'
+
+class TestDhcpRelayHelper(object):
+
+    @parameterized.expand(TEST_DATA)
+    @patchfs
+    def test_show_dhcpv6_helper(self, test_name, test_data, fs):
+        if not os.path.exists(DBCONFIG_PATH):
+            fs.create_file(DBCONFIG_PATH) 
+        MockConfigDb.set_config_db(test_data["config_db"])
+        runner = CliRunner()
+        table = MockConfigDb.get_table(self, "DHCP_RELAY")
+        result = show.get_data(table, "Vlan1000")
+        assert result == expected_table
+
diff --git a/dockers/docker-dhcp-relay/cli/clear/plugins/clear_dhcp6relay_counter.py b/dockers/docker-dhcp-relay/cli/clear/plugins/clear_dhcp6relay_counter.py
new file mode 100644
index 000000000000..7b0d5097d2ec
--- /dev/null
+++ b/dockers/docker-dhcp-relay/cli/clear/plugins/clear_dhcp6relay_counter.py
@@ -0,0 +1,32 @@
+import sys
+import click
+sys.path.insert(0, '../../show/plugins/')
+from show_dhcp6relay_counters import DHCPv6_Counter
+
+import utilities_common.cli as clicommon
+
+
+# sonic-clear dhcp6relay_counters
+@click.group(cls=clicommon.AliasedGroup)
+def dhcp6relay_clear():
+    pass
+
+@dhcp6relay_clear.command('dhcp6relay_counters')
+@click.option('-i', '--interface', required=False)
+def dhcp6relay_clear_counters(interface):
+    """ Clear dhcp6relay message counts """
+    
+    counter = DHCPv6_Counter()
+    counter_intf = counter.get_interface()
+
+    if interface:
+        counter.clear_table(interface)
+    else:
+        for intf in counter_intf:
+                counter.clear_table(intf)
+
+def register(cli):
+    cli.add_command(dhcp6relay_clear_counters)
+
+if __name__ == '__main__':
+    dhcp6relay_clear_counters()
diff --git a/dockers/docker-dhcp-relay/cli/show/plugins/show_dhcp_relay.py b/dockers/docker-dhcp-relay/cli/show/plugins/show_dhcp_relay.py
index 95698f276463..bad1de7d52bd 100644
--- a/dockers/docker-dhcp-relay/cli/show/plugins/show_dhcp_relay.py
+++ b/dockers/docker-dhcp-relay/cli/show/plugins/show_dhcp_relay.py
@@ -1,5 +1,22 @@
+import click
 from natsort import natsorted
+from tabulate import tabulate
 import show.vlan as vlan
+import utilities_common.cli as clicommon
+
+from swsscommon.swsscommon import ConfigDBConnector
+from swsscommon.swsscommon import SonicV2Connector
+
+
+# STATE_DB Table
+DHCPv6_COUNTER_TABLE = 'DHCPv6_COUNTER_TABLE'
+
+# DHCPv6 Counter Messages
+messages = ["Solicit", "Advertise", "Request", "Confirm", "Renew", "Rebind", "Reply", "Release", "Decline", "Relay-Forward", "Relay-Reply"]
+
+# DHCP_RELAY Config Table
+DHCP_RELAY = 'DHCP_RELAY'
+config_db = ConfigDBConnector()
 
 def get_dhcp_helper_address(ctx, vlan):
     cfg, _ = ctx
@@ -17,5 +34,99 @@ def get_dhcp_helper_address(ctx, vlan):
 vlan.VlanBrief.register_column('DHCP Helper Address', get_dhcp_helper_address)
 
 
-def register(cli):
+class DHCPv6_Counter(object):
+    def __init__(self):
+        self.db = SonicV2Connector(use_unix_socket_path=False)
+        self.db.connect(self.db.STATE_DB)
+        self.table_name = DHCPv6_COUNTER_TABLE + self.db.get_db_separator(self.db.STATE_DB)
+
+
+    def get_interface(self):
+        """ Get all names of all interfaces in DHCPv6_COUNTER_TABLE """
+        vlans = []
+        for key in self.db.keys(self.db.STATE_DB):
+            if DHCPv6_COUNTER_TABLE in key:
+                vlans.append(key[21:])
+        return vlans
+        
+
+    def get_dhcp6relay_msg_count(self, interface, msg):
+        """ Get count of a dhcp6relay message """
+        count = self.db.get(self.db.STATE_DB, self.table_name + str(interface), str(msg))
+        data = [str(msg), count]
+        return data
+
+    
+    def clear_table(self, interface):
+        """ Reset all message counts to 0 """
+        for msg in messages:
+            self.db.set(self.db.STATE_DB, self.table_name + str(interface), str(msg), '0') 
+
+def print_count(counter, intf):
+    """Print count of each message"""
+    data = []
+    for i in messages:
+        data.append(counter.get_dhcp6relay_msg_count(intf, i))
+    print(tabulate(data, headers = ["Message Type", intf], tablefmt='simple', stralign='right') + "\n")
+
+
+#
+# 'dhcp6relay_counters' group ###
+#
+
+
+@click.group(cls=clicommon.AliasedGroup, name="dhcp6relay_counters")
+def dhcp6relay_counters():
+    """Show DHCPv6 counter"""
     pass
+
+
+# 'counts' subcommand ("show dhcp6relay_counters counts")
+@dhcp6relay_counters.command('counts')
+@click.option('-i', '--interface', required=False)
+@click.option('--verbose', is_flag=True, help="Enable verbose output")
+def counts(interface, verbose):
+    """Show dhcp6relay message counts"""
+
+    counter = DHCPv6_Counter()
+    counter_intf = counter.get_interface()
+
+    if interface:
+        print_count(counter, interface)
+    else:
+        for intf in counter_intf:
+                print_count(counter, intf)
+
+
+
+@click.group(cls=clicommon.AliasedGroup, name="dhcprelay_helper")
+def dhcp_relay_helper():
+    """Show DHCP_Relay helper information"""
+    pass
+
+@dhcp_relay_helper.command('ipv6')
+def get_dhcpv6_helper_address():
+    """Parse through DHCP_RELAY table for each interface in config_db.json and print dhcpv6 helpers in table format"""
+    if config_db is not None:
+        config_db.connect()
+        table_data = config_db.get_table(DHCP_RELAY)
+        if table_data is not None:
+            vlans = config_db.get_keys(DHCP_RELAY)
+            for vlan in vlans: 
+                output = get_data(table_data, vlan)
+                print(output)
+            
+
+def get_data(table_data, vlan):
+    vlan_data = table_data.get(vlan)
+    helpers_data = vlan_data.get('dhcpv6_servers')
+    if helpers_data is not None:
+        addr = {vlan:[]}
+        for ip in helpers_data:
+            addr[vlan].append(ip)
+    output = tabulate({'Interface':[vlan], vlan:addr.get(vlan)}, tablefmt='simple', stralign='right') + '\n'
+    return output
+
+def register(cli):
+    cli.add_command(dhcp6relay_counters)
+    cli.add_command(dhcp_relay_helper)
diff --git a/files/build_templates/manifest.json.j2 b/files/build_templates/manifest.json.j2
index 9f3e872d2e65..431b2dd22432 100644
--- a/files/build_templates/manifest.json.j2
+++ b/files/build_templates/manifest.json.j2
@@ -29,6 +29,7 @@
     },
     "cli": {
         "config": "{{ config_cli_plugin|default('') }}",
-        "show": "{{ show_cli_plugin|default('') }}"
+        "show": "{{ show_cli_plugin|default('') }}",
+        "clear": "{{ clear_cli_plugin|default('') }}"
     }
 }
diff --git a/rules/docker-dhcp-relay.mk b/rules/docker-dhcp-relay.mk
index c3232e6da27a..42d7a3f86b33 100644
--- a/rules/docker-dhcp-relay.mk
+++ b/rules/docker-dhcp-relay.mk
@@ -53,5 +53,6 @@ $(DOCKER_DHCP_RELAY)_CONTAINER_TMPFS += /var/tmp/
 
 $(DOCKER_DHCP_RELAY)_CLI_CONFIG_PLUGIN = /cli/config/plugins/dhcp_relay.py
 $(DOCKER_DHCP_RELAY)_CLI_SHOW_PLUGIN = /cli/show/plugins/show_dhcp_relay.py
+$(DOCKER_DHCP_RELAY)_CLI_CLEAR_PLUGIN = /cli/clear/plugins/clear_dhcp6relay_counter.py
 
 $(DOCKER_DHCP_RELAY)_FILES += $(SUPERVISOR_PROC_EXIT_LISTENER_SCRIPT)
diff --git a/rules/functions b/rules/functions
index 5873b2db1b87..6f62b87797c8 100644
--- a/rules/functions
+++ b/rules/functions
@@ -196,6 +196,7 @@ define generate_manifest
     $(eval export tmpfs=$($(1).gz_CONTAINER_TMPFS))
     $(eval export config_cli_plugin=$($(1).gz_CLI_CONFIG_PLUGIN))
     $(eval export show_cli_plugin=$($(1).gz_CLI_SHOW_PLUGIN))
+    $(eval export clear_cli_plugin=$($(1).gz_CLI_CLEAR_PLUGIN))
     j2 $($*.gz_PATH)/Dockerfile$(2).j2 > $($(1).gz_PATH)/Dockerfile$(2)
     j2 --customize scripts/j2cli/json_filter.py files/build_templates/manifest.json.j2 > $($(1).gz_PATH)/manifest.common.json
     if [ -f $($*.gz_PATH)/manifest.part.json.j2 ]; then
diff --git a/sonic-slave-buster/Dockerfile.j2 b/sonic-slave-buster/Dockerfile.j2
index bbd68a31396e..f23c1aeb8c17 100644
--- a/sonic-slave-buster/Dockerfile.j2
+++ b/sonic-slave-buster/Dockerfile.j2
@@ -486,6 +486,10 @@ EXPOSE 22
 RUN git clone https://chromium.googlesource.com/chromium/tools/depot_tools.git /usr/share/depot_tools
 ENV PATH /usr/share/depot_tools:$PATH
 
+# Install dependencies for dhcp relay test 
+RUN pip3 install parameterized==0.8.1
+RUN pip3 install pyfakefs
+
 # Install docker engine 17.03.2~ce-0 inside docker and enable experimental feature
 RUN apt-get update
 RUN apt-get install -y \
diff --git a/sonic-slave-stretch/Dockerfile.j2 b/sonic-slave-stretch/Dockerfile.j2
index cc45afbc6f82..e077c69e01b3 100644
--- a/sonic-slave-stretch/Dockerfile.j2
+++ b/sonic-slave-stretch/Dockerfile.j2
@@ -291,6 +291,10 @@ RUN apt-get update && apt-get install -y \
         python-lxml \
         libexpat1-dev
 
+# Install dependencies for dhcp relay test 
+RUN pip3 install parameterized==0.8.1
+RUN pip3 install pyfakefs
+
 ## Config dpkg
 ## install the configuration file if it’s currently missing
 RUN sudo augtool --autosave "set /files/etc/dpkg/dpkg.cfg/force-confmiss"