fixed encrypted genshi/cheetah generators

6 files changed, 29 insertions, 30 deletions

diff --git a/src/lib/Bcfg2/Server/Plugins/Cfg/CfgCheetahGenerator.py b/src/lib/Bcfg2/Server/Plugins/Cfg/CfgCheetahGenerator.py
index 3edd1d8cb..e74b77e83 100644
--- a/src/lib/Bcfg2/Server/Plugins/Cfg/CfgCheetahGenerator.py
+++ b/src/lib/Bcfg2/Server/Plugins/Cfg/CfgCheetahGenerator.py
@@ -6,8 +6,7 @@ from Bcfg2.Server.Plugins.Cfg import CfgGenerator
 logger = logging.getLogger(__name__)
 
 try:
-    import Cheetah.Template
-    import Cheetah.Parser
+    from Cheetah.Template import Template
     have_cheetah = True
 except ImportError:
     have_cheetah = False
@@ -25,9 +24,8 @@ class CfgCheetahGenerator(CfgGenerator):
             raise Bcfg2.Server.Plugin.PluginExecutionError(msg)
 
     def get_data(self, entry, metadata):
-        template = Cheetah.Template.Template(self.data,
-                                             compilerSettings=self.settings)
+        template = Template(self.data, compilerSettings=self.settings)
         template.metadata = metadata
         template.path = entry.get('realname', entry.get('name'))
-        template.source_path = self.path
+        template.source_path = self.name
         return template.respond()
diff --git a/src/lib/Bcfg2/Server/Plugins/Cfg/CfgEncryptedCheetahGenerator.py b/src/lib/Bcfg2/Server/Plugins/Cfg/CfgEncryptedCheetahGenerator.py
index 3911cff62..a75329d2a 100644
--- a/src/lib/Bcfg2/Server/Plugins/Cfg/CfgEncryptedCheetahGenerator.py
+++ b/src/lib/Bcfg2/Server/Plugins/Cfg/CfgEncryptedCheetahGenerator.py
@@ -11,4 +11,4 @@ class CfgEncryptedCheetahGenerator(CfgCheetahGenerator, CfgEncryptedGenerator):
         CfgEncryptedGenerator.handle_event(self, event)
 
     def get_data(self, entry, metadata):
-        CfgCheetahGenerator.get_data(self, entry, metadata)
+        return CfgCheetahGenerator.get_data(self, entry, metadata)
diff --git a/src/lib/Bcfg2/Server/Plugins/Cfg/CfgEncryptedGenerator.py b/src/lib/Bcfg2/Server/Plugins/Cfg/CfgEncryptedGenerator.py
index 9b2db3100..0839e3536 100644
--- a/src/lib/Bcfg2/Server/Plugins/Cfg/CfgEncryptedGenerator.py
+++ b/src/lib/Bcfg2/Server/Plugins/Cfg/CfgEncryptedGenerator.py
@@ -18,12 +18,16 @@ def passphrases():
         return dict()
 
 def decrypt(crypted):
+    if not have_crypto:
+        msg = "Cfg: M2Crypto is not available: %s" % entry.get("name")
+        logger.error(msg)
+        raise Bcfg2.Server.Plugin.PluginExecutionError(msg)
     for passwd in passphrases().values():
         try:
             return ssl_decrypt(crypted, passwd)
         except EVPError:
             pass
-    raise EVPError("Failed to decrypt %s" % self.name)
+    raise EVPError("Failed to decrypt")
 
 class CfgEncryptedGenerator(CfgGenerator):
     __extensions__ = ["crypt"]
@@ -49,9 +53,9 @@ class CfgEncryptedGenerator(CfgGenerator):
         try:
             self.data = decrypt(crypted)
         except EVPError:
-            err = sys.exc_info()[1]
-            logger.error(err)
-            raise Bcfg2.Server.Plugin.PluginExecutionError(err)
+            msg = "Failed to decrypt %s" % self.name
+            logger.error(msg)
+            raise Bcfg2.Server.Plugin.PluginExecutionError(msg)
 
     def get_data(self, entry, metadata):
         if self.data is None:
diff --git a/src/lib/Bcfg2/Server/Plugins/Cfg/CfgEncryptedGenshiGenerator.py b/src/lib/Bcfg2/Server/Plugins/Cfg/CfgEncryptedGenshiGenerator.py
index 3845c438b..649bc38a5 100644
--- a/src/lib/Bcfg2/Server/Plugins/Cfg/CfgEncryptedGenshiGenerator.py
+++ b/src/lib/Bcfg2/Server/Plugins/Cfg/CfgEncryptedGenshiGenerator.py
@@ -1,6 +1,5 @@
 import logging
 from Bcfg2.Bcfg2Py3k import StringIO
-from Bcfg2.Server.Plugins.Cfg import SETUP
 from Bcfg2.Server.Plugins.Cfg.CfgGenshiGenerator import CfgGenshiGenerator
 from Bcfg2.Server.Plugins.Cfg.CfgEncryptedGenerator import decrypt, \
     CfgEncryptedGenerator
@@ -8,28 +7,20 @@ from Bcfg2.Server.Plugins.Cfg.CfgEncryptedGenerator import decrypt, \
 logger = logging.getLogger(__name__)
 
 try:
-    from genshi.template import TemplateLoader, loader
+    from genshi.template import TemplateLoader
 except ImportError:
     # CfgGenshiGenerator will raise errors if genshi doesn't exist
     pass
 
-def crypted_loader(filename):
-    loadfunc = loader.directory(os.path.dirname(filename))
-    filepath, filename, fileobj, uptodate = loadfunc(filename)
-    return (filepath, filename, StringIO(decrypt(fileobj.read())), uptodate)
-    
 
-class CfgEncryptedGenshiGenerator(CfgGenshiGenerator, CfgEncryptedGenerator):
-    __extensions__ = ['genshi.crypt', 'crypt.genshi']
-
-    def __init__(self, fname, spec, encoding):
-        CfgEncryptedGenerator.__init__(self, fname, spec, encoding)
-        CfgGenshiGenerator.__init__(self, fname, spec, encoding)
-        self.loader = TemplateLoader([crypted_loader])
+class EncryptedTemplateLoader(TemplateLoader):
+    def _instantiate(self, cls, fileobj, filepath, filename, encoding=None):
+        plaintext = StringIO(decrypt(fileobj.read()))
+        return TemplateLoader._instantiate(self, cls, plaintext, filepath,
+                                           filename, encoding=encoding)
+        
 
-    def handle_event(self, event):
-        CfgEncryptedGenerator.handle_event(self, event)
-        CfgGenshiGenerator.handle_event(self, event)
+class CfgEncryptedGenshiGenerator(CfgGenshiGenerator):
+    __extensions__ = ['genshi.crypt', 'crypt.genshi']
+    __loader_cls__ = EncryptedTemplateLoader
 
-    def get_data(self, entry, metadata):
-        CfgGenshiGenerator.get_data(self, entry, metadata)
diff --git a/src/lib/Bcfg2/Server/Plugins/Cfg/CfgGenshiGenerator.py b/src/lib/Bcfg2/Server/Plugins/Cfg/CfgGenshiGenerator.py
index 2c0a076d7..7c84648d7 100644
--- a/src/lib/Bcfg2/Server/Plugins/Cfg/CfgGenshiGenerator.py
+++ b/src/lib/Bcfg2/Server/Plugins/Cfg/CfgGenshiGenerator.py
@@ -23,10 +23,12 @@ def removecomment(stream):
 
 class CfgGenshiGenerator(CfgGenerator):
     __extensions__ = ['genshi']
+    __loader_cls__ = TemplateLoader
 
     def __init__(self, fname, spec, encoding):
         CfgGenerator.__init__(self, fname, spec, encoding)
-        self.loader = TemplateLoader()
+        self.loader = self.__loader_cls__()
+        self.template = None
         if not have_genshi:
             msg = "Cfg: Genshi is not available: %s" % entry.get("name")
             logger.error(msg)
diff --git a/src/lib/Bcfg2/Server/Plugins/Cfg/__init__.py b/src/lib/Bcfg2/Server/Plugins/Cfg/__init__.py
index 500c1e394..38400d984 100644
--- a/src/lib/Bcfg2/Server/Plugins/Cfg/__init__.py
+++ b/src/lib/Bcfg2/Server/Plugins/Cfg/__init__.py
@@ -287,6 +287,10 @@ class CfgEntrySet(Bcfg2.Server.Plugin.EntrySet):
                 logger.error("You need to specify base64 encoding for %s." %
                              entry.get('name'))
                 raise Bcfg2.Server.Plugin.PluginExecutionError(msg)
+            except TypeError:
+                # data is already unicode; newer versions of Cheetah
+                # seem to return unicode
+                pass
 
         if data:
             entry.text = data