T2614: add a key mangling option to vyos.config.get_config_dict()

3 files changed, 69 insertions, 2 deletions

diff --git a/python/vyos/config.py b/python/vyos/config.py
index f94293e2d..56353c322 100644
--- a/python/vyos/config.py
+++ b/python/vyos/config.py
@@ -68,6 +68,7 @@ import re
 import json
 import subprocess
 
+import vyos.util
 import vyos.configtree
 
 
@@ -287,7 +288,7 @@ class Config(object):
             self.__session_env = save_env
             return(default)
 
-    def get_config_dict(self, path=[], effective=False):
+    def get_config_dict(self, path=[], effective=False, key_mangling=None):
         """
         Args: path (str list): Configuration tree path, can be empty
         Returns: a dict representation of the config
@@ -299,6 +300,15 @@ class Config(object):
         else:
             config_dict = {}
 
+        if key_mangling:
+            if not (isinstance(key_mangling, tuple) and \
+                    (len(key_mangling) == 2) and \
+                    isinstance(key_mangling[0], str) and \
+                    isinstance(key_mangling[1], str)):
+                raise ValueError("key_mangling must be a tuple of two strings")
+            else:
+                config_dict = vyos.util.mangle_dict_keys(config_dict, key_mangling[0], key_mangling[1])
+
         return config_dict
 
     def is_multi(self, path):
diff --git a/python/vyos/util.py b/python/vyos/util.py
index c93f8d3f0..0ddc14963 100644
--- a/python/vyos/util.py
+++ b/python/vyos/util.py
@@ -1,4 +1,4 @@
-# Copyright 2019 VyOS maintainers and contributors <maintainers@vyos.io>
+# Copyright 2020 VyOS maintainers and contributors <maintainers@vyos.io>
 #
 # This library is free software; you can redistribute it and/or
 # modify it under the terms of the GNU Lesser General Public
@@ -14,6 +14,7 @@
 # License along with this library.  If not, see <http://www.gnu.org/licenses/>.
 
 import os
+import re
 import sys
 
 #
@@ -340,6 +341,28 @@ def colon_separated_to_dict(data_string, uniquekeys=False):
 
     return data
 
+def mangle_dict_keys(data, regex, replacement):
+    """ Mangles dict keys according to a regex and replacement character.
+    Some libraries like Jinja2 do not like certain characters in dict keys.
+    This function can be used for replacing all offending characters
+    with something acceptable.
+
+    Args:
+        data (dict): Original dict to mangle
+
+    Returns: dict
+    """
+    new_dict = {}
+    for key in data.keys():
+        new_key = re.sub(regex, replacement, key)
+
+        value = data[key]
+        if isinstance(value, dict):
+            new_dict[new_key] = mangle_dict_keys(value, regex, replacement)
+        else:
+            new_dict[new_key] = value
+
+    return new_dict
 
 def process_running(pid_file):
     """ Checks if a process with PID in pid_file is running """
diff --git a/src/tests/test_util.py b/src/tests/test_util.py
new file mode 100644
index 000000000..0e56a67a8
--- /dev/null
+++ b/src/tests/test_util.py
@@ -0,0 +1,34 @@
+#!/usr/bin/env python3
+#
+# Copyright (C) 2020 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
+# published by the Free Software Foundation.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU General Public License for more details.
+#
+# 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 unittest
+from unittest import TestCase
+
+import vyos.util
+
+
+class TestVyOSUtil(TestCase):
+    def setUp(self):
+        pass
+
+    def test_key_mangline(self):
+        data = {"foo-bar": {"baz-quux": None}}
+        expected_data = {"foo_bar": {"baz_quux": None}}
+        new_data = vyos.util.mangle_dict_keys(data, '-', '_')
+        self.assertEqual(new_data, expected_data)
+