fixed encrypted genshi/cheetah generators

1 files changed, 10 insertions, 19 deletions

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)