Utils: Generalize safe_module_name

Other plugins that also import user generated python source code, may want to use
safe_module_name from TemplateHelper.

4 files changed, 33 insertions, 23 deletions

diff --git a/src/lib/Bcfg2/Server/Lint/TemplateHelper.py b/src/lib/Bcfg2/Server/Lint/TemplateHelper.py
index 9d05516f1..ce6fdca74 100644
--- a/src/lib/Bcfg2/Server/Lint/TemplateHelper.py
+++ b/src/lib/Bcfg2/Server/Lint/TemplateHelper.py
@@ -4,8 +4,8 @@
 import sys
 import imp
 from Bcfg2.Server.Lint import ServerPlugin
-from Bcfg2.Server.Plugins.TemplateHelper import HelperModule, MODULE_RE, \
-    safe_module_name
+from Bcfg2.Server.Plugins.TemplateHelper import HelperModule, MODULE_RE
+from Bcfg2.Utils import safe_module_name
 
 
 class TemplateHelper(ServerPlugin):
@@ -44,7 +44,8 @@ class TemplateHelper(ServerPlugin):
         module_name = MODULE_RE.search(helper).group(1)
 
         try:
-            module = imp.load_source(safe_module_name(module_name), helper)
+            module = imp.load_source(
+                safe_module_name('TemplateHelper', module_name), helper)
         except:  # pylint: disable=W0702
             err = sys.exc_info()[1]
             self.LintError("templatehelper-import-error",
diff --git a/src/lib/Bcfg2/Server/Plugins/TemplateHelper.py b/src/lib/Bcfg2/Server/Plugins/TemplateHelper.py
index cec2de297..b0b8d0061 100644
--- a/src/lib/Bcfg2/Server/Plugins/TemplateHelper.py
+++ b/src/lib/Bcfg2/Server/Plugins/TemplateHelper.py
@@ -7,18 +7,11 @@ import lxml.etree
 from Bcfg2.Server.Plugin import Plugin, Connector, DirectoryBacked, \
     TemplateDataProvider, DefaultTemplateDataProvider
 from Bcfg2.Logger import Debuggable
+from Bcfg2.Utils import safe_module_name
 
 MODULE_RE = re.compile(r'(?P<filename>(?P<module>[^\/]+)\.py)$')
 
 
-def safe_module_name(module):
-    """ Munge the name of a TemplateHelper module to avoid collisions
-    with other Python modules.  E.g., if someone has a helper named
-    'ldap.py', it should not be added to ``sys.modules`` as ``ldap``,
-    but rather as something more obscure. """
-    return '__TemplateHelper_%s' % module
-
-
 class HelperModule(Debuggable):
     """ Representation of a TemplateHelper module """
 
@@ -52,8 +45,9 @@ class HelperModule(Debuggable):
             return
 
         try:
-            module = imp.load_source(safe_module_name(self._module_name),
-                                     self.name)
+            module = imp.load_source(
+                safe_module_name('TemplateHelper', self._module_name),
+                self.name)
         except:  # pylint: disable=W0702
             # this needs to be a blanket except because the
             # imp.load_source() call can raise literally any error,
diff --git a/src/lib/Bcfg2/Utils.py b/src/lib/Bcfg2/Utils.py
index 64d0d8b93..2fdc0c3e0 100644
--- a/src/lib/Bcfg2/Utils.py
+++ b/src/lib/Bcfg2/Utils.py
@@ -321,6 +321,15 @@ def safe_input(msg):
     return input(msg)
 
 
+def safe_module_name(prefix, module):
+    """ Munge the name of a module with prefix to avoid collisions
+    with other Python modules.  E.g., if you want to import user
+    defined helper modules and someone has a helper named 'ldap.py',
+    it should not be added to ``sys.modules`` as ``ldap``, but rather
+    as something more obscure. """
+    return '__%s_%s' % (prefix, module)
+
+
 class classproperty(object):  # pylint: disable=C0103
     """ Decorator that can be used to create read-only class
     properties. """
diff --git a/testsuite/Testsrc/Testlib/TestServer/TestPlugins/TestTemplateHelper.py b/testsuite/Testsrc/Testlib/TestServer/TestPlugins/TestTemplateHelper.py
index 128d6cae5..1c2fa23a2 100644
--- a/testsuite/Testsrc/Testlib/TestServer/TestPlugins/TestTemplateHelper.py
+++ b/testsuite/Testsrc/Testlib/TestServer/TestPlugins/TestTemplateHelper.py
@@ -2,6 +2,7 @@ import os
 import sys
 import Bcfg2.Server.Plugin
 from mock import Mock, MagicMock, patch
+from Bcfg2.Utils import safe_module_name
 from Bcfg2.Server.Plugins.TemplateHelper import *
 
 # add all parent testsuite directories to sys.path to allow (most)
@@ -39,8 +40,9 @@ class TestHelperModule(Bcfg2TestCase):
         mock_load_source.side_effect = ImportError
         attrs = dir(hm)
         hm.HandleEvent()
-        mock_load_source.assert_called_with(safe_module_name(hm._module_name),
-                                            hm.name)
+        mock_load_source.assert_called_with(
+            safe_module_name('TemplateHelper', hm._module_name),
+            hm.name)
         self.assertEqual(attrs, dir(hm))
         self.assertEqual(hm._attrs, [])
 
@@ -51,8 +53,9 @@ class TestHelperModule(Bcfg2TestCase):
         mock_load_source.return_value = Mock()
         attrs = dir(hm)
         hm.HandleEvent()
-        mock_load_source.assert_called_with(safe_module_name(hm._module_name),
-                                            hm.name)
+        mock_load_source.assert_called_with(
+            safe_module_name('TemplateHelper', hm._module_name),
+            hm.name)
         self.assertEqual(attrs, dir(hm))
         self.assertEqual(hm._attrs, [])
 
@@ -63,8 +66,9 @@ class TestHelperModule(Bcfg2TestCase):
         mock_load_source.return_value = module
         attrs = dir(hm)
         hm.HandleEvent()
-        mock_load_source.assert_called_with(safe_module_name(hm._module_name),
-                                            hm.name)
+        mock_load_source.assert_called_with(
+            safe_module_name('TemplateHelper', hm._module_name),
+            hm.name)
         self.assertEqual(attrs, dir(hm))
         self.assertEqual(hm._attrs, [])
 
@@ -74,8 +78,9 @@ class TestHelperModule(Bcfg2TestCase):
         mock_load_source.reset()
         mock_load_source.return_value = module
         hm.HandleEvent()
-        mock_load_source.assert_called_with(safe_module_name(hm._module_name),
-                                            hm.name)
+        mock_load_source.assert_called_with(
+            safe_module_name('TemplateHelper', hm._module_name),
+            hm.name)
         self.assertTrue(hasattr(hm, "foo"))
         self.assertTrue(hasattr(hm, "bar"))
         self.assertTrue(hasattr(hm, "baz"))
@@ -87,8 +92,9 @@ class TestHelperModule(Bcfg2TestCase):
         mock_load_source.reset()
         mock_load_source.return_value = module
         hm.HandleEvent()
-        mock_load_source.assert_called_with(safe_module_name(hm._module_name),
-                                            hm.name)
+        mock_load_source.assert_called_with(
+            safe_module_name('TemplateHelper', hm._module_name),
+            hm.name)
         self.assertTrue(hasattr(hm, "foo"))
         self.assertTrue(hasattr(hm, "bar"))
         self.assertTrue(hasattr(hm, "quux"))