make Bcfg2 automatically handle JSON, XML, and YAML output from probes

2 files changed, 126 insertions, 10 deletions

diff --git a/doc/server/plugins/probes/index.txt b/doc/server/plugins/probes/index.txt
index b9f698a0f..d52e7bb95 100644
--- a/doc/server/plugins/probes/index.txt
+++ b/doc/server/plugins/probes/index.txt
@@ -85,10 +85,12 @@ file looks like::
     #end if
 
 Any Probe script you run will store its output in
-``$self.metadata.Probes["scriptname"]``, so we get to our `scratchlocal`
-script's output as seen above.  Note that we had to wrap the output in an
-`int()` call; the script output is treated as a string, so it needs to
-be converted before it can be tested numerically.
+``$self.metadata.Probes["scriptname"]``, so we get to our
+`scratchlocal` script's output as seen above.  (See `Handling Probe
+Output`_, below, for more information on how this is done.)  Note that
+we had to wrap the output in an `int()` call; the script output is
+treated as a string, so it needs to be converted before it can be
+tested numerically.
 
 With all of these pieces in place, the following series of events will
 happen when the client is run:
@@ -110,8 +112,37 @@ is to add the ``/etc/auto.master`` to a Bundle:
 
     <Path name='/etc/auto.master'/>
 
-Host and Group Specific probes
-==============================
+Handling Probe Output
+=====================
+
+Bcfg2 stores output from probes in the ``Probes`` property of a
+client's metadata object.  To access this data in TGenshi, for
+instance, you could do::
+
+  ${metadata.Probes['script-name']}
+
+This is not the full output of the probe; any lines that start with
+"group:" have been stripped from the output.  The data is a
+string-like object that has some interesting and salient features:
+
+* If the data is a valid XML document, then
+  ``metadata.Probes['script-name'].xdata`` will be an
+  ``lxml.etree._Element`` object representing the XML data.
+* If the data is a valid JSON document, and the Python ``json`` module
+  is installed (included in Python 2.6 onward), then
+  ``metadata.Probes['script-name'].json`` will be a data structure
+  representing the JSON data.
+* If the data is a valid YAML document, and either the Python ``yaml``
+  module or ``syck`` module is installed, then
+  ``metadata.Probes['script-name'].yaml`` will be a data structure
+  representing the YAML data.
+
+If these conditions are not met, then the named properties will be
+``None``.  In all other fashions, the probe data objects should act
+like strings.
+
+Host- and Group-Specific probes
+===============================
 
 Bcfg2 has the ability to alter probes based on client hostname and group
 membership. These files work similarly to files in Cfg.
diff --git a/src/lib/Server/Plugins/Probes.py b/src/lib/Server/Plugins/Probes.py
index 3599f2af1..4a7171297 100644
--- a/src/lib/Server/Plugins/Probes.py
+++ b/src/lib/Server/Plugins/Probes.py
@@ -2,12 +2,96 @@ import lxml.etree
 import operator
 import re
 
+try:
+    import json
+    has_json = True
+except ImportError:
+    has_json = False
+
+try:
+    import syck
+    has_syck = True
+except ImportError:
+    has_syck = False
+
+try:
+    import yaml
+    has_yaml = True
+except ImportError:
+    has_yaml = False
+
 import Bcfg2.Server.Plugin
 
 specific_probe_matcher = re.compile("(.*/)?(?P<basename>\S+)(.(?P<mode>[GH](\d\d)?)_\S+)")
 probe_matcher = re.compile("(.*/)?(?P<basename>\S+)")
 
 
+class ProbeData (object):
+    """ a ProbeData object emulates a str object, but also has .xdata
+    and .json properties to provide convenient ways to use ProbeData
+    objects as XML or JSON data """
+    def __init__(self, data):
+        self.data = data
+        self._xdata = None
+        self._json = None
+        self._yaml = None
+
+    def __str__(self):
+        return str(self.data)
+    
+    def __repr__(self):
+        return repr(self.data)
+    
+    def __getattr__(self, name):
+        """ make ProbeData act like a str object """
+        return getattr(self.data, name)
+
+    def __complex__(self):
+        return complex(self.data)
+
+    def __int__(self):
+        return int(self.data)
+
+    def __long__(self):
+        return long(self.data)
+
+    def __float__(self):
+        return float(self.data)
+
+    @property
+    def xdata(self):
+        if self._xdata is None:
+            try:
+                self._xdata = lxml.etree.XML(self.data)
+            except lxml.etree.XMLSyntaxError:
+                pass
+        return self._xdata
+
+    @property
+    def json(self):
+        if self._json is None and has_json:
+            try:
+                self._json = json.loads(self.data)
+            except ValueError:
+                pass
+        return self._json
+
+    @property
+    def yaml(self):
+        if self._yaml is None:
+            if has_yaml:
+                try:
+                    self._yaml = yaml.load(self.data)
+                except yaml.YAMLError:
+                    pass
+            elif has_syck:
+                try:
+                    self._yaml = syck.load(self.data)
+                except syck.error:
+                    pass
+        return self._yaml
+
+
 class ProbeSet(Bcfg2.Server.Plugin.EntrySet):
     ignore = re.compile("^(\.#.*|.*~|\\..*\\.(tmp|sw[px])|probed\\.xml)$")
 
@@ -106,7 +190,8 @@ class Probes(Bcfg2.Server.Plugin.Plugin,
             self.cgroups[client.get('name')] = []
             for pdata in client:
                 if (pdata.tag == 'Probe'):
-                    self.probedata[client.get('name')][pdata.get('name')] = pdata.get('value')
+                    self.probedata[client.get('name')][pdata.get('name')] = \
+                        ProbeData(pdata.get('value'))
                 elif (pdata.tag == 'Group'):
                     self.cgroups[client.get('name')].append(pdata.get('name'))
 
@@ -142,11 +227,11 @@ class Probes(Bcfg2.Server.Plugin.Plugin,
                 if newgroup not in self.cgroups[client.hostname]:
                     self.cgroups[client.hostname].append(newgroup)
                 dlines.remove(line)
-        dtext = "\n".join(dlines)
+        dobj = ProbeData("\n".join(dlines))
         try:
-            self.probedata[client.hostname].update({data.get('name'): dtext})
+            self.probedata[client.hostname].update({data.get('name'): dobj})
         except KeyError:
-            self.probedata[client.hostname] = {data.get('name'): dtext}
+            self.probedata[client.hostname] = {data.get('name'): dobj}
 
     def get_additional_groups(self, meta):
         return self.cgroups.get(meta.hostname, list())